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

[Arjun Bhamra]

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

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

>From 0f202c8dd27d9f7644fe55f58885b7ea58fa05fc Mon Sep 17 00:00:00 2001
From: Arjun Bhamra <arjun.bhamra25 at gmail.com>
Date: Thu, 12 Feb 2026 20:07:16 -0500
Subject: [PATCH 3/3] add clarifying comment to test file

---
 mlir/test/Dialect/SparseTensor/codegen_invalid.mlir | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/mlir/test/Dialect/SparseTensor/codegen_invalid.mlir b/mlir/test/Dialect/SparseTensor/codegen_invalid.mlir
index 9fb8bee7a5e66..792272796a2b5 100644
--- a/mlir/test/Dialect/SparseTensor/codegen_invalid.mlir
+++ b/mlir/test/Dialect/SparseTensor/codegen_invalid.mlir
@@ -1,5 +1,11 @@
 // RUN: mlir-opt %s -sparse-tensor-codegen -verify-diagnostics
 
+// NOTE: This test has valid IR, however we are testing whether
+// the legalization failure occurs when important passes are
+// missing. Notably, using --lower-sparse-ops-to-foreach
+// followed by --lower-sparse-foreach-to-scf prior to
+// sparse codegen will convert the dense tensor correctly.
+
 #SparseVector = #sparse_tensor.encoding<{
   map = (d0) -> (d0 : compressed)
 }>