[mlir][sparse] add verification of absent value in sparse_tensor.unary #70248
Conversation
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
This value should always be a plain contant or something invariant computed outside the surrounding linalg operation, since there is no co-iteration defined on anything done in this branch. Fixes: llvm#69395
|
@llvm/pr-subscribers-mlir-sparse @llvm/pr-subscribers-mlir Author: Aart Bik (aartbik) ChangesThis value should always be a plain contant or something invariant computed outside the surrounding linalg operation, since there is no co-iteration defined on anything done in this branch. Fixes: Full diff: https://github.com/llvm/llvm-project/pull/70248.diff 3 Files Affected:
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
index 8c33e8651b1694e..50f5e7335dc923b 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorOps.td
@@ -624,11 +624,11 @@ def SparseTensor_InsertOp : SparseTensor_Op<"insert",
string summary = "Inserts a value into the sparse tensor";
string description = [{
Inserts the value into the underlying storage of the tensor at the
- given level-coordinates. The arity of `lvlCoords` must match the
- level-rank of the tensor. This operation can only be applied when
- the tensor materializes unintialized from a `bufferization.alloc_tensor`
- operation and the final tensor is constructed with a `load` operation
- which has the `hasInserts` attribute set.
+ given level-coordinates. The arity of `lvlCoords` must match the
+ level-rank of the tensor. This operation can only be applied when
+ the tensor materializes unintialized from a `tensor.empty` operation
+ and the final tensor is constructed with a `load` operation which
+ has the `hasInserts` attribute set.
The level-properties of the sparse tensor type fully describe what
kind of insertion order is allowed. When all levels have "unique"
@@ -974,7 +974,7 @@ def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [Pure]>,
Example of isEqual applied to intersecting elements only:
```mlir
- %C = bufferization.alloc_tensor...
+ %C = tensor.empty(...)
%0 = linalg.generic #trait
ins(%A: tensor<?xf64, #SparseVector>,
%B: tensor<?xf64, #SparseVector>)
@@ -996,7 +996,7 @@ def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [Pure]>,
Example of A+B in upper triangle, A-B in lower triangle:
```mlir
- %C = bufferization.alloc_tensor...
+ %C = tensor.empty(...)
%1 = linalg.generic #trait
ins(%A: tensor<?x?xf64, #CSR>, %B: tensor<?x?xf64, #CSR>
outs(%C: tensor<?x?xf64, #CSR> {
@@ -1029,7 +1029,7 @@ def SparseTensor_BinaryOp : SparseTensor_Op<"binary", [Pure]>,
because we never use its values, only its sparse structure:
```mlir
- %C = bufferization.alloc_tensor...
+ %C = tensor.empty(...)
%2 = linalg.generic #trait
ins(%A: tensor<?x?xf64, #CSR>, %B: tensor<?x?xi32, #CSR>
outs(%C: tensor<?x?xf64, #CSR> {
@@ -1069,7 +1069,9 @@ def SparseTensor_UnaryOp : SparseTensor_Op<"unary", [Pure]>,
Each region contains a single block describing the computation and result.
A non-empty block must end with a sparse_tensor.yield and the return type
must match the type of `output`. The primary region's block has one
- argument, while the missing region's block has zero arguments.
+ argument, while the missing region's block has zero arguments. The
+ absent region may only generate constants or values already computed
+ on entry of the `linalg.generic` operation.
A region may also be declared empty (i.e. `absent={}`), indicating that the
region does not contribute to the output.
@@ -1082,17 +1084,17 @@ def SparseTensor_UnaryOp : SparseTensor_Op<"unary", [Pure]>,
Example of A+1, restricted to existing elements:
```mlir
- %C = bufferization.alloc_tensor...
+ %C = tensor.empty(...) : tensor<?xf64, #SparseVector>
%0 = linalg.generic #trait
ins(%A: tensor<?xf64, #SparseVector>)
outs(%C: tensor<?xf64, #SparseVector>) {
^bb0(%a: f64, %c: f64) :
%result = sparse_tensor.unary %a : f64 to f64
present={
- ^bb0(%arg0: f64):
- %cf1 = arith.constant 1.0 : f64
- %ret = arith.addf %arg0, %cf1 : f64
- sparse_tensor.yield %ret : f64
+ ^bb0(%arg0: f64):
+ %cf1 = arith.constant 1.0 : f64
+ %ret = arith.addf %arg0, %cf1 : f64
+ sparse_tensor.yield %ret : f64
}
absent={}
linalg.yield %result : f64
@@ -1102,41 +1104,42 @@ def SparseTensor_UnaryOp : SparseTensor_Op<"unary", [Pure]>,
Example returning +1 for existing values and -1 for missing values:
```mlir
- %C = bufferization.alloc_tensor...
+ %p1 = arith.constant 1 : i32
+ %m1 = arith.constant -1 : i32
+ %C = tensor.empty(...) : tensor<?xi32, #SparseVector>
%1 = linalg.generic #trait
ins(%A: tensor<?xf64, #SparseVector>)
- outs(%C: tensor<?xf64, #SparseVector>) {
- ^bb0(%a: f64, %c: f64) :
+ outs(%C: tensor<?xi32, #SparseVector>) {
+ ^bb0(%a: f64, %c: i32) :
%result = sparse_tensor.unary %a : f64 to i32
present={
^bb0(%x: f64):
- %ret = arith.constant 1 : i32
- sparse_tensor.yield %ret : i32
- }
- absent={
- %ret = arith.constant -1 : i32
- sparse_tensor.yield %ret : i32
- }
- linalg.yield %result : f64
- } -> tensor<?xf64, #SparseVector>
+ sparse_tensor.yield %p1 : i32
+ }
+ absent={
+ sparse_tensor.yield %m1 : i32
+ }
+ linalg.yield %result : i32
+ } -> tensor<?xi32, #SparseVector>
```
Example showing a structural inversion (existing values become missing in
the output, while missing values are filled with 1):
```mlir
- %C = bufferization.alloc_tensor...
+ %c1 = arith.constant 1 : i64
+ %C = tensor.empty(...) : tensor<?xi64, #SparseVector>
%2 = linalg.generic #trait
- ins(%A: tensor<?xf64, #SparseVector>)
- outs(%C: tensor<?xf64, #SparseVector>) {
- %result = sparse_tensor.unary %a : f64 to i64
- present={}
- absent={
- %ret = arith.constant 1 : i64
- sparse_tensor.yield %ret : i64
- }
- linalg.yield %result : f64
- } -> tensor<?xf64, #SparseVector>
+ ins(%A: tensor<?xf64, #SparseVector>)
+ outs(%C: tensor<?xi64, #SparseVector>) {
+ ^bb0(%a: f64, %c: i64) :
+ %result = sparse_tensor.unary %a : f64 to i64
+ present={}
+ absent={
+ sparse_tensor.yield %c1 : i64
+ }
+ linalg.yield %result : i64
+ } -> tensor<?xi64, #SparseVector>
```
}];
@@ -1177,7 +1180,7 @@ def SparseTensor_ReduceOp : SparseTensor_Op<"reduce", [Pure, SameOperandsAndResu
```mlir
%cf1 = arith.constant 1.0 : f64
%cf100 = arith.constant 100.0 : f64
- %C = bufferization.alloc_tensor...
+ %C = tensor.empty(...)
%0 = linalg.generic #trait
ins(%A: tensor<?x?xf64, #SparseMatrix>)
outs(%C: tensor<?xf64, #SparseVector>) {
@@ -1220,7 +1223,7 @@ def SparseTensor_SelectOp : SparseTensor_Op<"select", [Pure, SameOperandsAndResu
Example of selecting A >= 4.0:
```mlir
- %C = bufferization.alloc_tensor...
+ %C = tensor.empty(...)
%0 = linalg.generic #trait
ins(%A: tensor<?xf64, #SparseVector>)
outs(%C: tensor<?xf64, #SparseVector>) {
@@ -1238,7 +1241,7 @@ def SparseTensor_SelectOp : SparseTensor_Op<"select", [Pure, SameOperandsAndResu
Example of selecting lower triangle of a matrix:
```mlir
- %C = bufferization.alloc_tensor...
+ %C = tensor.empty(...)
%1 = linalg.generic #trait
ins(%A: tensor<?x?xf64, #CSR>)
outs(%C: tensor<?x?xf64, #CSR>) {
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 17e6ef53fe596e0..f05cbd8d16d9a76 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -34,8 +34,13 @@
using namespace mlir;
using namespace mlir::sparse_tensor;
+#define RETURN_FAILURE_IF_FAILED(X) \
+ if (failed(X)) { \
+ return failure(); \
+ }
+
//===----------------------------------------------------------------------===//
-// Additional convenience methods.
+// Local convenience methods.
//===----------------------------------------------------------------------===//
static constexpr bool acceptBitWidth(unsigned bitWidth) {
@@ -52,7 +57,7 @@ static constexpr bool acceptBitWidth(unsigned bitWidth) {
}
//===----------------------------------------------------------------------===//
-// StorageLayout
+// SparseTensorDialect StorageLayout.
//===----------------------------------------------------------------------===//
static constexpr Level kInvalidLevel = -1u;
@@ -183,7 +188,7 @@ StorageLayout::getFieldIndexAndStride(SparseTensorFieldKind kind,
}
//===----------------------------------------------------------------------===//
-// TensorDialect Attribute Methods.
+// SparseTensorDialect Attribute Methods.
//===----------------------------------------------------------------------===//
std::optional<uint64_t> SparseTensorDimSliceAttr::getStatic(int64_t v) {
@@ -658,11 +663,6 @@ SparseTensorEncodingAttr::verify(function_ref<InFlightDiagnostic()> emitError,
return success();
}
-#define RETURN_FAILURE_IF_FAILED(X) \
- if (failed(X)) { \
- return failure(); \
- }
-
LogicalResult SparseTensorEncodingAttr::verifyEncoding(
ArrayRef<DynSize> dimShape, Type elementType,
function_ref<InFlightDiagnostic()> emitError) const {
@@ -685,7 +685,7 @@ LogicalResult SparseTensorEncodingAttr::verifyEncoding(
}
//===----------------------------------------------------------------------===//
-// Convenience Methods.
+// Convenience methods.
//===----------------------------------------------------------------------===//
SparseTensorEncodingAttr
@@ -1365,10 +1365,6 @@ LogicalResult SetStorageSpecifierOp::verify() {
return success();
}
-//===----------------------------------------------------------------------===//
-// TensorDialect Linalg.Generic Operations.
-//===----------------------------------------------------------------------===//
-
template <class T>
static LogicalResult verifyNumBlockArgs(T *op, Region ®ion,
const char *regionName,
@@ -1445,6 +1441,18 @@ LogicalResult UnaryOp::verify() {
if (!absent.empty()) {
RETURN_FAILURE_IF_FAILED(
verifyNumBlockArgs(this, absent, "absent", TypeRange{}, outputType))
+ // Absent branch can only yield invariant values.
+ Block *absentBlock = &absent.front();
+ Block *parent = getOperation()->getBlock();
+ Value absentVal = cast<YieldOp>(absentBlock->getTerminator()).getResult();
+ if (auto arg = dyn_cast<BlockArgument>(absentVal)) {
+ if (arg.getOwner() == parent)
+ return emitError("absent region cannot yield linalg argument");
+ } else if (Operation *def = absentVal.getDefiningOp()) {
+ if (!isa<arith::ConstantOp>(def) &&
+ (def->getBlock() == absentBlock || def->getBlock() == parent))
+ return emitError("absent region cannot yield locally computed value");
+ }
}
return success();
}
@@ -1719,10 +1727,6 @@ LogicalResult YieldOp::verify() {
#undef RETURN_FAILURE_IF_FAILED
-//===----------------------------------------------------------------------===//
-// TensorDialect Methods.
-//===----------------------------------------------------------------------===//
-
/// Materialize a single constant operation from a given attribute value with
/// the desired resultant type.
Operation *SparseTensorDialect::materializeConstant(OpBuilder &builder,
diff --git a/mlir/test/Dialect/SparseTensor/invalid.mlir b/mlir/test/Dialect/SparseTensor/invalid.mlir
index 33aa81c5a747d9b..0217ef152be6a0d 100644
--- a/mlir/test/Dialect/SparseTensor/invalid.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid.mlir
@@ -544,6 +544,57 @@ func.func @invalid_unary_wrong_yield(%arg0: f64) -> f64 {
// -----
+
+#SparseVector = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+#trait = {
+ indexing_maps = [ affine_map<(i) -> (i)>, affine_map<(i) -> (i)> ],
+ iterator_types = ["parallel"]
+}
+
+func.func @invalid_absent_value(%arg0 : tensor<100xf64, #SparseVector>) -> tensor<100xf64, #SparseVector> {
+ %C = tensor.empty() : tensor<100xf64, #SparseVector>
+ %0 = linalg.generic #trait
+ ins(%arg0: tensor<100xf64, #SparseVector>)
+ outs(%C: tensor<100xf64, #SparseVector>) {
+ ^bb0(%a: f64, %c: f64) :
+ // expected-error@+1 {{absent region cannot yield linalg argument}}
+ %result = sparse_tensor.unary %a : f64 to f64
+ present={}
+ absent={ sparse_tensor.yield %a : f64 }
+ linalg.yield %result : f64
+ } -> tensor<100xf64, #SparseVector>
+ return %0 : tensor<100xf64, #SparseVector>
+}
+
+// -----
+
+#SparseVector = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+#trait = {
+ indexing_maps = [ affine_map<(i) -> (i)>, affine_map<(i) -> (i)> ],
+ iterator_types = ["parallel"]
+}
+
+func.func @invalid_absent_computation(%arg0 : tensor<100xf64, #SparseVector>) -> tensor<100xf64, #SparseVector> {
+ %f0 = arith.constant 0.0 : f64
+ %C = tensor.empty() : tensor<100xf64, #SparseVector>
+ %0 = linalg.generic #trait
+ ins(%arg0: tensor<100xf64, #SparseVector>)
+ outs(%C: tensor<100xf64, #SparseVector>) {
+ ^bb0(%a: f64, %c: f64) :
+ %v = arith.addf %a, %f0 : f64
+ // expected-error@+1 {{absent region cannot yield locally computed value}}
+ %result = sparse_tensor.unary %a : f64 to f64
+ present={}
+ absent={ sparse_tensor.yield %v : f64 }
+ linalg.yield %result : f64
+ } -> tensor<100xf64, #SparseVector>
+ return %0 : tensor<100xf64, #SparseVector>
+}
+
+// -----
+
func.func @invalid_reduce_num_args_mismatch(%arg0: f64, %arg1: f64) -> f64 {
%cf1 = arith.constant 1.0 : f64
// expected-error@+1 {{reduce region must have exactly 2 arguments}}
|
PeimingLiu
approved these changes
Oct 25, 2023
zahiraam
pushed a commit
to zahiraam/llvm-project
that referenced
this pull request
Oct 26, 2023
llvm#70248) This value should always be a plain contant or something invariant computed outside the surrounding linalg operation, since there is no co-iteration defined on anything done in this branch. Fixes: llvm#69395
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
This value should always be a plain contant or something invariant computed outside the surrounding linalg operation, since there is no co-iteration defined on anything done in this branch.
Fixes:
#69395