[Mlir-commits] [mlir] [mlir][SparseTensor] Reject sparse encoding on non-primary types (PR #209591)
Vito Secona
llvmlistbot at llvm.org
Tue Jul 14 11:54:59 PDT 2026
https://github.com/secona created https://github.com/llvm/llvm-project/pull/209591
SparseTensorConversion rewrite patterns assume element types are valid primary types, but the dialect does not actually enforce this restriction. This change makes it so that sparse tensor encodings gracefully reject invalid primary types by verifying the encoding.
>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] [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>
}
// -----
More information about the Mlir-commits
mailing list