diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td @@ -493,6 +493,8 @@ ArrayAttr:$iterator_types, OptionalAttr:$doc, OptionalAttr:$library_call, + // ArrayAttr of StrArrayAttr: + OptionalAttr:$sparse, Confined, [IntMinValue<0>]> :$symbol_source); let results = (outs Variadic:$result_tensors); @@ -549,6 +551,8 @@ 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 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td @@ -678,6 +678,18 @@ 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 --- a/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h +++ b/mlir/include/mlir/Dialect/Utils/StructuredOpsUtils.h @@ -66,6 +66,9 @@ /// 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 @@ 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 --- a/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp +++ b/mlir/lib/Dialect/Linalg/EDSC/Builders.cpp @@ -69,6 +69,7 @@ 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 --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -110,7 +110,7 @@ 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 @@ 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 @@ } namespace { + template struct BlockArgsVerifier { static LogicalResult verify(GenericOpType op, Block &block); @@ -405,6 +406,48 @@ } 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 @@ 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 --- a/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Bufferize.cpp @@ -131,7 +131,8 @@ /*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 --- a/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/FusionOnTensors.cpp @@ -227,6 +227,7 @@ consumer.iterator_types(), /*doc=*/nullptr, /*library_call=*/nullptr, + /*sparse=*/nullptr, /*symbol_source=*/nullptr) .getOperation(); } else { @@ -241,6 +242,7 @@ consumer.iterator_types(), /*doc=*/nullptr, /*library_call=*/nullptr, + /*sparse=*/nullptr, /*symbol_source=*/nullptr) .getOperation(); } @@ -820,6 +822,7 @@ 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 @@ 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 --- /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> +}