[LAYOUTS] Ensure that we are consistent about what indices are constant (#7308)

apgoucher · web-flow · commit 48229b770ec9 · 2025-06-25T15:23:21.000Z
diff --git a/lib/Conversion/TritonGPUToLLVM/Utility.cpp b/lib/Conversion/TritonGPUToLLVM/Utility.cpp
@@ -216,10 +216,10 @@ applyLinearLayout(Location loc, RewriterBase &rewriter,
   }
 
   for (auto [inDimName, idx] : indices) {
-    if (idx.getDefiningOp<LLVM::ConstantOp>()) {
+    APInt constant;
+    if (matchPattern(idx, m_ConstantInt(&constant))) {
       continue;
     }
-
     int nBits = layout.getInDimSizeLog2(inDimName);
     for (int i = 0; i < nBits; i++) {
       Value bit = b.and_(idx, b.i32_val(1 << i));

Original file line number	Diff line number	Diff line change
`@@ -216,10 +216,10 @@ applyLinearLayout(Location loc, RewriterBase &rewriter,`
`216`	`216`	`}`
`217`	`217`
`218`	`218`	`for (auto [inDimName, idx] : indices) {`
`219`		`- if (idx.getDefiningOp<LLVM::ConstantOp>()) {`
	`219`	`+ APInt constant;`
	`220`	`+ if (matchPattern(idx, m_ConstantInt(&constant))) {`
`220`	`221`	`continue;`
`221`	`222`	`}`
`222`		`-`
`223`	`223`	`int nBits = layout.getInDimSizeLog2(inDimName);`
`224`	`224`	`for (int i = 0; i < nBits; i++) {`
`225`	`225`	`Value bit = b.and_(idx, b.i32_val(1 << i));`