[mlir][sparse] Fix crash in sparse_tensor.new with unsupported element type#183898
Merged
Conversation
…t 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 llvm#180310
joker-eph
requested review from
PeimingLiu,
aartbik,
matthias-springer and
yinying-lisa-li
as code owners
Member
|
@llvm/pr-subscribers-mlir-sparse Author: Mehdi Amini (joker-eph) Changessparse_tensor.new with an element type not supported by the sparse tensor runtime library (e.g. 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 Full diff: https://github.com/llvm/llvm-project/pull/183898.diff 4 Files Affected:
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@+1 {{failed to legalize operation 'sparse_tensor.new'}}
+ %0 = sparse_tensor.new %arg0 : index to tensor<?xindex, #sparse>
+ return
+}
|
Member
|
@llvm/pr-subscribers-mlir Author: Mehdi Amini (joker-eph) Changessparse_tensor.new with an element type not supported by the sparse tensor runtime library (e.g. 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 Full diff: https://github.com/llvm/llvm-project/pull/183898.diff 4 Files Affected:
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@+1 {{failed to legalize operation 'sparse_tensor.new'}}
+ %0 = sparse_tensor.new %arg0 : index to tensor<?xindex, #sparse>
+ return
+}
|
aartbik
approved these changes
Mar 2, 2026
sahas3
pushed a commit
to sahas3/llvm-project
that referenced
this pull request
Mar 4, 2026
…t type (llvm#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 llvm#180310
sujianIBM
pushed a commit
to sujianIBM/llvm-project
that referenced
this pull request
Mar 5, 2026
…t type (llvm#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 llvm#180310
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 asindexhave 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