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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From 7a5dd5c37015681aeb0e7f90ad12df0e8b83be37 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Feb 2026 03:03:29 -0800
Subject: [PATCH] [mlir][sparse] Fix crash in sparse_tensor.new with
 unsupported element type

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
---
 .../Transforms/SparseTensorConversion.cpp          |  3 +++
 .../SparseTensor/Transforms/Utils/CodegenUtils.cpp | 12 ++++++++++++
 .../SparseTensor/Transforms/Utils/CodegenUtils.h   |  5 +++++
 .../Dialect/SparseTensor/conversion_invalid.mlir   | 14 ++++++++++++++
 4 files changed, 34 insertions(+)
 create mode 100644 mlir/test/Dialect/SparseTensor/conversion_invalid.mlir

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