Fix crash when processing getMakeTensorPtr op through while loop

alexbaden · alexbaden · commit aaa66531be04 · 2025-06-25T01:05:24.000Z
diff --git a/lib/Dialect/Triton/IR/Utility.cpp b/lib/Dialect/Triton/IR/Utility.cpp
@@ -89,6 +89,9 @@ tt::MakeTensorPtrOp tt::getMakeTensorPtrOp(Value v) {
                : condBr.getFalseDestOperands()[argNum]);
     return tt::getMakeTensorPtrOp(argOwner->getOperand(argNum));
   }
+  if (auto whileOp = dyn_cast<scf::WhileOp>(argOwner)) {
+    return tt::getMakeTensorPtrOp(whileOp.getOperand(argNum));
+  }
   llvm_unreachable("Unable to getMakeTensorPtr()");
 }
 

Original file line number	Diff line number	Diff line change
`@@ -89,6 +89,9 @@ tt::MakeTensorPtrOp tt::getMakeTensorPtrOp(Value v) {`
`89`	`89`	`: condBr.getFalseDestOperands()[argNum]);`
`90`	`90`	`return tt::getMakeTensorPtrOp(argOwner->getOperand(argNum));`
`91`	`91`	`}`
	`92`	`+ if (auto whileOp = dyn_cast<scf::WhileOp>(argOwner)) {`
	`93`	`+ return tt::getMakeTensorPtrOp(whileOp.getOperand(argNum));`
	`94`	`+ }`
`92`	`95`	`llvm_unreachable("Unable to getMakeTensorPtr()");`
`93`	`96`	`}`
`94`	`97`