[Mlir-commits] [mlir] 8d082c7 - [mlir][sparse] Fix crash in sparse_tensor.new with unsupported element type (#183898)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-03T12:42:14+01:00
New Revision: 8d082c7c3144c0a635031387bf7bf6d60960ab9e

URL: https://github.com/llvm/llvm-project/commit/8d082c7c3144c0a635031387bf7bf6d60960ab9e
DIFF: https://github.com/llvm/llvm-project/commit/8d082c7c3144c0a635031387bf7bf6d60960ab9e.diff

LOG: [mlir][sparse] Fix crash in sparse_tensor.new with unsupported element type (#183898)

sparse_tensor.new with an element type not supported by the sparse
tensor runtime library (e.g. `index`) would call primaryTypeEncoding(),
which hits an llvm_unreachable("Unknown primary type"). The runtime
library only supports fixed-width types (f64/f32/f16/bf16,
i64/i32/i16/i8, and complex<f32/f64>); platform-dependent types such as
`index` have no corresponding PrimaryType encoding.

Fix by adding an isValidPrimaryType() helper and checking it in
SparseTensorNewConverter::matchAndRewrite() before calling genReader(),
so the conversion fails gracefully instead of crashing.

Fixes #180310

Added: 
    mlir/test/Dialect/SparseTensor/conversion_invalid.mlir

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index 0e88d31dae8e8..c9bcefe92de79 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -383,6 +383,9 @@ class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
     const auto stt = getSparseTensorType(op);
     if (!stt.hasEncoding())
       return failure();
+    // Verify that the element type is supported by the runtime library.
+    if (!isValidPrimaryType(stt.getElementType()))
+      return rewriter.notifyMatchFailure(op, "unsupported element type");
     // Construct the `reader` opening method calls.
     SmallVector<Value> dimSizesValues;
     Value dimSizesBuffer;

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
index f57f7f7fc0946..ffc898c9933c3 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
@@ -99,6 +99,18 @@ StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(Type tp) {
   return overheadTypeFunctionSuffix(overheadTypeEncoding(tp));
 }
 
+bool mlir::sparse_tensor::isValidPrimaryType(Type elemTp) {
+  if (elemTp.isF64() || elemTp.isF32() || elemTp.isF16() || elemTp.isBF16() ||
+      elemTp.isInteger(64) || elemTp.isInteger(32) || elemTp.isInteger(16) ||
+      elemTp.isInteger(8))
+    return true;
+  if (auto complexTp = dyn_cast<ComplexType>(elemTp)) {
+    Type elt = complexTp.getElementType();
+    return elt.isF64() || elt.isF32();
+  }
+  return false;
+}
+
 PrimaryType mlir::sparse_tensor::primaryTypeEncoding(Type elemTp) {
   if (elemTp.isF64())
     return PrimaryType::kF64;

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h
index 1c10dd5566184..98e351a371ac1 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h
@@ -60,6 +60,11 @@ StringRef overheadTypeFunctionSuffix(OverheadType ot);
 /// Converts an overhead storage type to its function-name suffix.
 StringRef overheadTypeFunctionSuffix(Type overheadTp);
 
+/// Returns true if the given type is a valid sparse tensor element type
+/// supported by the runtime library (i.e., maps to a PrimaryType).
+/// Use this to guard calls to primaryTypeEncoding() with invalid types.
+bool isValidPrimaryType(Type elemTp);
+
 /// Converts a primary storage type to its internal type-encoding.
 PrimaryType primaryTypeEncoding(Type elemTp);
 

diff  --git a/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir b/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
new file mode 100644
index 0000000000000..456c1898a9311
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt %s --sparse-tensor-conversion -verify-diagnostics -split-input-file
+
+// Regression test for https://github.com/llvm/llvm-project/issues/180310:
+// sparse_tensor.new with an unsupported element type (e.g. index) must not
+// crash with llvm_unreachable in primaryTypeEncoding; the conversion should
+// fail gracefully.
+
+#sparse = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+func.func @new_index_elem_type(%arg0: index) {
+  // expected-error at +1 {{failed to legalize operation 'sparse_tensor.new'}}
+  %0 = sparse_tensor.new %arg0 : index to tensor<?xindex, #sparse>
+  return
+}