[Mlir-commits] [mlir] [mlir][SparseTensor] handle uninitialized transMap when translating shape (PR #195506)

[Vito Secona]

https://github.com/secona updated https://github.com/llvm/llvm-project/pull/195506

>From 82d8c391dfe734ac73ffdf97f88438135fbc12c7 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Sat, 2 May 2026 22:45:18 +0700
Subject: [PATCH 1/3] handle unitialized transMap

---
 mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index eab2d14797257..df3160ab700c4 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -513,6 +513,15 @@ SparseTensorEncodingAttr::translateShape(ArrayRef<int64_t> srcShape,
   AffineMap transMap =
       dir == CrdTransDirectionKind::dim2lvl ? getDimToLvl() : getLvlToDim();
 
+  // Check if transMap is valid. There are cases where the lvlToDim map is
+  // uninitialized due to the format used, e.g. ELL. This is visible as
+  // inferring lvlToDim (see inferLvlToDim function below) may return an
+  // uninitialized affine map. Fallback to dynamic shapes.
+  if (!transMap) {
+    ret.resize(rank, ShapedType::kDynamic);
+    return ret;
+  }
+
   SmallVector<AffineExpr> dimRep;
   dimRep.reserve(srcShape.size());
   for (int64_t sz : srcShape) {

>From 1eb3a6a7b666918e0fdaec5776484ebc378efa8e Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Sat, 2 May 2026 23:28:12 +0700
Subject: [PATCH 2/3] add regression test

---
 .../SparseTensor/encoding_with_symbols.mlir   | 31 ++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
index 7cd68ee00dd09..ae216a127048f 100644
--- a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
+++ b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -sparsification-and-bufferization | FileCheck %s
+// RUN: mlir-opt %s -split-input-file -sparsification-and-bufferization -verify-diagnostics | FileCheck %s
 
 // Tests that mlir-opt does not crash when parsing sparse tensor encodings with symbols.
 
@@ -24,3 +24,32 @@ func.func @tensor_add(%arg0: tensor<8x8xf32, #Sparse>) -> tensor<8x8xf32> {
   // CHECK: return %{{.*}} : memref<8x8xf32>
   return %result : tensor<8x8xf32>
 }
+
+// -----
+
+// This section makes sure that using the following encoding does not result in
+// an assertion error, but instead the expected error. Ultimately, we want to
+// make this section pass without any expected errors.
+
+#Sparse = #sparse_tensor.encoding<{
+  map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
+}>
+
+func.func @tensor_convert() -> memref<?xindex> {
+  %I = tensor.generate {
+  ^bb0(%i: index, %j: index):
+    %is_diag = arith.cmpi eq, %i, %j : index
+    %f0 = arith.constant 0.0 : f32
+    %f1 = arith.constant 1.0 : f32
+    %val = arith.select %is_diag, %f1, %f0 : f32
+    tensor.yield %val : f32
+  } : tensor<32x32xf32>
+
+  // expected-error at +1 {{'bufferization.alloc_tensor' op operand count (1) does not match with the total size (0) specified in attribute 'operandSegmentSizes'}}
+  %J = sparse_tensor.convert %I : tensor<32x32xf32> to tensor<32x32xf32, #Sparse>
+
+  %result = sparse_tensor.positions %J { level = 0 : index }
+    : tensor<32x32xf32, #Sparse> to memref<?xindex>
+
+  return %result : memref<?xindex>
+}

>From b61684e9b51ede8ddd598fda3cb1bcf8e4993b96 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Mon, 4 May 2026 14:37:40 +0700
Subject: [PATCH 3/3] fix bufferize.alloc_tensor rewrite

---
 .../Transforms/SparseReinterpretMap.cpp             | 13 ++++++-------
 .../Dialect/SparseTensor/encoding_with_symbols.mlir |  2 +-
 2 files changed, 7 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index 0fc5cc76de39c..00ca33d14313a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -633,14 +633,13 @@ struct TensorAllocDemapper : public OpRewritePattern<AllocOp> {
     }
 
     assert(dynSz.empty()); // should have consumed all.
-    rewriter.startOpModification(op);
-    op->setOperands(dynLvlSzs);
-    op.getResult().setType(stt.getDemappedType());
-    rewriter.finalizeOpModification(op);
-    rewriter.setInsertionPointAfter(op);
 
-    Value t = genRemap(rewriter, stt.getEncoding(), op.getResult());
-    rewriter.replaceAllUsesExcept(op.getResult(), t, t.getDefiningOp());
+    // create a new op to let the MLIR builder calculate the correct metadata.
+    auto allocOp = bufferization::AllocTensorOp::create(
+        rewriter, loc, stt.getDemappedType(), dynLvlSzs);
+
+    Value t = genRemap(rewriter, stt.getEncoding(), allocOp.getResult());
+    rewriter.replaceOp(op, t);
     return success();
   }
 };
diff --git a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
index ae216a127048f..a019a0e9476b6 100644
--- a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
+++ b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
@@ -45,7 +45,7 @@ func.func @tensor_convert() -> memref<?xindex> {
     tensor.yield %val : f32
   } : tensor<32x32xf32>
 
-  // expected-error at +1 {{'bufferization.alloc_tensor' op operand count (1) does not match with the total size (0) specified in attribute 'operandSegmentSizes'}}
+  // expected-error at +1 {{Level size mismatch between source/dest tensors}}
   %J = sparse_tensor.convert %I : tensor<32x32xf32> to tensor<32x32xf32, #Sparse>
 
   %result = sparse_tensor.positions %J { level = 0 : index }