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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/188959

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


>From c36b3ccdcc0de346f9e8f62331dd1b40e8c9d066 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Mar 2026 15:57:35 -0700
Subject: [PATCH] [MLIR][SparseTensor] Fix domination violation in co-iteration
 for dense iterators

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
Co-Authored-By: Claude Sonnet 4.6 <noreply at anthropic.com>
Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.
---
 .../Transforms/Utils/LoopEmitter.cpp          | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

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,