[Mlir-commits] [mlir] [MLIR][SparseTensor] Fix domination violation in co-iteration for dense iterators (PR #188959)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

In exitWhileLoop, random-accessible (dense) iterators were being located using whileOp.getResults().back() while the insertion point was still inside the while loop's after block. This caused a domination violation: the ADDI created by locate() was inside the after block, but it was later used (via derefImpl's SUBI) after the while loop exits.

Move the locate() calls for random-accessible iterators to after builder.setInsertionPointAfter(whileOp), where the while results are properly in scope.

Fixes 10 failing tests under MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

Assisted-by: Claude Code


---
Full diff: https://github.com/llvm/llvm-project/pull/188959.diff


1 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp (+14-5) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp
index 684c088eb9b0f..650939f856514 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/LoopEmitter.cpp
@@ -824,12 +824,9 @@ void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
       // Following loops continue iteration from the break point of the
       // current while loop.
       whileRes = it.linkNewScope(whileRes);
-    } else {
-      // Make sure randomly accessible (dense) iterator is set to the right
-      // position according to the universal index.
-      Value uniIdx = whileOp.getResults().back();
-      it.locate(builder, loc, uniIdx);
     }
+    // Note: random-accessible (dense) iterators are located after the while
+    // loop exits; see below.
   }
 
   // Reduction value from users.
@@ -853,6 +850,18 @@ void LoopEmitter::exitWhileLoop(OpBuilder &builder, Location loc,
     YIELD(operands);
 
   builder.setInsertionPointAfter(whileOp);
+
+  // Locate random-accessible (dense) iterators at the final universal index.
+  // This must happen after the while loop exits so that the while results are
+  // in scope. Doing it inside the after block (using while results there) would
+  // cause domination violations when the iterator is later dereferenced.
+  for (auto [tid, lvl] : unpackTensorLevelRange(loopInfo.tidLvls)) {
+    SparseIterator &it = getCurIterator(tid, lvl);
+    if (it.randomAccessible()) {
+      Value uniIdx = whileOp.getResults().back();
+      it.locate(builder, loc, uniIdx);
+    }
+  }
 }
 
 void LoopEmitter::exitCurrentLoop(RewriterBase &rewriter, Location loc,

``````````

</details>


https://github.com/llvm/llvm-project/pull/188959