[Mlir-commits] [mlir] [mlir][sparse] Reject dense level after non-unique level in encoding verifier (PR #184157)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sparse

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

The sparse tensor iteration model for dense levels requires exactly one parent position to linearize into a contiguous range. However, when a non-unique level (e.g. compressed(nonunique)) precedes a dense level, the DedupIterator provides a two-element cursor {posLo, segHi}, causing DenseLevel::peekRangeAt to assert and crash with:

  Assertion `parentPos.size() == 1 && "Dense level can not be non-unique."'

Fix this by adding a check in SparseTensorEncodingAttr::verify that rejects any encoding where a dense level directly follows a non-unique level, emitting a proper diagnostic instead of crashing during lowering.

Fixes #<!-- -->130008

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp (+9) 
- (modified) mlir/test/Dialect/SparseTensor/invalid_encoding.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 3a34ad90941b0..05b2e07b88ba5 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -826,6 +826,15 @@ LogicalResult SparseTensorEncodingAttr::verify(
     return emitError() << "SoA is only applicable to singleton lvlTypes.";
   }
 
+  // Dense levels cannot follow a non-unique level. The iteration model for
+  // dense levels requires exactly one parent position to linearize into a
+  // contiguous range, but a non-unique parent provides two cursor values
+  // (segment start and end), which the dense level cannot handle.
+  for (auto [i, lt] : llvm::drop_begin(llvm::enumerate(lvlTypes))) {
+    if (isDenseLT(lt) && !isUniqueLT(lvlTypes[i - 1]))
+      return emitError() << "dense level cannot follow a non-unique level";
+  }
+
   // TODO: audit formats that actually are supported by backend.
   if (auto it = llvm::find_if(lvlTypes, isNOutOfMLT);
       it != std::end(lvlTypes)) {
diff --git a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
index a3f72bd3ae971..0d0234149fb3e 100644
--- a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
@@ -256,6 +256,17 @@ func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
 
 // -----
 
+// expected-error at +1{{dense level cannot follow a non-unique level}}
+#COO_Dense = #sparse_tensor.encoding<{
+  map = (i, j) -> (
+    i : compressed(nonunique),
+    j : dense
+  )
+}>
+func.func private @dense_after_nonunique(tensor<?x?xf32, #COO_Dense>)
+
+// -----
+
 // expected-error at +2 {{use of undeclared identifier 'l1'}}
 #TooFewLvlDecl = #sparse_tensor.encoding<{
   map = {l0} (d0, d1) -> (l0 = d0 : dense, l1 = d1 : compressed)

``````````

</details>


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