[MLIR][SparseTensor] Reject unsupported symbolic demapping - #217476
Merged
Conversation
Symbolic maps cannot be lowered without runtime symbol operands. Decline the demapping rewrites before they create invalid reinterpretations, and report a legalization failure instead of asserting. Assisted-by: Codex
joker-eph
requested review from
PeimingLiu,
aartbik,
matthias-springer and
yinying-lisa-li
as code owners
August 19, 2026 22:05
|
@llvm/pr-subscribers-mlir Author: Mehdi Amini (joker-eph) ChangesSymbolic maps cannot be lowered without runtime symbol operands. Decline the demapping rewrites before they create invalid reinterpretations, and report a legalization failure instead of asserting. Assisted-by: Codex Full diff: https://github.com/llvm/llvm-project/pull/217476.diff 5 Files Affected:
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index 85f64648c212e..f280377f2cfc2 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -38,6 +38,12 @@ struct DemapInsRewriter : public OpRewritePattern<SourceOp> {
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
+ for (Value in : op->getOperands())
+ if (auto stt = tryGetSparseTensorType(in);
+ stt && !stt->isIdentity() &&
+ stt->getEncoding().getDimToLvl().getNumSymbols() != 0)
+ return failure();
+
// Demaps non-trivial inputs.
bool changed = false;
SmallVector<Value> deMappedIns(op->getOperands());
@@ -605,6 +611,8 @@ struct TensorAllocDemapper : public OpRewritePattern<AllocOp> {
Location loc = op.getLoc();
auto stt = getSparseTensorType(op.getResult());
+ if (stt.getEncoding().getDimToLvl().getNumSymbols() != 0)
+ return failure();
SmallVector<Value> maxDimCrds;
maxDimCrds.reserve(stt.getDimRank());
@@ -674,6 +682,8 @@ struct SparseAssembleDemapper : public OpRewritePattern<AssembleOp> {
assert(hasAnySparseResult(op));
auto stt = getSparseTensorType(op.getResult());
+ if (stt.getEncoding().getDimToLvl().getNumSymbols() != 0)
+ return failure();
rewriter.modifyOpInPlace(
op, [&op, &stt]() { op.getResult().setType(stt.getDemappedType()); });
rewriter.setInsertionPointAfter(op);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 73d3ec86775db..42a6c4d77dbdd 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -1340,6 +1340,8 @@ struct CrdTranslateRewriter : public OpRewritePattern<CrdTranslateOp> {
AffineMap map = op.getDirection() == CrdTransDirectionKind::dim2lvl
? op.getEncoder().getDimToLvl()
: op.getEncoder().getLvlToDim();
+ if (!map || map.getNumSymbols() != 0)
+ return failure();
SmallVector<Value> outCrds;
for (AffineExpr result : map.getResults()) {
diff --git a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
index 653ad63cae8c5..ba5c526ade1d7 100644
--- a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
+++ b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
@@ -1,8 +1,7 @@
// RUN: mlir-opt %s -split-input-file -sparsification-and-bufferization -verify-diagnostics | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
-// Tests that mlir-opt does not crash when parsing sparse tensor encodings with symbols.
+// Tests that processing sparse tensor encodings with symbols does not crash and
+// reports a diagnostic when lowering is unsupported.
// CHECK-DAG: #[[$SPARSE_0:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d1 : dense, d2 : compressed) }>
// CHECK-DAG: #[[$SPARSE_1:.*]] = #sparse_tensor.encoding<{ map = [s0](d0, d1) -> (d0 * (s0 * 3) : dense, d0 : dense, d1 : compressed) }>
@@ -47,7 +46,7 @@ func.func @tensor_convert() -> memref<?xindex> {
tensor.yield %val : f32
} : tensor<32x32xf32>
- // expected-error@+1 {{Level size mismatch between source/dest tensors}}
+ // expected-error@+1 {{failed to legalize operation 'bufferization.alloc_tensor'}}
%J = sparse_tensor.convert %I : tensor<32x32xf32> to tensor<32x32xf32, #Sparse>
%result = sparse_tensor.positions %J { level = 0 : index }
diff --git a/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
index c4ebec368a9ce..e03c003085eee 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
@@ -170,3 +170,16 @@ func.func @foreach_bcoo(%A: tensor<4x4x4xf64, #BCOO>) {
}
return
}
+
+#Symbolic = #sparse_tensor.encoding<{
+ map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
+}>
+
+// CHECK-LABEL: func.func @symbolic_crd_translate(
+// CHECK: sparse_tensor.crd_translate dim_to_lvl
+func.func @symbolic_crd_translate(%i: index, %j: index)
+ -> (index, index, index) {
+ %l0, %l1, %l2 = sparse_tensor.crd_translate dim_to_lvl [%i, %j]
+ as #Symbolic : index, index, index
+ return %l0, %l1, %l2 : index, index, index
+}
diff --git a/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir b/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
index 97c668716eecc..89ad74bd7000e 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
@@ -129,3 +129,21 @@ func.func @sparse_disassemble_reinterpret_map(%sp : tensor<2x4xf64, #BSR>,
-> (tensor<?xindex>, tensor<?xindex>), tensor<?xf64>, (index, index), index
return %rd, %rp, %ri : tensor<?xf64>, tensor<?xindex>, tensor<?xindex>
}
+
+// -----
+
+#Symbolic = #sparse_tensor.encoding<{
+ map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
+}>
+
+// CHECK-LABEL: func.func @sparse_assemble_symbolic_map(
+// CHECK: %[[ASSEMBLED:.*]] = sparse_tensor.assemble {{.*}} to tensor<32x32xf32, #sparse{{[0-9]*}}>
+// CHECK-NEXT: return %[[ASSEMBLED]] : tensor<32x32xf32, #sparse{{[0-9]*}}>
+func.func @sparse_assemble_symbolic_map(
+ %val: tensor<?xf32>, %pos: tensor<?xindex>, %crd: tensor<?xindex>)
+ -> tensor<32x32xf32, #Symbolic> {
+ %0 = sparse_tensor.assemble (%pos, %crd), %val
+ : (tensor<?xindex>, tensor<?xindex>), tensor<?xf32>
+ to tensor<32x32xf32, #Symbolic>
+ return %0 : tensor<32x32xf32, #Symbolic>
+}
|
|
@llvm/pr-subscribers-mlir-sparse Author: Mehdi Amini (joker-eph) ChangesSymbolic maps cannot be lowered without runtime symbol operands. Decline the demapping rewrites before they create invalid reinterpretations, and report a legalization failure instead of asserting. Assisted-by: Codex Full diff: https://github.com/llvm/llvm-project/pull/217476.diff 5 Files Affected:
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index 85f64648c212e..f280377f2cfc2 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -38,6 +38,12 @@ struct DemapInsRewriter : public OpRewritePattern<SourceOp> {
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
+ for (Value in : op->getOperands())
+ if (auto stt = tryGetSparseTensorType(in);
+ stt && !stt->isIdentity() &&
+ stt->getEncoding().getDimToLvl().getNumSymbols() != 0)
+ return failure();
+
// Demaps non-trivial inputs.
bool changed = false;
SmallVector<Value> deMappedIns(op->getOperands());
@@ -605,6 +611,8 @@ struct TensorAllocDemapper : public OpRewritePattern<AllocOp> {
Location loc = op.getLoc();
auto stt = getSparseTensorType(op.getResult());
+ if (stt.getEncoding().getDimToLvl().getNumSymbols() != 0)
+ return failure();
SmallVector<Value> maxDimCrds;
maxDimCrds.reserve(stt.getDimRank());
@@ -674,6 +682,8 @@ struct SparseAssembleDemapper : public OpRewritePattern<AssembleOp> {
assert(hasAnySparseResult(op));
auto stt = getSparseTensorType(op.getResult());
+ if (stt.getEncoding().getDimToLvl().getNumSymbols() != 0)
+ return failure();
rewriter.modifyOpInPlace(
op, [&op, &stt]() { op.getResult().setType(stt.getDemappedType()); });
rewriter.setInsertionPointAfter(op);
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
index 73d3ec86775db..42a6c4d77dbdd 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp
@@ -1340,6 +1340,8 @@ struct CrdTranslateRewriter : public OpRewritePattern<CrdTranslateOp> {
AffineMap map = op.getDirection() == CrdTransDirectionKind::dim2lvl
? op.getEncoder().getDimToLvl()
: op.getEncoder().getLvlToDim();
+ if (!map || map.getNumSymbols() != 0)
+ return failure();
SmallVector<Value> outCrds;
for (AffineExpr result : map.getResults()) {
diff --git a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
index 653ad63cae8c5..ba5c526ade1d7 100644
--- a/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
+++ b/mlir/test/Dialect/SparseTensor/encoding_with_symbols.mlir
@@ -1,8 +1,7 @@
// RUN: mlir-opt %s -split-input-file -sparsification-and-bufferization -verify-diagnostics | FileCheck %s
-// XFAIL: mlir-expensive-checks
-
-// Tests that mlir-opt does not crash when parsing sparse tensor encodings with symbols.
+// Tests that processing sparse tensor encodings with symbols does not crash and
+// reports a diagnostic when lowering is unsupported.
// CHECK-DAG: #[[$SPARSE_0:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : dense, d1 : dense, d2 : compressed) }>
// CHECK-DAG: #[[$SPARSE_1:.*]] = #sparse_tensor.encoding<{ map = [s0](d0, d1) -> (d0 * (s0 * 3) : dense, d0 : dense, d1 : compressed) }>
@@ -47,7 +46,7 @@ func.func @tensor_convert() -> memref<?xindex> {
tensor.yield %val : f32
} : tensor<32x32xf32>
- // expected-error@+1 {{Level size mismatch between source/dest tensors}}
+ // expected-error@+1 {{failed to legalize operation 'bufferization.alloc_tensor'}}
%J = sparse_tensor.convert %I : tensor<32x32xf32> to tensor<32x32xf32, #Sparse>
%result = sparse_tensor.positions %J { level = 0 : index }
diff --git a/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
index c4ebec368a9ce..e03c003085eee 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_foreach.mlir
@@ -170,3 +170,16 @@ func.func @foreach_bcoo(%A: tensor<4x4x4xf64, #BCOO>) {
}
return
}
+
+#Symbolic = #sparse_tensor.encoding<{
+ map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
+}>
+
+// CHECK-LABEL: func.func @symbolic_crd_translate(
+// CHECK: sparse_tensor.crd_translate dim_to_lvl
+func.func @symbolic_crd_translate(%i: index, %j: index)
+ -> (index, index, index) {
+ %l0, %l1, %l2 = sparse_tensor.crd_translate dim_to_lvl [%i, %j]
+ as #Symbolic : index, index, index
+ return %l0, %l1, %l2 : index, index, index
+}
diff --git a/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir b/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
index 97c668716eecc..89ad74bd7000e 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
@@ -129,3 +129,21 @@ func.func @sparse_disassemble_reinterpret_map(%sp : tensor<2x4xf64, #BSR>,
-> (tensor<?xindex>, tensor<?xindex>), tensor<?xf64>, (index, index), index
return %rd, %rp, %ri : tensor<?xf64>, tensor<?xindex>, tensor<?xindex>
}
+
+// -----
+
+#Symbolic = #sparse_tensor.encoding<{
+ map = [c](i, j) -> (c * 3 * i : dense, i : dense, j : compressed)
+}>
+
+// CHECK-LABEL: func.func @sparse_assemble_symbolic_map(
+// CHECK: %[[ASSEMBLED:.*]] = sparse_tensor.assemble {{.*}} to tensor<32x32xf32, #sparse{{[0-9]*}}>
+// CHECK-NEXT: return %[[ASSEMBLED]] : tensor<32x32xf32, #sparse{{[0-9]*}}>
+func.func @sparse_assemble_symbolic_map(
+ %val: tensor<?xf32>, %pos: tensor<?xindex>, %crd: tensor<?xindex>)
+ -> tensor<32x32xf32, #Symbolic> {
+ %0 = sparse_tensor.assemble (%pos, %crd), %val
+ : (tensor<?xindex>, tensor<?xindex>), tensor<?xf32>
+ to tensor<32x32xf32, #Symbolic>
+ return %0 : tensor<32x32xf32, #Symbolic>
+}
|
aartbik
approved these changes
Aug 20, 2026
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
2 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.
Symbolic maps cannot be lowered without runtime symbol operands. Decline the demapping rewrites before they create invalid reinterpretations, and report a legalization failure instead of asserting.
Assisted-by: Codex