[Mlir-commits] [mlir] [mlir][SparseTensor] Fix crash in SparseTensorEmptyConverter (PR #208913)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Vito Secona (secona)

<details>
<summary>Changes</summary>

Fixes #<!-- -->202787

The crash comes from a call to `primaryTypeEncoding` where it reaches an `llvm_unreachable`. Previous change in #<!-- -->183898 fixes a the same crash issue of `primaryTypeEncoding` by adding `isValidPrimaryType()` check in SparseTensorNewConverter. The same fix is applied here by adding a call to `isValidPrimaryType()` in SparseTensorEmptyConverter. Now the pass errors instead of crashing.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp (+2) 
- (modified) mlir/test/Dialect/SparseTensor/conversion_invalid.mlir (+23) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index c9bcefe92de79..fb336c2f02f52 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -449,6 +449,8 @@ class SparseTensorEmptyConverter : public OpConversionPattern<tensor::EmptyOp> {
     const auto stt = getSparseTensorType(op);
     if (!stt.hasEncoding())
       return failure();
+    if (!isValidPrimaryType(stt.getElementType()))
+      return rewriter.notifyMatchFailure(op, "unsupported element type");
     // Gather all dimension sizes as SSA values.
     const Dimension dimRank = stt.getDimRank();
     SmallVector<Value> dimSizesValues;
diff --git a/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir b/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
index 456c1898a9311..9166e4c3a0512 100644
--- a/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
+++ b/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
@@ -12,3 +12,26 @@ func.func @new_index_elem_type(%arg0: index) {
   %0 = sparse_tensor.new %arg0 : index to tensor<?xindex, #sparse>
   return
 }
+
+// -----
+
+#map = affine_map<(d0) -> (0, d0)>
+#map1 = affine_map<(d0) -> (d0)>
+#sparse = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+module {
+  func.func @main(%arg0: tensor<1x77xi1>, %arg1: tensor<1x77xi1>) -> tensor<77xi1, #sparse> {
+    %0 = tensor.empty() : tensor<77xi1, #sparse>
+    %1 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel"]}
+      ins(%arg0, %arg1 : tensor<1x77xi1>, tensor<1x77xi1>)
+      outs(%0 : tensor<77xi1, #sparse>)
+    {
+    ^bb0(%in: i1, %in_0: i1, %out: i1):
+      %2 = arith.addi %in, %in_0 : i1
+      linalg.yield %2 : i1
+    } -> tensor<77xi1, #sparse>
+
+    // expected-error at +2 {{failed to legalize unresolved materialization}}
+    // expected-note at +1 {{see existing live user here}}
+    return %1 : tensor<77xi1, #sparse>
+  }
+}

``````````

</details>


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