[Mlir-commits] [mlir] [mlir][sparse] Handle dense iterators in sparse iteration lowering (PR #208963)

[llvmlistbot at llvm.org]

https://github.com/qyingwu updated https://github.com/llvm/llvm-project/pull/208963

>From cc7eade3da850dd6070b3468c2fcd236d920d7e7 Mon Sep 17 00:00:00 2001
From: qyingwu <qiyingwu at utexas.edu>
Date: Sat, 11 Jul 2026 15:44:01 -0700
Subject: [PATCH] [mlir][sparse] Handle dense iterators in sparse iteration
 lowering

---
 .../Transforms/SparseIterationToScf.cpp           |  5 ++++-
 .../SparseTensor/sparse_iteration_to_scf.mlir     | 15 +++++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
index 9cd489653a0f3..e4cecef2d44e5 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp
@@ -134,8 +134,11 @@ static ValueRange genLoopWithIterator(
         });
     {
       OpBuilder::InsertionGuard guard(rewriter);
-      it->linkNewScope(forOp.getInductionVar());
       rewriter.setInsertionPointToStart(forOp.getBody());
+      if (it->randomAccessible())
+        it->locate(rewriter, loc, forOp.getInductionVar());
+      else
+        it->linkNewScope(forOp.getInductionVar());
       SmallVector<Value> ret = bodyBuilder(rewriter, loc, forOp.getBodyRegion(),
                                            it, forOp.getRegionIterArgs());
 
diff --git a/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir b/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir
index 6d8f9018ad3a5..855f1e99e7396 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir
@@ -50,3 +50,18 @@ func.func @sparse_iteration_to_scf(%sp : tensor<4x8xf32, #COO>) -> index {
   }
   return %r1 : index
 }
+
+#DenseCompressed = #sparse_tensor.encoding<{
+  map = (d0, d1) -> (d0 : dense, d1 : compressed)
+}>
+
+// CHECK-LABEL:   @sparse_iteration_dense_level
+// CHECK:           scf.for
+func.func @sparse_iteration_dense_level(%sp: tensor<?x?xf64, #DenseCompressed>) {
+  %0 = sparse_tensor.extract_iteration_space %sp lvls = 0
+      : tensor<?x?xf64, #DenseCompressed> -> !sparse_tensor.iter_space<#DenseCompressed, lvls = 0>
+  sparse_tensor.iterate %it in %0 at(%i)
+      : !sparse_tensor.iter_space<#DenseCompressed, lvls = 0> {
+  }
+  return
+}