[Mlir-commits] [mlir] [MLIR] Graceful handling of unint sparse tensor encodings with `sparse-tensor-codegen` (PR #181145)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-sparse

Author: Arjun Bhamra (abhamra)

<details>
<summary>Changes</summary>

This PR handles the case where users call the `--sparse-tensor-codegen` pass without sufficiently lowering dense tensors to sparse ones (with passes like `--lower-sparse-ops-to-foreach` and `--lower-sparse-foreach-to-scf` among others). This results in dense tensors having a null `SparseTensorEncodingAttr`, which was originally assumed to be true in the SparseTensor's `ConvertOp` lowering, but is now checked against.

This PR closes #<!-- -->177779.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp (+7) 
- (added) mlir/test/Dialect/SparseTensor/codegen_invalid.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 4c28658dec67d..0ca63b46e25a8 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1150,6 +1150,13 @@ class SparseConvertConverter : public OpConversionPattern<ConvertOp> {
     SparseTensorEncodingAttr encDst = getSparseTensorEncoding(op.getType());
     SparseTensorEncodingAttr encSrc =
         getSparseTensorEncoding(op.getSource().getType());
+
+    // If either the source or the destination don't have a valid sparse
+    // tensor encoding, we should fail to legalize. This should be handled
+    // by another set of passes before reaching here.
+    if (!encSrc || !encDst)
+      return failure();
+
     // The output tensor can not be a slice and those cases should have been
     // rejected by ConvertOp::verify() already.
     assert(!encDst.isSlice() && "Cannot convert to a sparse tensor slices.");
diff --git a/mlir/test/Dialect/SparseTensor/codegen_invalid.mlir b/mlir/test/Dialect/SparseTensor/codegen_invalid.mlir
new file mode 100644
index 0000000000000..9fb8bee7a5e66
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/codegen_invalid.mlir
@@ -0,0 +1,20 @@
+// RUN: mlir-opt %s -sparse-tensor-codegen -verify-diagnostics
+
+#SparseVector = #sparse_tensor.encoding<{
+  map = (d0) -> (d0 : compressed)
+}>
+
+module {
+  func.func @main() -> tensor<8xf32, #SparseVector> {
+    %dense = arith.constant dense<[1.0, 0.0, 0.0, 2.0, 0.0, 3.0, 0.0, 0.0]>
+      : tensor<8xf32>
+
+    // expected-error at +1 {{failed to legalize operation 'sparse_tensor.convert' that was explicitly marked illegal}}
+    %sparse = sparse_tensor.convert %dense
+      : tensor<8xf32> to tensor<8xf32, #SparseVector>
+
+    return %sparse : tensor<8xf32, #SparseVector>
+  }
+}
+
+// -----

``````````

</details>


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