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

[Vito Secona]

https://github.com/secona created https://github.com/llvm/llvm-project/pull/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 

>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/2] 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/2] 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>
+}