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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: qyingwu

<details>
<summary>Changes</summary>

Fixes #<!-- -->205980.

`lower-sparse-iteration-to-scf` always called `linkNewScope()` when lowering
  iterators handled by `scf.for`. Dense levels are random-access iterators, and
  `linkNewScope()` asserts for random-access iterators because those should be
  traversed by coordinate.

  Use `locate()` for random-access iterators and keep `linkNewScope()` for
  non-random-access iterators.

  Verification:
  - `cmake --build /tmp/mlir-208198 --target mlir-opt`
  - `/tmp/mlir-208198/bin/mlir-opt -lower-sparse-iteration-to-scf /tmp/issue-
  205980.mlir`
  - `/tmp/mlir-208198/bin/llvm-lit -sv mlir/test/Dialect/SparseTensor/
  sparse_iteration_to_scf.mlir mlir/test/Dialect/SparseTensor/
  sparse_kernels_to_iterator.mlir mlir/test/Dialect/SparseTensor/
  sparse_space_collapse.mlir`

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseIterationToScf.cpp (+4-1) 
- (modified) mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf.mlir (+15) 


``````````diff
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
+}

``````````

</details>


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