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

[llvmlistbot at llvm.org]

Author: Vito Secona
Date: 2026-05-07T11:46:46-07:00
New Revision: 02c8eb80ae0d5bd841c81b5439cee197b29fe4fc

URL: https://github.com/llvm/llvm-project/commit/02c8eb80ae0d5bd841c81b5439cee197b29fe4fc
DIFF: https://github.com/llvm/llvm-project/commit/02c8eb80ae0d5bd841c81b5439cee197b29fe4fc.diff

LOG: [mlir][SparseTensor] handle uninitialized transMap when translating shape (#195506)

When translating a shape using
`SparseTensorEncodingAttr::translateShape` from lvl to dim, there is a
possibility that the `transMap` map (`lvlToDim` map under the hood) is
uninitialized. This leads to an assertion error when calling the
`.getResults()` method.

This change adds a guard to check if the `transMap` map is uninitialized
and return early with dynamic shapes. This change also adds a regression
test based on the reproduce MLIR code.

Closes #195464

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
    mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir

Removed: 
    


################################################################################
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) {

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index 0fc5cc76de39c..85f64648c212e 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 =
+        AllocOp::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 7cd68ee00dd09..a019a0e9476b6 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 {{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 }
+    : tensor<32x32xf32, #Sparse> to memref<?xindex>
+
+  return %result : memref<?xindex>
+}