[AMD] fixed getLeafForOps implementation

triton-lang · Nov 25, 2024 · f62acd8 · f62acd8
1 parent 576bee1
commit f62acd8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.cpp b/third_party/amd/lib/TritonAMDGPUToLLVM/Utility.cpp
@@ -366,9 +366,9 @@ SmallVector<scf::ForOp> getLeafForOps(triton::FuncOp funcOp) {
 
   SmallVector<scf::ForOp> leafOps;
   for (scf::ForOp forOp : allOps) {
-    auto searchResult =
-        forOp->walk([](scf::ForOp) { return WalkResult::interrupt(); });
-    if (searchResult.wasInterrupted())
+    auto searchResult = forOp.getBody()->walk(
+        [](scf::ForOp) { return WalkResult::interrupt(); });
+    if (!searchResult.wasInterrupted())
       leafOps.push_back(forOp);
   }
   return leafOps;