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

[Arjun Bhamra]

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

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.

>From 6eab80d6580a6f13c67e1c9632b8b5891aa4a046 Mon Sep 17 00:00:00 2001
From: Arjun Bhamra <arjun.bhamra25 at gmail.com>
Date: Tue, 10 Feb 2026 15:15:04 -0500
Subject: [PATCH 1/2] added legalization check for non-null encSrc, encDst

---
 .../SparseTensor/Transforms/SparseTensorCodegen.cpp        | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 1fbcf5fdc68db..540ed44239f4a 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1149,6 +1149,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.");

>From e385e524790e820d4ed5173c3585d9c611511295 Mon Sep 17 00:00:00 2001
From: Arjun Bhamra <arjun.bhamra25 at gmail.com>
Date: Thu, 12 Feb 2026 09:06:08 -0500
Subject: [PATCH 2/2] Added codegen_invalid test file, test for legalization
 failure

---
 .../Dialect/SparseTensor/codegen_invalid.mlir | 20 +++++++++++++++++++
 1 file changed, 20 insertions(+)
 create mode 100644 mlir/test/Dialect/SparseTensor/codegen_invalid.mlir

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