diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td index 29ce9efc2e98b6..2c200fe08b1056 100644 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td @@ -493,6 +493,8 @@ class GenericOpBase : LinalgStructuredBase_Op:$doc, OptionalAttr:$library_call, + // ArrayAttr of StrArrayAttr: + OptionalAttr:$sparse, Confined, [IntMinValue<0>]> :$symbol_source); let results = (outs Variadic:$result_tensors); @@ -549,6 +551,8 @@ def GenericOp : GenericOpBase<"generic"> { Each element of the list represents and iterator of one of the following types: parallel, reduction, window + - sparse: an optional list with per-dimension sparsity annotations (either + "D" for dense or "S" for sparse) for each input and output view. - symbol_source: index of the operand whose dimensions will be propagated as symbols to the indexing maps. When specified the number of symbols in each of the indexing maps has to be either 0 or the rank of the diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td index 85e0e3c9f56a98..6646964a983e7d 100644 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td @@ -678,6 +678,18 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> { llvm::all_of(this->getOperation()->getResults(), isTensorType); }] >, + InterfaceMethod< + /*desc=*/[{ + Return whether the op has sparse tensor semantics. + }], + /*retTy=*/"bool", + /*methodName=*/"hasSparseSemantics", + /*args=*/(ins), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + return $_op.getAttr(getSparseAttrName()).template dyn_cast_or_null() != nullptr; + }] + >, InterfaceMethod< /*desc=*/[{ Return the name registered for this op when lowering to an external diff --git a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h index 805db03330748c..21c311ed6d1810 100644 --- a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h +++ b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h @@ -66,6 +66,9 @@ constexpr StringRef getDocAttrName() { return "doc"; } /// function that implements the structured op. constexpr StringRef getLibraryCallAttrName() { return "library_call"; } +/// Attribute name for the ArrayAttr of StrArrayAttr that encodes sparsity. +constexpr StringRef getSparseAttrName() { return "sparse"; } + /// Attribute name for the StrArrayAttr which encodes the value of strides. constexpr StringRef getStridesAttrName() { return "strides"; } @@ -134,6 +137,18 @@ inline StringRef toString(IteratorType t) { llvm_unreachable("Unsupported IteratorType"); } +/// Use to encode a dense or sparse dimension. +constexpr StringRef getSparseDimName() { return "S"; } +inline bool isSparseDim(Attribute attr) { + auto strAttr = attr.dyn_cast_or_null(); + return strAttr && strAttr.getValue() == getSparseDimName(); +} +constexpr StringRef getDenseDimName() { return "D"; } +inline bool isDenseDim(Attribute attr) { + auto strAttr = attr.dyn_cast_or_null(); + return strAttr && strAttr.getValue() == getDenseDimName(); +} + } // end namespace mlir #endif // MLIR_UTILS_STRUCTUREDOPSUTILS_H diff --git a/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp b/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp index 366aa0fdcc5aa4..11ac845f0ec71d 100644 --- a/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp +++ b/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp @@ -69,6 +69,7 @@ Operation *mlir::edsc::makeGenericLinalgOp( builder.getStrArrayAttr(iteratorStrTypes), StringAttr() /*doc*/, StringAttr() /*library_call*/, + ArrayAttr() /*sparse*/, IntegerAttr() /*symbol_source*/ /* TODO: other attributes in op */ ) diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp index 4e7fef14955130..2cd52029832db7 100644 --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -110,7 +110,7 @@ void GenericOp::build( builder.getStrArrayAttr(iteratorTypes), doc.empty() ? StringAttr() : builder.getStringAttr(doc), libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall), - symbolSource); + ArrayAttr(), symbolSource); if (!bodyBuild) return; @@ -170,7 +170,7 @@ void IndexedGenericOp::build( builder.getStrArrayAttr(iteratorTypes), doc.empty() ? StringAttr() : builder.getStringAttr(doc), libraryCall.empty() ? StringAttr() : builder.getStringAttr(libraryCall), - symbolSource); + ArrayAttr(), symbolSource); if (!bodyBuild) return; @@ -349,6 +349,7 @@ void IndexedGenericOp::getEffects( } namespace { + template struct BlockArgsVerifier { static LogicalResult verify(GenericOpType op, Block &block); @@ -405,6 +406,48 @@ LogicalResult BlockArgsVerifier::verify(IndexedGenericOp op, } return success(); } + +template +struct AnnotationsVerifier { + static LogicalResult verify(GenericOpType op) { return success(); } +}; + +template <> +LogicalResult AnnotationsVerifier::verify(GenericOp op) { + ArrayAttr sparseAttr = op.sparseAttr(); + if (!sparseAttr) + return success(); + // Verify consistency of sparse annotations. + if (!op.hasTensorSemantics()) + return op.emitOpError("expected sparse annotations on tensors only"); + unsigned numTensors = op.getNumInputsAndOutputs(); + if (sparseAttr.size() != numTensors) + return op.emitOpError("expected one sparse annotation for each tensor"); + for (unsigned t = 0; t < numTensors; t++) { + auto dimAttr = sparseAttr[t].dyn_cast_or_null(); + if (!dimAttr) + return op.emitOpError("expected sparse annotation array for tensor ") + << t; + unsigned rank = op.getShapedType(t).getRank(); + if (dimAttr.size() != rank) + return op.emitOpError("expected sparse annotation with rank ") + << rank << " for tensor " << t; + // Per-dimension annotations for each tensor consist of only "D" or "S". + for (unsigned d = 0; d < rank; d++) { + if (isDenseDim(dimAttr[d])) { + continue; + } else if (isSparseDim(dimAttr[d])) { + if (t == numTensors - 1) + return op.emitOpError("sparse output tensors not supported (yet)"); + continue; + } + return op.emitOpError("expected sparse annotation at position ") + << d << " for tensor " << t; + } + } + return success(); +} + } // namespace template @@ -466,6 +509,9 @@ static LogicalResult verifyGenericOp(GenericOpType op) { return op.emitOpError("expected the concatenation of maps in indexing_map " "to be invertible"); + if (failed(AnnotationsVerifier::verify(op))) + return failure(); + return success(); } diff --git a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp index f0e1de7094f12d..3672b80730b8c0 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp @@ -131,7 +131,8 @@ static void finalizeBufferAllocation(ConversionPatternRewriter &rewriter, /*outputBuffers=*/outputs, /*initTensors=*/llvm::None, genericOp.indexing_maps(), genericOp.iterator_types(), genericOp.docAttr(), - genericOp.library_callAttr(), genericOp.symbol_sourceAttr()); + genericOp.library_callAttr(), genericOp.sparseAttr(), + genericOp.symbol_sourceAttr()); // Create a new block in the region of the new Generic Op. Block *oldBlock = genericOp.getBody(); diff --git a/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp b/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp index 3e3392d849754b..7cb9bb5b13bfcb 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp @@ -227,6 +227,7 @@ fuseTensorOpsImpl(LinalgOp producer, LinalgOp consumer, unsigned consumerIdx, consumer.iterator_types(), /*doc=*/nullptr, /*library_call=*/nullptr, + /*sparse=*/nullptr, /*symbol_source=*/nullptr) .getOperation(); } else { @@ -241,6 +242,7 @@ fuseTensorOpsImpl(LinalgOp producer, LinalgOp consumer, unsigned consumerIdx, consumer.iterator_types(), /*doc=*/nullptr, /*library_call=*/nullptr, + /*sparse=*/nullptr, /*symbol_source=*/nullptr) .getOperation(); } @@ -820,6 +822,7 @@ struct FoldConsumerReshapeOpByLinearization producer.iterator_types(), /*doc=*/nullptr, /*library_call=*/nullptr, + /*sparse=*/nullptr, /*symbol_source=*/nullptr); auto &fusedRegion = fusedOp.getOperation()->getRegion(0); rewriter.cloneRegionBefore(producer.getOperation()->getRegion(0), @@ -903,6 +906,7 @@ struct FoldSplatConstants : public OpRewritePattern { linalgOp.iterator_types(), /*doc=*/nullptr, /*library_call=*/nullptr, + /*sparse=*/nullptr, /*symbol_source=*/nullptr); // Map the block argument corresponding to the replaced argument with the diff --git a/mlir/test/Dialect/Linalg/sparse_invalid.mlir b/mlir/test/Dialect/Linalg/sparse_invalid.mlir new file mode 100644 index 00000000000000..985667ce433b6e --- /dev/null +++ b/mlir/test/Dialect/Linalg/sparse_invalid.mlir @@ -0,0 +1,172 @@ +// RUN: mlir-opt %s -split-input-file -verify-diagnostics + +#trait_memref = { + indexing_maps = [ + affine_map<(i) -> (i)>, // a + affine_map<(i) -> (i)> // x (out) + ], + sparse = [ + [ "S" ], // a + [ "D" ] // x + ], + iterator_types = ["parallel"] +} + +func @invalid_memref(%arga: memref<32xf32>, %argb: f32) -> tensor<32xf32> { + // expected-error@+1 {{'linalg.generic' op expected sparse annotations on tensors only}} + %0 = linalg.generic #trait_memref + ins(%arga: memref<32xf32>) { + ^bb(%a: f32): + %0 = addf %a, %argb : f32 + linalg.yield %0 : f32 + } -> tensor<32xf32> + return %0 : tensor<32xf32> +} + +// ----- + +#trait_too_many = { + indexing_maps = [ + affine_map<(i) -> (i)>, // a + affine_map<(i) -> (i)> // x (out) + ], + sparse = [ + [ "S" ], // a + [ "S" ], // b + [ "D" ] // x + ], + iterator_types = ["parallel"] +} + +func @invalid_too_many(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> { + // expected-error@+1 {{'linalg.generic' op expected one sparse annotation for each tensor}} + %0 = linalg.generic #trait_too_many + ins(%arga: tensor<32xf32>) { + ^bb(%a: f32): + %0 = addf %a, %argb : f32 + linalg.yield %0 : f32 + } -> tensor<32xf32> + return %0 : tensor<32xf32> +} + +// ----- + +#trait_no_array = { + indexing_maps = [ + affine_map<(i) -> (i)>, // a + affine_map<(i) -> (i)> // x (out) + ], + sparse = [ 1, 2 ], + iterator_types = ["parallel"] +} + +func @invalid_no_array(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> { + // expected-error@+1 {{'linalg.generic' op expected sparse annotation array for tensor 0}} + %0 = linalg.generic #trait_no_array + ins(%arga: tensor<32xf32>) { + ^bb(%a: f32): + %0 = addf %a, %argb : f32 + linalg.yield %0 : f32 + } -> tensor<32xf32> + return %0 : tensor<32xf32> +} + +// ----- + +#trait_wrong_rank = { + indexing_maps = [ + affine_map<(i) -> (i)>, // a + affine_map<(i) -> (i)> // x (out) + ], + sparse = [ + [ "S" ], + [ "D", "D" ] + ], + iterator_types = ["parallel"] +} + +func @invalid_wrong_rank(%arga: tensor<32xf32>, %argb: f32) -> tensor<32xf32> { + // expected-error@+1 {{'linalg.generic' op expected sparse annotation with rank 1 for tensor 1}} + %0 = linalg.generic #trait_wrong_rank + ins(%arga: tensor<32xf32>) { + ^bb(%a: f32): + %0 = addf %a, %argb : f32 + linalg.yield %0 : f32 + } -> tensor<32xf32> + return %0 : tensor<32xf32> +} + +// ----- + +#trait_no_string = { + indexing_maps = [ + affine_map<(i,j) -> (i,j)>, // a + affine_map<(i,j) -> (i,j)> // x (out) + ], + sparse = [ + [ "S", 1 ], + [ "D", "D" ] + ], + iterator_types = ["parallel","parallel"] +} + +func @invalid_no_string(%arga: tensor<32x16xf32>, %argb: f32) -> tensor<32x16xf32> { + // expected-error@+1 {{'linalg.generic' op expected sparse annotation at position 1 for tensor 0}} + %0 = linalg.generic #trait_no_string + ins(%arga: tensor<32x16xf32>) { + ^bb(%a: f32): + %0 = addf %a, %argb : f32 + linalg.yield %0 : f32 + } -> tensor<32x16xf32> + return %0 : tensor<32x16xf32> +} + +// ----- + +#trait_wrong_symbol = { + indexing_maps = [ + affine_map<(i,j) -> (i,j)>, // a + affine_map<(i,j) -> (i,j)> // x (out) + ], + sparse = [ + [ "S", "S" ], + [ "D", "X" ] + ], + iterator_types = ["parallel","parallel"] +} + +func @invalid_wrong_symbol(%arga: tensor<32x16xf32>, %argb: f32) -> tensor<32x16xf32> { + // expected-error@+1 {{'linalg.generic' op expected sparse annotation at position 1 for tensor 1}} + %0 = linalg.generic #trait_wrong_symbol + ins(%arga: tensor<32x16xf32>) { + ^bb(%a: f32): + %0 = addf %a, %argb : f32 + linalg.yield %0 : f32 + } -> tensor<32x16xf32> + return %0 : tensor<32x16xf32> +} + +// ----- + +#trait_no_sparse_output = { + indexing_maps = [ + affine_map<(i,j) -> (i,j)>, // a + affine_map<(i,j) -> (i,j)> // x (out) + ], + sparse = [ + [ "S", "S" ], + [ "D", "S" ] + ], + iterator_types = ["parallel","parallel"] +} + +func @invalid_no_sparse_output(%arga: tensor<32x16xf32>, %argb: f32) -> tensor<32x16xf32> { + // expected-error@+1 {{'linalg.generic' op sparse output tensors not supported (yet)}} + %0 = linalg.generic #trait_no_sparse_output + ins(%arga: tensor<32x16xf32>) { + ^bb(%a: f32): + %0 = addf %a, %argb : f32 + linalg.yield %0 : f32 + } -> tensor<32x16xf32> + return %0 : tensor<32x16xf32> +}