[Mlir-commits] [mlir] [mlir][SparseTensor] Reject sparse encoding on non-primary types (PR #209591)
Vito Secona
llvmlistbot at llvm.org
Mon Jul 20 07:13:39 PDT 2026
https://github.com/secona updated https://github.com/llvm/llvm-project/pull/209591
>From c3ecdc74839369e26f73f60e3733e6f070cc4429 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Wed, 15 Jul 2026 01:42:06 +0700
Subject: [PATCH 1/2] [mlir][SparseTensor] Reject sparse encoding on
non-primary types
---
.../SparseTensor/IR/SparseTensorDialect.cpp | 14 ++
.../Transforms/SparseTensorConversion.cpp | 3 -
.../Transforms/Utils/CodegenUtils.cpp | 12 --
.../Transforms/Utils/CodegenUtils.h | 5 -
.../SparseTensor/constant_index_map.mlir | 44 ++---
.../SparseTensor/conversion_invalid.mlir | 99 ++++++++++-
.../SparseTensor/vectorize_reduction.mlir | 160 +++++++++---------
7 files changed, 214 insertions(+), 123 deletions(-)
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 7c9f552d62255..6c26a31bf88a0 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -919,6 +919,18 @@ LogicalResult SparseTensorEncodingAttr::verify(
return success();
}
+static bool 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;
+}
+
LogicalResult SparseTensorEncodingAttr::verifyEncoding(
ArrayRef<Size> dimShape, Type elementType,
function_ref<InFlightDiagnostic()> emitError) const {
@@ -964,6 +976,8 @@ LogicalResult SparseTensorEncodingAttr::verifyEncoding(
return emitError() << "implicit value must be zero";
}
}
+ if (!isValidPrimaryType(elementType))
+ return emitError() << "invalid primary type";
return success();
}
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
index c9bcefe92de79..0e88d31dae8e8 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp
@@ -383,9 +383,6 @@ 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 614c29ba68481..3c6a905ebb696 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.cpp
@@ -99,18 +99,6 @@ 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 98e351a371ac1..1c10dd5566184 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/CodegenUtils.h
@@ -60,11 +60,6 @@ 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/constant_index_map.mlir b/mlir/test/Dialect/SparseTensor/constant_index_map.mlir
index cf1eb3e9e44f5..6c6e4579d4f11 100644
--- a/mlir/test/Dialect/SparseTensor/constant_index_map.mlir
+++ b/mlir/test/Dialect/SparseTensor/constant_index_map.mlir
@@ -8,34 +8,34 @@
#SpVec = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
// CHECK-LABEL: func.func @main(
-// CHECK-SAME: %[[VAL_0:.*0]]: tensor<1x77xi1>,
-// CHECK-SAME: %[[VAL_1:.*1]]: tensor<1x77xi1>) -> tensor<77xi1, #{{.*}}> {
+// CHECK-SAME: %[[VAL_0:.*0]]: tensor<1x77xi32>,
+// CHECK-SAME: %[[VAL_1:.*1]]: tensor<1x77xi32>) -> tensor<77xi32, #{{.*}}> {
// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 77 : index
// CHECK-DAG: %[[VAL_3:.*]] = arith.constant 0 : index
// CHECK-DAG: %[[VAL_4:.*]] = arith.constant 1 : index
-// CHECK-DAG: %[[VAL_5:.*]] = tensor.empty() : tensor<77xi1, #{{.*}}>
-// CHECK-DAG: %[[VAL_6:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<1x77xi1>
-// CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_buffer %[[VAL_1]] : tensor<1x77xi1>
-// CHECK: %[[VAL_8:.*]] = scf.for %[[VAL_9:.*]] = %[[VAL_3]] to %[[VAL_2]] step %[[VAL_4]] iter_args(%[[VAL_10:.*]] = %[[VAL_5]]) -> (tensor<77xi1, #{{.*}}>) {
-// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]], %[[VAL_9]]] : memref<1x77xi1>
-// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_3]], %[[VAL_9]]] : memref<1x77xi1>
-// CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_11]], %[[VAL_12]] : i1
-// CHECK: %[[VAL_14:.*]] = tensor.insert %[[VAL_13]] into %[[VAL_10]]{{\[}}%[[VAL_9]]] : tensor<77xi1, #{{.*}}>
-// CHECK: scf.yield %[[VAL_14]] : tensor<77xi1, #{{.*}}>
+// CHECK-DAG: %[[VAL_5:.*]] = tensor.empty() : tensor<77xi32, #{{.*}}>
+// CHECK-DAG: %[[VAL_6:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<1x77xi32>
+// CHECK-DAG: %[[VAL_7:.*]] = bufferization.to_buffer %[[VAL_1]] : tensor<1x77xi32>
+// CHECK: %[[VAL_8:.*]] = scf.for %[[VAL_9:.*]] = %[[VAL_3]] to %[[VAL_2]] step %[[VAL_4]] iter_args(%[[VAL_10:.*]] = %[[VAL_5]]) -> (tensor<77xi32, #{{.*}}>) {
+// CHECK: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_3]], %[[VAL_9]]] : memref<1x77xi32>
+// CHECK: %[[VAL_12:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_3]], %[[VAL_9]]] : memref<1x77xi32>
+// CHECK: %[[VAL_13:.*]] = arith.addi %[[VAL_11]], %[[VAL_12]] : i32
+// CHECK: %[[VAL_14:.*]] = tensor.insert %[[VAL_13]] into %[[VAL_10]]{{\[}}%[[VAL_9]]] : tensor<77xi32, #{{.*}}>
+// CHECK: scf.yield %[[VAL_14]] : tensor<77xi32, #{{.*}}>
// CHECK: }
-// CHECK: %[[VAL_15:.*]] = sparse_tensor.load %[[VAL_16:.*]] hasInserts : tensor<77xi1, #{{.*}}>
-// CHECK: return %[[VAL_15]] : tensor<77xi1, #{{.*}}>
+// CHECK: %[[VAL_15:.*]] = sparse_tensor.load %[[VAL_16:.*]] hasInserts : tensor<77xi32, #{{.*}}>
+// CHECK: return %[[VAL_15]] : tensor<77xi32, #{{.*}}>
// CHECK: }
-func.func @main(%arg0: tensor<1x77xi1>, %arg1: tensor<1x77xi1>) -> tensor<77xi1, #SpVec> {
- %0 = tensor.empty() : tensor<77xi1, #SpVec>
+func.func @main(%arg0: tensor<1x77xi32>, %arg1: tensor<1x77xi32>) -> tensor<77xi32, #SpVec> {
+ %0 = tensor.empty() : tensor<77xi32, #SpVec>
%1 = linalg.generic {
indexing_maps = [#map1, #map1, #map2],
iterator_types = ["parallel"]}
- ins(%arg0, %arg1 : tensor<1x77xi1>, tensor<1x77xi1>)
- outs(%0 : tensor<77xi1, #SpVec>) {
- ^bb0(%in: i1, %in_0: i1, %out: i1):
- %2 = arith.addi %in, %in_0 : i1
- linalg.yield %2 : i1
- } -> tensor<77xi1, #SpVec>
- return %1 : tensor<77xi1, #SpVec>
+ ins(%arg0, %arg1 : tensor<1x77xi32>, tensor<1x77xi32>)
+ outs(%0 : tensor<77xi32, #SpVec>) {
+ ^bb0(%in: i32, %in_0: i32, %out: i32):
+ %2 = arith.addi %in, %in_0 : i32
+ linalg.yield %2 : i32
+ } -> tensor<77xi32, #SpVec>
+ return %1 : tensor<77xi32, #SpVec>
}
diff --git a/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir b/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
index 456c1898a9311..7e95be495b333 100644
--- a/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
+++ b/mlir/test/Dialect/SparseTensor/conversion_invalid.mlir
@@ -8,7 +8,104 @@
#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'}}
+ // expected-error at +1 {{invalid primary type}}
%0 = sparse_tensor.new %arg0 : index to tensor<?xindex, #sparse>
return
}
+
+// -----
+
+#sparse1d = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+// expected-error at +1 {{invalid primary type}}
+func.func @alloc_tensor_i1_elem_type(%sz: index) -> tensor<?xi1, #sparse1d> {
+ %0 = bufferization.alloc_tensor(%sz) : tensor<?xi1, #sparse1d>
+ return %0 : tensor<?xi1, #sparse1d>
+}
+
+// -----
+
+#unordered_coo = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed(nonunique, nonordered), d1 : singleton(nonordered))}>
+#ordered_coo = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton)}>
+
+// expected-error at +1 {{invalid primary type}}
+func.func @reorder_coo_i1_elem_type(%arg0: tensor<?x?xi1, #unordered_coo>)
+ -> tensor<?x?xi1, #ordered_coo> {
+ %0 = sparse_tensor.reorder_coo quick_sort %arg0
+ : tensor<?x?xi1, #unordered_coo> to tensor<?x?xi1, #ordered_coo>
+ return %0 : tensor<?x?xi1, #ordered_coo>
+}
+
+// -----
+
+#sparse1d = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+// expected-error at +2 {{invalid primary type}}
+func.func @assemble_i1_elem_type(%pos: tensor<2xindex>, %coords: tensor<4x1xindex>,
+ %vals: tensor<4xi1>) -> tensor<8xi1, #sparse1d> {
+ %0 = sparse_tensor.assemble (%pos, %coords), %vals
+ : (tensor<2xindex>, tensor<4x1xindex>), tensor<4xi1> to tensor<8xi1, #sparse1d>
+ return %0 : tensor<8xi1, #sparse1d>
+}
+
+// -----
+
+#sparse1d = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+// expected-error at +1 {{invalid primary type}}
+func.func @values_i1_elem_type(%arg0: tensor<8xi1, #sparse1d>) -> memref<?xi1> {
+ %0 = sparse_tensor.values %arg0 : tensor<8xi1, #sparse1d> to memref<?xi1>
+ return %0 : memref<?xi1>
+}
+
+// -----
+
+#sparse1d = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+// expected-error at +1 {{invalid primary type}}
+func.func @number_of_entries_i1_elem_type(%arg0: tensor<8xi1, #sparse1d>) -> index {
+ %0 = sparse_tensor.number_of_entries %arg0 : tensor<8xi1, #sparse1d>
+ return %0 : index
+}
+
+// -----
+
+#sparse1d = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+// expected-error at +1 {{invalid primary type}}
+func.func @disassemble_i1_elem_type(%sp: tensor<8xi1, #sparse1d>,
+ %op: tensor<2xindex>, %oi: tensor<4x1xindex>,
+ %od: tensor<4xi1>)
+ -> (tensor<2xindex>, tensor<4x1xindex>, tensor<4xi1>) {
+ %rp, %ri, %d, %rpl, %ril, %dl = sparse_tensor.disassemble %sp
+ : tensor<8xi1, #sparse1d>
+ out_lvls(%op, %oi : tensor<2xindex>, tensor<4x1xindex>)
+ out_vals(%od : tensor<4xi1>)
+ -> (tensor<2xindex>, tensor<4x1xindex>), tensor<4xi1>, (index, index), index
+ return %rp, %ri, %d : tensor<2xindex>, tensor<4x1xindex>, tensor<4xi1>
+}
+
+// -----
+
+#sparse1d = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+// expected-error at +1 {{invalid primary type}}
+func.func @insert_i1_elem_type(%arg0: tensor<8xi1, #sparse1d>, %idx: index, %val: i1)
+ -> tensor<8xi1, #sparse1d> {
+ %0 = tensor.insert %val into %arg0[%idx] : tensor<8xi1, #sparse1d>
+ return %0 : tensor<8xi1, #sparse1d>
+}
+
+// -----
+
+#csr = #sparse_tensor.encoding<{map = (d0, d1) -> (d0 : dense, d1 : compressed)}>
+
+// expected-error at +1 {{invalid primary type}}
+func.func @compress_i1_elem_type(%tensor: tensor<8x8xi1, #csr>,
+ %values: memref<?xi1>, %filled: memref<?xi1>,
+ %added: memref<?xindex>, %count: index, %i: index)
+ -> tensor<8x8xi1, #csr> {
+ %0 = sparse_tensor.compress %values, %filled, %added, %count into %tensor[%i]
+ : memref<?xi1>, memref<?xi1>, memref<?xindex>, tensor<8x8xi1, #csr>
+ return %0 : tensor<8x8xi1, #csr>
+}
diff --git a/mlir/test/Dialect/SparseTensor/vectorize_reduction.mlir b/mlir/test/Dialect/SparseTensor/vectorize_reduction.mlir
index 6effbbf98abb7..5eb428f197a75 100644
--- a/mlir/test/Dialect/SparseTensor/vectorize_reduction.mlir
+++ b/mlir/test/Dialect/SparseTensor/vectorize_reduction.mlir
@@ -8,52 +8,52 @@
// Check that we vectorize reductions with ori.
// CHECK-ON-LABEL: func.func @sparse_reduction_ori(
-// CHECK-ON-SAME: %[[VAL_0:.*]]: tensor<i13>,
-// CHECK-ON-SAME: %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
+// CHECK-ON-SAME: %[[VAL_0:.*]]: tensor<i32>,
+// CHECK-ON-SAME: %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
// CHECK-ON-DAG: %[[VAL_2:.*]] = arith.constant 8 : index
-// CHECK-ON-DAG: %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi13>
+// CHECK-ON-DAG: %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi32>
// CHECK-ON-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-ON-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
-// CHECK-ON-DAG: %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
-// CHECK-ON-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
-// CHECK-ON-DAG: %[[VAL_8:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i13> to memref<i13>
-// CHECK-ON: %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i13>
+// CHECK-ON-DAG: %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
+// CHECK-ON-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
+// CHECK-ON-DAG: %[[VAL_8:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i32> to memref<i32>
+// CHECK-ON: %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i32>
// CHECK-ON: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-ON: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK-ON: %[[VAL_12:.*]] = vector.broadcast %[[VAL_9]] : i13 to vector<8xi13>
-// CHECK-ON: %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi13>) {
+// CHECK-ON: %[[VAL_12:.*]] = vector.broadcast %[[VAL_9]] : i32 to vector<8xi32>
+// CHECK-ON: %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi32>) {
// CHECK-ON: %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
// CHECK-ON: %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
-// CHECK-ON: %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi13>, vector<8xi1>, vector<8xi13> into vector<8xi13>
-// CHECK-ON: %[[VAL_19:.*]] = arith.ori %[[VAL_15]], %[[VAL_18]] : vector<8xi13>
-// CHECK-ON: %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi13>
-// CHECK-ON: scf.yield %[[VAL_20]] : vector<8xi13>
+// CHECK-ON: %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi32>, vector<8xi1>, vector<8xi32> into vector<8xi32>
+// CHECK-ON: %[[VAL_19:.*]] = arith.ori %[[VAL_15]], %[[VAL_18]] : vector<8xi32>
+// CHECK-ON: %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi32>
+// CHECK-ON: scf.yield %[[VAL_20]] : vector<8xi32>
// CHECK-ON: } {"Emitted from" = "linalg.generic"}
-// CHECK-ON: %[[VAL_21:.*]] = vector.reduction <or>, %[[VAL_22:.*]] : vector<8xi13> into i13
-// CHECK-ON: memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i13>
-// CHECK-ON: %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i13>
-// CHECK-ON: return %[[VAL_23]] : tensor<i13>
+// CHECK-ON: %[[VAL_21:.*]] = vector.reduction <or>, %[[VAL_22:.*]] : vector<8xi32> into i32
+// CHECK-ON: memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i32>
+// CHECK-ON: %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i32>
+// CHECK-ON: return %[[VAL_23]] : tensor<i32>
// CHECK-ON: }
//
// CHECK-OFF-LABEL: func.func @sparse_reduction_ori(
-// CHECK-OFF-SAME: %[[VAL_0:.*]]: tensor<i13>,
-// CHECK-OFF-SAME: %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
+// CHECK-OFF-SAME: %[[VAL_0:.*]]: tensor<i32>,
+// CHECK-OFF-SAME: %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
// CHECK-OFF-DAG: %[[VAL_2:.*]] = arith.constant 0 : index
// CHECK-OFF-DAG: %[[VAL_3:.*]] = arith.constant 1 : index
-// CHECK-OFF-DAG: %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
-// CHECK-OFF-DAG: %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
-// CHECK-OFF-DAG: %[[VAL_6:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i13> to memref<i13>
-// CHECK-OFF: %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i13>
+// CHECK-OFF-DAG: %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
+// CHECK-OFF-DAG: %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
+// CHECK-OFF-DAG: %[[VAL_6:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i32> to memref<i32>
+// CHECK-OFF: %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i32>
// CHECK-OFF: %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
// CHECK-OFF: %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK-OFF: %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i13) {
-// CHECK-OFF: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi13>
-// CHECK-OFF: %[[VAL_14:.*]] = arith.ori %[[VAL_12]], %[[VAL_13]] : i13
-// CHECK-OFF: scf.yield %[[VAL_14]] : i13
+// CHECK-OFF: %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i32) {
+// CHECK-OFF: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi32>
+// CHECK-OFF: %[[VAL_14:.*]] = arith.ori %[[VAL_12]], %[[VAL_13]] : i32
+// CHECK-OFF: scf.yield %[[VAL_14]] : i32
// CHECK-OFF: } {"Emitted from" = "linalg.generic"}
-// CHECK-OFF: memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i13>
-// CHECK-OFF: %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i13>
-// CHECK-OFF: return %[[VAL_16]] : tensor<i13>
+// CHECK-OFF: memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i32>
+// CHECK-OFF: %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i32>
+// CHECK-OFF: return %[[VAL_16]] : tensor<i32>
// CHECK-OFF: }
#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
@@ -65,17 +65,17 @@
iterator_types = ["reduction"]
}
-func.func @sparse_reduction_ori(%argx: tensor<i13>,
- %arga: tensor<?xi13, #SparseVector>)
- -> tensor<i13> {
+func.func @sparse_reduction_ori(%argx: tensor<i32>,
+ %arga: tensor<?xi32, #SparseVector>)
+ -> tensor<i32> {
%0 = linalg.generic #trait
- ins(%arga: tensor<?xi13, #SparseVector>)
- outs(%argx: tensor<i13>) {
- ^bb(%a: i13, %x: i13):
- %t = arith.ori %x, %a: i13
- linalg.yield %t : i13
- } -> tensor<i13>
- return %0 : tensor<i13>
+ ins(%arga: tensor<?xi32, #SparseVector>)
+ outs(%argx: tensor<i32>) {
+ ^bb(%a: i32, %x: i32):
+ %t = arith.ori %x, %a: i32
+ linalg.yield %t : i32
+ } -> tensor<i32>
+ return %0 : tensor<i32>
}
// -----
@@ -85,52 +85,52 @@ func.func @sparse_reduction_ori(%argx: tensor<i13>,
// irrespective to where the accumulator appears on commutative operations.
// CHECK-ON-LABEL: func.func @sparse_reduction_ori_accumulator_on_rhs(
-// CHECK-ON-SAME: %[[VAL_0:.*]]: tensor<i13>,
-// CHECK-ON-SAME: %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
+// CHECK-ON-SAME: %[[VAL_0:.*]]: tensor<i32>,
+// CHECK-ON-SAME: %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
// CHECK-ON-DAG: %[[VAL_2:.*]] = arith.constant 8 : index
-// CHECK-ON-DAG: %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi13>
+// CHECK-ON-DAG: %[[VAL_3:.*]] = arith.constant dense<0> : vector<8xi32>
// CHECK-ON-DAG: %[[VAL_4:.*]] = arith.constant 0 : index
// CHECK-ON-DAG: %[[VAL_5:.*]] = arith.constant 1 : index
-// CHECK-ON-DAG: %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
-// CHECK-ON-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
-// CHECK-ON-DAG: %[[VAL_8:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i13> to memref<i13>
-// CHECK-ON: %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i13>
+// CHECK-ON-DAG: %[[VAL_6:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
+// CHECK-ON-DAG: %[[VAL_7:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
+// CHECK-ON-DAG: %[[VAL_8:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i32> to memref<i32>
+// CHECK-ON: %[[VAL_9:.*]] = memref.load %[[VAL_8]][] : memref<i32>
// CHECK-ON: %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_4]]] : memref<?xindex>
// CHECK-ON: %[[VAL_11:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_5]]] : memref<?xindex>
-// CHECK-ON: %[[VAL_12:.*]] = vector.broadcast %[[VAL_9]] : i13 to vector<8xi13>
-// CHECK-ON: %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi13>) {
+// CHECK-ON: %[[VAL_12:.*]] = vector.broadcast %[[VAL_9]] : i32 to vector<8xi32>
+// CHECK-ON: %[[VAL_13:.*]] = scf.for %[[VAL_14:.*]] = %[[VAL_10]] to %[[VAL_11]] step %[[VAL_2]] iter_args(%[[VAL_15:.*]] = %[[VAL_12]]) -> (vector<8xi32>) {
// CHECK-ON: %[[VAL_16:.*]] = affine.min #map(%[[VAL_11]], %[[VAL_14]]){{\[}}%[[VAL_2]]]
// CHECK-ON: %[[VAL_17:.*]] = vector.create_mask %[[VAL_16]] : vector<8xi1>
-// CHECK-ON: %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi13>, vector<8xi1>, vector<8xi13> into vector<8xi13>
-// CHECK-ON: %[[VAL_19:.*]] = arith.ori %[[VAL_18]], %[[VAL_15]] : vector<8xi13>
-// CHECK-ON: %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi13>
-// CHECK-ON: scf.yield %[[VAL_20]] : vector<8xi13>
+// CHECK-ON: %[[VAL_18:.*]] = vector.maskedload %[[VAL_7]]{{\[}}%[[VAL_14]]], %[[VAL_17]], %[[VAL_3]] : memref<?xi32>, vector<8xi1>, vector<8xi32> into vector<8xi32>
+// CHECK-ON: %[[VAL_19:.*]] = arith.ori %[[VAL_18]], %[[VAL_15]] : vector<8xi32>
+// CHECK-ON: %[[VAL_20:.*]] = arith.select %[[VAL_17]], %[[VAL_19]], %[[VAL_15]] : vector<8xi1>, vector<8xi32>
+// CHECK-ON: scf.yield %[[VAL_20]] : vector<8xi32>
// CHECK-ON: } {"Emitted from" = "linalg.generic"}
-// CHECK-ON: %[[VAL_21:.*]] = vector.reduction <or>, %[[VAL_22:.*]] : vector<8xi13> into i13
-// CHECK-ON: memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i13>
-// CHECK-ON: %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i13>
-// CHECK-ON: return %[[VAL_23]] : tensor<i13>
+// CHECK-ON: %[[VAL_21:.*]] = vector.reduction <or>, %[[VAL_22:.*]] : vector<8xi32> into i32
+// CHECK-ON: memref.store %[[VAL_21]], %[[VAL_8]][] : memref<i32>
+// CHECK-ON: %[[VAL_23:.*]] = bufferization.to_tensor %[[VAL_8]] : memref<i32>
+// CHECK-ON: return %[[VAL_23]] : tensor<i32>
// CHECK-ON: }
//
// CHECK-OFF-LABEL: func.func @sparse_reduction_ori_accumulator_on_rhs(
-// CHECK-OFF-SAME: %[[VAL_0:.*]]: tensor<i13>,
-// CHECK-OFF-SAME: %[[VAL_1:.*]]: tensor<?xi13, #sparse{{[0-9]*}}>) -> tensor<i13> {
+// CHECK-OFF-SAME: %[[VAL_0:.*]]: tensor<i32>,
+// CHECK-OFF-SAME: %[[VAL_1:.*]]: tensor<?xi32, #sparse{{[0-9]*}}>) -> tensor<i32> {
// CHECK-OFF-DAG: %[[VAL_2:.*]] = arith.constant 0 : index
// CHECK-OFF-DAG: %[[VAL_3:.*]] = arith.constant 1 : index
-// CHECK-OFF-DAG: %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xindex>
-// CHECK-OFF-DAG: %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi13, #sparse{{[0-9]*}}> to memref<?xi13>
-// CHECK-OFF-DAG: %[[VAL_6:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i13> to memref<i13>
-// CHECK-OFF: %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i13>
+// CHECK-OFF-DAG: %[[VAL_4:.*]] = sparse_tensor.positions %[[VAL_1]] {level = 0 : index} : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xindex>
+// CHECK-OFF-DAG: %[[VAL_5:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?xi32, #sparse{{[0-9]*}}> to memref<?xi32>
+// CHECK-OFF-DAG: %[[VAL_6:.*]] = bufferization.to_buffer %[[VAL_0]] : tensor<i32> to memref<i32>
+// CHECK-OFF: %[[VAL_7:.*]] = memref.load %[[VAL_6]][] : memref<i32>
// CHECK-OFF: %[[VAL_8:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_2]]] : memref<?xindex>
// CHECK-OFF: %[[VAL_9:.*]] = memref.load %[[VAL_4]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK-OFF: %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i13) {
-// CHECK-OFF: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi13>
-// CHECK-OFF: %[[VAL_14:.*]] = arith.ori %[[VAL_13]], %[[VAL_12]] : i13
-// CHECK-OFF: scf.yield %[[VAL_14]] : i13
+// CHECK-OFF: %[[VAL_10:.*]] = scf.for %[[VAL_11:.*]] = %[[VAL_8]] to %[[VAL_9]] step %[[VAL_3]] iter_args(%[[VAL_12:.*]] = %[[VAL_7]]) -> (i32) {
+// CHECK-OFF: %[[VAL_13:.*]] = memref.load %[[VAL_5]]{{\[}}%[[VAL_11]]] : memref<?xi32>
+// CHECK-OFF: %[[VAL_14:.*]] = arith.ori %[[VAL_13]], %[[VAL_12]] : i32
+// CHECK-OFF: scf.yield %[[VAL_14]] : i32
// CHECK-OFF: } {"Emitted from" = "linalg.generic"}
-// CHECK-OFF: memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i13>
-// CHECK-OFF: %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i13>
-// CHECK-OFF: return %[[VAL_16]] : tensor<i13>
+// CHECK-OFF: memref.store %[[VAL_15:.*]], %[[VAL_6]][] : memref<i32>
+// CHECK-OFF: %[[VAL_16:.*]] = bufferization.to_tensor %[[VAL_6]] : memref<i32>
+// CHECK-OFF: return %[[VAL_16]] : tensor<i32>
// CHECK-OFF: }
#SparseVector = #sparse_tensor.encoding<{map = (d0) -> (d0 : compressed)}>
@@ -142,17 +142,17 @@ func.func @sparse_reduction_ori(%argx: tensor<i13>,
iterator_types = ["reduction"]
}
-func.func @sparse_reduction_ori_accumulator_on_rhs(%argx: tensor<i13>,
- %arga: tensor<?xi13, #SparseVector>)
- -> tensor<i13> {
+func.func @sparse_reduction_ori_accumulator_on_rhs(%argx: tensor<i32>,
+ %arga: tensor<?xi32, #SparseVector>)
+ -> tensor<i32> {
%0 = linalg.generic #trait
- ins(%arga: tensor<?xi13, #SparseVector>)
- outs(%argx: tensor<i13>) {
- ^bb(%a: i13, %x: i13):
- %t = arith.ori %a, %x: i13
- linalg.yield %t : i13
- } -> tensor<i13>
- return %0 : tensor<i13>
+ ins(%arga: tensor<?xi32, #SparseVector>)
+ outs(%argx: tensor<i32>) {
+ ^bb(%a: i32, %x: i32):
+ %t = arith.ori %a, %x: i32
+ linalg.yield %t : i32
+ } -> tensor<i32>
+ return %0 : tensor<i32>
}
// -----
>From e31a1ea800b60840e81ea3cdc51abcc8d6f1d6b0 Mon Sep 17 00:00:00 2001
From: Vito Secona <secona00 at gmail.com>
Date: Mon, 20 Jul 2026 21:11:23 +0700
Subject: [PATCH 2/2] docs: add docs regarding supported types
---
.../mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index 46fdff08466c1..b13a9928ae29d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -126,6 +126,10 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
encoding. Compiler passes that run before this sparsifier pass need to be aware
of the semantics of tensor types with such a sparsity encoding.
+ This encoding can only be applied to tensors with supported primary element
+ types (e.g., `f16`, `f32`, `f64`, `bf16`, `i8`, `i16`, `i32`, `i64`, and
+ complex variants of `f32` and `f64`).
+
In this encoding, we use **dimension** to refer to the axes of the semantic tensor,
and **level** to refer to the axes of the actual storage format, i.e., the
operational representation of the sparse tensor in memory. The number of
More information about the Mlir-commits
mailing list