[mlir][SparseTensor] handle uninitialized transMap when translating shape

[secona]

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

[llvmorg-github-actions]

@llvm/pr-subscribers-mlir

@llvm/pr-subscribers-mlir-sparse

Author: Vito Secona (secona)

Changes

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

---

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

2 Files Affected:

• (modified) mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp (+9)
• (modified) mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir (+30-1)

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/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@+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>
+}

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[secona]

I need help with the merge :)