From e1dbc25ee2311b840f70905cd5f4dfafdeb5233f Mon Sep 17 00:00:00 2001 From: Aart Bik Date: Wed, 11 Nov 2020 16:37:48 -0800 Subject: [PATCH] [mlir][sparse] integrate sparse annotation into generic linalg op This CL integrates the new sparse annotations (hereto merely added as fully transparent attributes) more tightly to the generic linalg op in order to add verification of the annotations' consistency as well as to make make other passes more aware of their presence (in the long run, rewriting rules must preserve the integrity of the annotations). Reviewed By: nicolasvasilache Differential Revision: https://reviews.llvm.org/D91224 --- .../mlir/Dialect/Linalg/IR/LinalgStructuredOps.td | 4 + .../Linalg/IR/LinalgStructuredOpsInterface.td | 12 ++ .../mlir/Dialect/Utils/StructuredOpsUtils.h | 15 ++ mlir/lib/Dialect/Linalg/EDSC/Builders.cpp | 1 + mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 50 +++++- mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp | 3 +- .../Dialect/Linalg/Transforms/FusionOnTensors.cpp | 4 + mlir/test/Dialect/Linalg/sparse_invalid.mlir | 172 +++++++++++++++++++++ 8 files changed, 258 insertions(+), 3 deletions(-) create mode 100644 mlir/test/Dialect/Linalg/sparse_invalid.mlir diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td index 29ce9ef..2c200fe 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 85e0e3c..6646964 100644 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td @@ -680,6 +680,18 @@ def LinalgStructuredInterface : OpInterface<"LinalgOp"> { >, 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 library call. }], diff --git a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h index 805db03..21c311e 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 366aa0f..11ac845 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 4e7fef1..2cd5202 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 f0e1de7..3672b80 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 3e3392d8..7cb9bb5 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 0000000..985667ce --- /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> +} -- 2.7.4