[Mlir-commits] [mlir] [mlir][sparse] Collapse dense and compressed iteration ranges (PR #211154)

[llvmlistbot at llvm.org]

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

>From 1e1d4b9258b80472e784dc6f05c1495369059fcc Mon Sep 17 00:00:00 2001
From: qyingwu <qiyingwu at utexas.edu>
Date: Tue, 21 Jul 2026 18:35:09 -0700
Subject: [PATCH] [mlir][sparse] Diagnose unsupported collapsed iteration
 spaces

---
 .../Transforms/SparseTensorPasses.cpp         | 27 ++++++++++++++++
 .../sparse_iteration_to_scf_invalid.mlir      | 31 +++++++++++++++++++
 2 files changed, 58 insertions(+)
 create mode 100644 mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf_invalid.mlir

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
index b660e22154688..98dffad8166b4 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
@@ -48,6 +48,28 @@ namespace {
 // Passes implementation.
 //===----------------------------------------------------------------------===//
 
+static bool supportsCollapsedRangeBetween(LevelType lt) {
+  return isSingletonLT(lt);
+}
+
+static LogicalResult verifyLowerableCollapsedIterSpaces(Operation *op) {
+  WalkResult result = op->walk([&](ExtractIterSpaceOp op) {
+    auto [lvlLo, lvlHi] = op.getLvlRange();
+    SparseTensorEncodingAttr enc = cast<SparseTensorEncodingAttr>(
+        getRankedTensorType(op.getTensor()).getEncoding());
+    for (Level lvl = lvlLo + 1; lvl < lvlHi; ++lvl) {
+      LevelType lt = enc.getLvlType(lvl);
+      if (!supportsCollapsedRangeBetween(lt)) {
+        op.emitOpError() << "cannot lower collapsed iteration space with "
+                         << toMLIRString(lt) << " level after the first level";
+        return WalkResult::interrupt();
+      }
+    }
+    return WalkResult::advance();
+  });
+  return failure(result.wasInterrupted());
+}
+
 struct SparseAssembler : public impl::SparseAssemblerBase<SparseAssembler> {
   SparseAssembler() = default;
   SparseAssembler(const SparseAssembler &pass) = default;
@@ -167,6 +189,11 @@ struct LowerSparseIterationToSCFPass
       default;
 
   void runOnOperation() override {
+    if (failed(verifyLowerableCollapsedIterSpaces(getOperation()))) {
+      signalPassFailure();
+      return;
+    }
+
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseIterationTypeConverter converter;
diff --git a/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf_invalid.mlir b/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf_invalid.mlir
new file mode 100644
index 0000000000000..d9909628d3a0b
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/sparse_iteration_to_scf_invalid.mlir
@@ -0,0 +1,31 @@
+// RUN: mlir-opt %s --split-input-file --lower-sparse-iteration-to-scf -verify-diagnostics
+
+#DenseCompressed = #sparse_tensor.encoding<{
+  map = (d0, d1) -> (d0 : dense, d1 : compressed)
+}>
+
+func.func @unsupported_dense_compressed_space(%sp: tensor<?x?xf64, #DenseCompressed>) {
+  // expected-error at below {{'sparse_tensor.extract_iteration_space' op cannot lower collapsed iteration space with compressed level after the first level}}
+  %0 = sparse_tensor.extract_iteration_space %sp lvls = 0 to 2
+      : tensor<?x?xf64, #DenseCompressed> -> !sparse_tensor.iter_space<#DenseCompressed, lvls = 0 to 2>
+  sparse_tensor.iterate %it in %0 at(%i, %j)
+      : !sparse_tensor.iter_space<#DenseCompressed, lvls = 0 to 2> {
+  }
+  return
+}
+
+// -----
+
+#CompressedDense = #sparse_tensor.encoding<{
+  map = (d0, d1) -> (d0 : compressed, d1 : dense)
+}>
+
+func.func @unsupported_compressed_dense_space(%sp: tensor<?x?xf64, #CompressedDense>) {
+  // expected-error at below {{'sparse_tensor.extract_iteration_space' op cannot lower collapsed iteration space with dense level after the first level}}
+  %0 = sparse_tensor.extract_iteration_space %sp lvls = 0 to 2
+      : tensor<?x?xf64, #CompressedDense> -> !sparse_tensor.iter_space<#CompressedDense, lvls = 0 to 2>
+  sparse_tensor.iterate %it in %0 at(%i, %j)
+      : !sparse_tensor.iter_space<#CompressedDense, lvls = 0 to 2> {
+  }
+  return
+}