diff --git a/mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt b/mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt --- a/mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt +++ b/mlir/include/mlir/Dialect/Linalg/IR/CMakeLists.txt @@ -2,7 +2,13 @@ set(LLVM_TARGET_DEFINITIONS LinalgStructuredOps.td) mlir_tablegen(LinalgStructuredOps.h.inc -gen-op-decls) mlir_tablegen(LinalgStructuredOps.cpp.inc -gen-op-defs) +add_public_tablegen_target(MLIRLinalgStructuredOpsIncGen) + +set(LLVM_TARGET_DEFINITIONS LinalgStructuredOpsInterface.td) mlir_tablegen(LinalgStructuredOpsInterfaces.h.inc -gen-op-interface-decls) mlir_tablegen(LinalgStructuredOpsInterfaces.cpp.inc -gen-op-interface-defs) -add_public_tablegen_target(MLIRLinalgStructuredOpsIncGen) +add_public_tablegen_target(MLIRLinalgStructuredOpsInterfaceIncGen) +set(LLVM_TARGET_DEFINITIONS LinalgStructuredOps.td) +mlir_tablegen(LinalgNamedStructuredOpsInterface.cpp.inc -gen-linalg-named-structured-ops-defs) +add_public_tablegen_target(MLIRLinalgNamedStructuredOpsInterfaceIncGen) 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 @@ -16,6 +16,7 @@ include "mlir/Dialect/AffineOps/AffineOpsBase.td" include "mlir/Dialect/Linalg/IR/LinalgBase.td" +include "mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td" // The Linalg `NInputs` trait provides the API for ops that are known // to have a specified number of inputs, all passed as operands. @@ -31,184 +32,6 @@ def StructuredOpTraits : NativeOpTrait<"linalg::StructuredOpTraits">; -// The linalg 'LinalgStructuredInterface' provides access to the 'LinalgOp' -// interface. -def LinalgStructuredInterface : OpInterface<"LinalgOp"> { - let methods = [ - //========================================================================// - // Loop types handling. - //========================================================================// - InterfaceMethod< - "Return the number of parallel loops within the current operation.", - "unsigned", "getNumParallelLoops" - >, - InterfaceMethod< - "Return the number of reduction loops within the current operation.", - "unsigned", "getNumReductionLoops" - >, - InterfaceMethod< - "Return the number of window loops within the current operation.", - "unsigned", "getNumWindowLoops" - >, - InterfaceMethod< - "Return the number of loops within the current operation.", - "unsigned", "getNumLoops">, - - InterfaceMethod< - [{Returns true if the current operation has only one loop and it's a - reduction loop}], - "unsigned", "hasSingleReductionLoop">, - - //========================================================================// - // Input arguments handling. - //========================================================================// - InterfaceMethod< - "Return the number of inputs from the current operation.", - "unsigned", "getNumInputs" - >, - InterfaceMethod<"Return the input view at the given index.", - "Value ", "getInput", (ins "unsigned":$i) - >, - InterfaceMethod<[{ - Return the index of the given input value `v`, or `None` if the value is - not an input. - }], - "llvm::Optional", "getIndexOfInput", (ins "Value ":$v) - >, - InterfaceMethod< - "Return the input operands from the current operation.", - "Operation::operand_range", "getInputs" - >, - InterfaceMethod<[{ - Return the type of the input shape at the given index. - }], "ShapedType", "getInputShapedType", (ins "unsigned":$i)>, - InterfaceMethod<[{ - Return the subset of input operands that are of ranked tensor type. - }], "SmallVector", "getInputTensorTypes">, - - //========================================================================// - // Output arguments handling. - //========================================================================// - InterfaceMethod< - "Return the number of outputs from the current operation.", - "unsigned", "getNumOutputs" - >, - InterfaceMethod<"Return the output buffer at the given index.", - "Value ", "getOutputBuffer", (ins "unsigned":$i) - >, - InterfaceMethod<[{ - Return the index of the given buffer value, or `None` if the value is - not part of the output buffers. - }], - "llvm::Optional", "getIndexOfOutputBuffer", (ins "Value ":$view) - >, - InterfaceMethod<[{ - Return the type of the output buffer at the given index. - }], "MemRefType", "getOutputBufferType", (ins "unsigned":$i)>, - InterfaceMethod<[{ - Return the results that are of ranked tensor type. - }], "SmallVector", "getOutputTensorTypes">, - InterfaceMethod< - "Return the output buffers (operands) from the current operation.", - "Operation::operand_range", "getOutputBuffers" - >, - - //========================================================================// - // Input and Output arguments handling. - //========================================================================// - InterfaceMethod< - "Return the number of inputs and outputs, irrespective of their buffer " - "or tensor type.", - "unsigned", "getNumInputsAndOutputs" - >, - InterfaceMethod< - "Return the number of inputs, irrespective of their buffer or tensor " - "type, and output buffers", - "unsigned", "getNumInputsAndOutputBuffers" - >, - InterfaceMethod< - "Return the range over inputs (irrespective of type) and output buffers.", - "Operation::operand_range", "getInputsAndOutputBuffers" - >, - - //========================================================================// - // Other interface methods. - //========================================================================// - InterfaceMethod< - "Return the reference iterators for this named op (if any are specied). " - "These reference iterators are used to specify the default behavior of " - "the op. Typically this would be a static method but in order to allow " - "rank-polymorphic ops, this needs to be per object instance. Named ops " - "must define referenceIterators, even if empty for the 0-D case. " - "Generic ops on the other hand have a None `referenceIterators`", - "llvm::Optional>", "referenceIterators" - >, - InterfaceMethod< - "Return the reference indexing maps for this named op (if any are " - "specified). Typically this would be a static method but in order to " - "allow rank-polymorphic ops, this needs to be per object instance. Named " - "ops must define referenceIterators, even if empty for the 0-D case. " - "Generic ops on the other hand have a None `referenceIndexingMaps`", - "llvm::Optional>", "referenceIndexingMaps" - >, - InterfaceMethod< - "Return the iterator types attribute within the current operation.", - "ArrayAttr", "iterator_types" - >, - InterfaceMethod< - "Return the indexing maps attribute within the current operation.", - "ArrayAttr", "indexing_maps" - >, - InterfaceMethod<"Return the input or output indexing map at index `i`.", - "AffineMap", "getIndexingMap", (ins "unsigned":$i) - >, - InterfaceMethod<"Return the input indexing map at index `i`.", - "AffineMap", "getInputIndexingMap", (ins "unsigned":$i) - >, - InterfaceMethod<"Return the output indexing map at index `i`.", - "AffineMap", "getOutputIndexingMap", (ins "unsigned":$i) - >, - InterfaceMethod<[{ - Return whether the op has only MemRef input and outputs. - }], "bool", "hasBufferSemantics">, - InterfaceMethod<[{ - Return whether the op has only RankedTensor input and outputs. - }], "bool", "hasTensorSemantics">, - - //========================================================================// - // Other static interface methods. - //========================================================================// - StaticInterfaceMethod<[{ - Create an operation of the current type with the given location, - operands, and attributes. - }], - "Operation *", "create", - (ins "OpBuilder &":$builder, "Location":$loc, - "ValueRange":$operands, - "ArrayRef":$attributes), [{ - return builder.create(loc, ArrayRef{}, operands, - attributes); - }] - >, - InterfaceMethod<[{ - Clone the current operation with the given location and operands. This - is used to abstract away the optional underlying region creation. - }], - "Operation *", "clone", - (ins "OpBuilder &":$b, "Location":$loc, "ValueRange":$operands), [{ - BlockAndValueMapping map; - unsigned numRegions = op.getOperation()->getNumRegions(); - Operation *res = create(b, loc, operands, op.getAttrs()); - assert(res->getNumRegions() == numRegions && "inconsistent # regions"); - for (unsigned ridx = 0; ridx < numRegions; ++ridx) - op.getOperation()->getRegion(ridx).cloneInto( - &res->getRegion(ridx), map); - return res; - }] - > - ]; -} - // Base Tablegen class for Linalg ops. // Linalg ops that correspond to library calls operate on linalg::View as their // first operands. These may be optionally followed by non-view operands @@ -229,7 +52,9 @@ } //////////////////////////////////////////////////////////////////////////////// -// Named Linalg ops, implemented as special configurations of a generic op. +// Named Linalg ops, implemented as special configurations of generic ops. +// At the moment these are not declarative and require a bunch of C++ code. +// These should be migrated to the declarative configurations as they evolve. //////////////////////////////////////////////////////////////////////////////// def CopyOp : LinalgStructured_Op<"copy", [NInputs<1>, NOutputs<1>]> { let description = [{ @@ -311,30 +136,6 @@ let hasFolder = 1; } -def FillOp : LinalgStructured_Op<"fill", [NInputs<0>, NOutputs<1>]> { - let arguments = (ins AnyStridedMemRef:$output, - AnyTypeOf<[AnyFloat, AnyInteger, AnyVector]>:$value); - let extraClassDeclaration = libraryCallName # [{ - // Defined in C++ for now. - // TODO(ntv): auto-generate. - ArrayAttr indexing_maps(); - - // Rank-polymorphic. - // filling_value -> O(ivs) with parallel iterators. - llvm::Optional> referenceIterators() { - unsigned nPar = output().getType().cast().getRank(); - return SmallVector(nPar, getParallelIteratorTypeName()); - } - - llvm::Optional> referenceIndexingMaps() { - llvm_unreachable("NYI referenceIndexingMaps for CopyOp"); - } - }]; - let verifier = [{ return ::verify(*this); }]; - - let hasFolder = 1; -} - def DotOp : LinalgStructured_Op<"dot", [NInputs<2>, NOutputs<1>]> { let arguments = (ins AnyStridedMemRefOfRank<1>, AnyStridedMemRefOfRank<1>, @@ -356,29 +157,6 @@ let hasFolder = 1; } -def MatvecOp : LinalgStructured_Op<"matvec", [NInputs<2>, NOutputs<1>]> { - let arguments = (ins AnyStridedMemRefOfRank<2>, - AnyStridedMemRefOfRank<1>, - AnyStridedMemRefOfRank<1>); - let extraClassDeclaration = libraryCallName # [{ - // Defined in C++ for now. - // TODO(ntv): auto-generate. - ArrayAttr indexing_maps(); - - llvm::Optional> referenceIterators() { - return SmallVector{ - getParallelIteratorTypeName(), - getReductionIteratorTypeName()}; - } - - llvm::Optional> referenceIndexingMaps() { - llvm_unreachable("NYI referenceIndexingMaps for MatvecOp"); - } - }]; - - let hasFolder = 1; -} - def MatmulOp : LinalgStructured_Op<"matmul", [NInputs<2>, NOutputs<1>]> { let arguments = (ins AnyStridedMemRefOfRank<2>, AnyStridedMemRefOfRank<2>, @@ -431,7 +209,9 @@ // TODO(ntv) extend to support more than 1 dimensions and potentially // grouping too. unsigned getNumBatchDimensions() { return 1; } + unsigned getNumInputFeatureDimensions() { return 1; } + unsigned getNumOutputFeatureDimensions() { return 1; } // Defined in C++ for now. @@ -480,6 +260,92 @@ let hasFolder = 1; } +//////////////////////////////////////////////////////////////////////////////// +// Named Linalg ops, implemented as special configurations of generic ops. +// At the moment these are not fully declarative special C++ code in the form of +// specifying the reference iterators and indexing maps. +// This is necessary to obtain rank-polymorphic behavior and until a better +// solution is implemented. +//////////////////////////////////////////////////////////////////////////////// + +def NamedStructuredOpTraits : NativeOpTrait<"linalg::NamedStructuredOpTraits">; +def NamedStructuredOpLibraryCallTraits : + NativeOpTrait<"linalg::NamedStructuredOpLibraryCallTraits">; + +def FillOp : LinalgStructured_Op<"fill", [ + NInputs<0>, + NOutputs<1>, + NamedStructuredOpTraits]> { + let arguments = (ins AnyStridedMemRef:$output, + AnyTypeOf<[AnyFloat, AnyInteger, AnyVector]>:$value); + let extraClassDeclaration = libraryCallName # [{ + // Rank-polymorphic. + // filling_value -> O(ivs) with parallel iterators. + llvm::Optional> referenceIterators() { + unsigned nPar = output().getType().cast().getRank(); + return SmallVector(nPar, getParallelIteratorTypeName()); + } + + llvm::Optional> referenceIndexingMaps() { + auto ctx = getContext(); + unsigned rank = output().getType().cast().getRank(); + if (rank == 0) + return SmallVector{ + AffineMap::get(0, 0, {getAffineConstantExpr(0, ctx)})}; + return SmallVector{ + AffineMap::getMultiDimIdentityMap(rank, ctx)}; + } + }]; + + let hasFolder = 1; +} + +//////////////////////////////////////////////////////////////////////////////// +// Named Linalg ops, implemented as a declarative configurations of generic ops. +// For now these ops have fixed rank which allows us to immediately specify +// their core properties declaratively. +// The big missing piece atm is the region which is currently still C++. +// In the future we will also want to automatically generate the matchers and +// rewriters to allow rewriting a linalg.generic into a named op. +//////////////////////////////////////////////////////////////////////////////// + +class AffineExpressions map> { + list indexing_maps = map; +} + +class LinalgNamedStructured_Op props> + : Op { + list iterators = ?; + list iterators_types = ?; + list input_indexing_maps = ?; + list output_indexing_maps = ?; + bit hasLibraryImpl = 0; + let assemblyFormat = "`(` operands `)` attr-dict `:` type(operands)"; +} + +def MatvecOp : LinalgNamedStructured_Op<"matvec", [ + NInputs<2>, + NOutputs<1>, + NamedStructuredOpTraits, + NamedStructuredOpLibraryCallTraits]> { + let arguments = (ins AnyStridedMemRefOfRank<2>, + AnyStridedMemRefOfRank<1>, + AnyStridedMemRefOfRank<1>); + let iterators = ["i", "r_j"]; + let iterators_types = ["parallel", "reduction"]; + // A(i, r_j) * B(r_j) -> C(i) + let input_indexing_maps = [AffineExpressions<["i", "r_j"]>, + AffineExpressions<["r_j"]>]; + let output_indexing_maps = [AffineExpressions<["i"]>]; + + let hasLibraryImpl = 1; + let hasFolder = 1; +} + +//////////////////////////////////////////////////////////////////////////////// +// Generic Linalg ops. +//////////////////////////////////////////////////////////////////////////////// def LinalgOperand: Type< Or<[AnyRankedTensor.predicate, AnyStridedMemRef.predicate]>>; @@ -489,9 +355,6 @@ CPred<"$_self.cast().getRank() == " # rank>] >>; -//////////////////////////////////////////////////////////////////////////////// -// Generic Linalg ops. -//////////////////////////////////////////////////////////////////////////////// class GenericOpBase : LinalgStructuredBase_Op { let arguments = (ins Variadic:$views, I64Attr:$args_in, diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td new file mode 100644 --- /dev/null +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOpsInterface.td @@ -0,0 +1,196 @@ +//===- LinalgStructuredInterface.td- Linalg StructuredIfce -*- tablegen -*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// This is the definition file for the structured interface for Linalg ops. +// +//===----------------------------------------------------------------------===// + +#ifndef LINALG_STRUCTURED_INTERFACE +#define LINALG_STRUCTURED_INTERFACE + +include "mlir/Dialect/Linalg/IR/LinalgBase.td" + +// The linalg 'LinalgStructuredInterface' provides access to the 'LinalgOp' +// interface. +def LinalgStructuredInterface : OpInterface<"LinalgOp"> { + let methods = [ + //========================================================================// + // Loop types handling. + //========================================================================// + InterfaceMethod< + "Return the number of parallel loops within the current operation.", + "unsigned", "getNumParallelLoops" + >, + InterfaceMethod< + "Return the number of reduction loops within the current operation.", + "unsigned", "getNumReductionLoops" + >, + InterfaceMethod< + "Return the number of window loops within the current operation.", + "unsigned", "getNumWindowLoops" + >, + InterfaceMethod< + "Return the number of loops within the current operation.", + "unsigned", "getNumLoops">, + + InterfaceMethod< + [{Returns true if the current operation has only one loop and it's a + reduction loop}], + "unsigned", "hasSingleReductionLoop">, + + //========================================================================// + // Input arguments handling. + //========================================================================// + InterfaceMethod< + "Return the number of inputs from the current operation.", + "unsigned", "getNumInputs" + >, + InterfaceMethod<"Return the input view at the given index.", + "Value ", "getInput", (ins "unsigned":$i) + >, + InterfaceMethod<[{ + Return the index of the given input value `v`, or `None` if the value is + not an input. + }], + "llvm::Optional", "getIndexOfInput", (ins "Value ":$v) + >, + InterfaceMethod< + "Return the input operands from the current operation.", + "Operation::operand_range", "getInputs" + >, + InterfaceMethod<[{ + Return the type of the input shape at the given index. + }], "ShapedType", "getInputShapedType", (ins "unsigned":$i)>, + InterfaceMethod<[{ + Return the subset of input operands that are of ranked tensor type. + }], "SmallVector", "getInputTensorTypes">, + + //========================================================================// + // Output arguments handling. + //========================================================================// + InterfaceMethod< + "Return the number of outputs from the current operation.", + "unsigned", "getNumOutputs" + >, + InterfaceMethod<"Return the output buffer at the given index.", + "Value ", "getOutputBuffer", (ins "unsigned":$i) + >, + InterfaceMethod<[{ + Return the index of the given buffer value, or `None` if the value is + not part of the output buffers. + }], + "llvm::Optional", "getIndexOfOutputBuffer", (ins "Value ":$view) + >, + InterfaceMethod<[{ + Return the type of the output buffer at the given index. + }], "MemRefType", "getOutputBufferType", (ins "unsigned":$i)>, + InterfaceMethod<[{ + Return the results that are of ranked tensor type. + }], "SmallVector", "getOutputTensorTypes">, + InterfaceMethod< + "Return the output buffers (operands) from the current operation.", + "Operation::operand_range", "getOutputBuffers" + >, + + //========================================================================// + // Input and Output arguments handling. + //========================================================================// + InterfaceMethod< + "Return the number of inputs and outputs, irrespective of their buffer " + "or tensor type.", + "unsigned", "getNumInputsAndOutputs" + >, + InterfaceMethod< + "Return the number of inputs, irrespective of their buffer or tensor " + "type, and output buffers", + "unsigned", "getNumInputsAndOutputBuffers" + >, + InterfaceMethod< + "Return the range over inputs (irrespective of type) and output buffers.", + "Operation::operand_range", "getInputsAndOutputBuffers" + >, + + //========================================================================// + // Other interface methods. + //========================================================================// + InterfaceMethod< + "Return the reference iterators for this named op (if any are specied). " + "These reference iterators are used to specify the default behavior of " + "the op. Typically this would be a static method but in order to allow " + "rank-polymorphic ops, this needs to be per object instance. Named ops " + "must define referenceIterators, even if empty for the 0-D case. " + "Generic ops on the other hand have a None `referenceIterators`", + "llvm::Optional>", "referenceIterators" + >, + InterfaceMethod< + "Return the reference indexing maps for this named op (if any are " + "specified). Typically this would be a static method but in order to " + "allow rank-polymorphic ops, this needs to be per object instance. Named " + "ops must define referenceIterators, even if empty for the 0-D case. " + "Generic ops on the other hand have a None `referenceIndexingMaps`", + "llvm::Optional>", "referenceIndexingMaps" + >, + InterfaceMethod< + "Return the iterator types attribute within the current operation.", + "ArrayAttr", "iterator_types" + >, + InterfaceMethod< + "Return the indexing maps attribute within the current operation.", + "ArrayAttr", "indexing_maps" + >, + InterfaceMethod<"Return the input or output indexing map at index `i`.", + "AffineMap", "getIndexingMap", (ins "unsigned":$i) + >, + InterfaceMethod<"Return the input indexing map at index `i`.", + "AffineMap", "getInputIndexingMap", (ins "unsigned":$i) + >, + InterfaceMethod<"Return the output indexing map at index `i`.", + "AffineMap", "getOutputIndexingMap", (ins "unsigned":$i) + >, + InterfaceMethod<[{ + Return whether the op has only MemRef input and outputs. + }], "bool", "hasBufferSemantics">, + InterfaceMethod<[{ + Return whether the op has only RankedTensor input and outputs. + }], "bool", "hasTensorSemantics">, + + //========================================================================// + // Other static interface methods. + //========================================================================// + StaticInterfaceMethod<[{ + Create an operation of the current type with the given location, + operands, and attributes. + }], + "Operation *", "create", + (ins "OpBuilder &":$builder, "Location":$loc, + "ValueRange":$operands, + "ArrayRef":$attributes), [{ + return builder.create(loc, ArrayRef{}, operands, + attributes); + }] + >, + InterfaceMethod<[{ + Clone the current operation with the given location and operands. This + is used to abstract away the optional underlying region creation. + }], + "Operation *", "clone", + (ins "OpBuilder &":$b, "Location":$loc, "ValueRange":$operands), [{ + BlockAndValueMapping map; + unsigned numRegions = op.getOperation()->getNumRegions(); + Operation *res = create(b, loc, operands, op.getAttrs()); + assert(res->getNumRegions() == numRegions && "inconsistent # regions"); + for (unsigned ridx = 0; ridx < numRegions; ++ridx) + op.getOperation()->getRegion(ridx).cloneInto( + &res->getRegion(ridx), map); + return res; + }] + > + ]; +} + +#endif // LINALG_STRUCTURED_INTERFACE diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h b/mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgTraits.h @@ -234,10 +234,11 @@ if (!maybeReferenceIteratorTypes && name != "generic" && name != "indexed_generic") { this->getOperation()->dump(); - llvm_unreachable("Op missing "); + llvm_unreachable("Op missing referenceIterators"); } - // If we have a reference, build the reference attribute. + // If we have a reference, build the reference attribute and set it in the + // op before returning. auto *ctx = this->getOperation()->getContext(); auto attrRange = llvm::map_range(*maybeReferenceIteratorTypes, [ctx](StringRef str) -> Attribute { @@ -339,6 +340,25 @@ } }; +/// This class provides the API for named Linalg StructuredOps. +template +class NamedStructuredOpTraits + : public OpTrait::TraitBase { +public: + llvm::Optional> referenceIterators(); + llvm::Optional> referenceIndexingMaps(); +}; + +/// This class provides the API mapping named Linalg StructuredOps to a library +/// call. +template +class NamedStructuredOpLibraryCallTraits + : public OpTrait::TraitBase { +public: + std::string getLibraryCallName(); +}; + } // namespace linalg } // namespace OpTrait } // namespace mlir diff --git a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt --- a/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt +++ b/mlir/lib/Dialect/Linalg/IR/CMakeLists.txt @@ -19,5 +19,7 @@ ${LIBS} MLIRLinalgOpsIncGen MLIRLinalgStructuredOpsIncGen + MLIRLinalgStructuredOpsInterfaceIncGen + MLIRLinalgNamedStructuredOpsInterfaceIncGen ) target_link_libraries(MLIRLinalgOps ${LIBS}) 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 @@ -875,6 +875,17 @@ return AffineMap::getMultiDimIdentityMap(rank, context); } +namespace mlir { +namespace OpTrait { +namespace linalg { + +// Linalg "named" traits depend on +#include "mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.cpp.inc" + +} // namespace linalg +} // namespace OpTrait +} // namespace mlir + namespace mlir { namespace linalg { @@ -938,12 +949,6 @@ extractOrIdentityMap(maybeInputMap, inputRank, context), extractOrIdentityMap(maybeOutputMap, outputRank, context)}; } - if (auto fillOp = dyn_cast(op)) { - // filling_value -> O(ivs) - unsigned rank = fillOp.getNumParallelLoops(); - return SmallVector{ - extractOrIdentityMap(llvm::None, rank, context)}; - } auto i = getAffineDimExpr(0, context); auto j = getAffineDimExpr(1, context); auto k = getAffineDimExpr(2, context); @@ -951,11 +956,6 @@ // A(r_i) * B(r_i) -> C() return SmallVector{AffineMap::get(1, 0, {i}), AffineMap::get(1, 0, {i}), AffineMap()}; - if (isa(op)) - // A(i, r_j) * B(r_j) -> C(i) - return SmallVector{AffineMap::get(2, 0, {i, j}), - AffineMap::get(2, 0, {j}), - AffineMap::get(2, 0, {i})}; if (isa(op)) // A(i, r_k) * B(r_k, j) -> C(i, j) return SmallVector{AffineMap::get(3, 0, {i, k}), @@ -1060,15 +1060,9 @@ ArrayAttr mlir::linalg::DotOp::indexing_maps() { return getIndexingMaps(getOperation()); } -ArrayAttr mlir::linalg::FillOp::indexing_maps() { - return getIndexingMaps(getOperation()); -} ArrayAttr mlir::linalg::MatmulOp::indexing_maps() { return getIndexingMaps(getOperation()); } -ArrayAttr mlir::linalg::MatvecOp::indexing_maps() { - return getIndexingMaps(getOperation()); -} // TODO(ntv, rriddle): Consider making all this boilerplate easy to autogenerate // with Tablegen. This seems a desirable property in the context of OpInterfaces diff --git a/mlir/tools/mlir-tblgen/CMakeLists.txt b/mlir/tools/mlir-tblgen/CMakeLists.txt --- a/mlir/tools/mlir-tblgen/CMakeLists.txt +++ b/mlir/tools/mlir-tblgen/CMakeLists.txt @@ -5,6 +5,7 @@ add_tablegen(mlir-tblgen MLIR EnumsGen.cpp + LinalgNamedOpsGen.cpp LLVMIRConversionGen.cpp LLVMIRIntrinsicGen.cpp mlir-tblgen.cpp diff --git a/mlir/tools/mlir-tblgen/LinalgNamedOpsGen.cpp b/mlir/tools/mlir-tblgen/LinalgNamedOpsGen.cpp new file mode 100644 --- /dev/null +++ b/mlir/tools/mlir-tblgen/LinalgNamedOpsGen.cpp @@ -0,0 +1,122 @@ +//===- LinalgNamedOpsGen.cpp - MLIR Linalg op generator -------------------===// +// +// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// +// +// LinalgNamedOpsGen uses the description of structured operations in the Linalg +// dialect to generate op definitions, parsers, pretty-printers and matchers +// (e.g. linalg.matmul). +// +//===----------------------------------------------------------------------===// + +#include "mlir/Support/STLExtras.h" +#include "mlir/TableGen/GenInfo.h" +#include "mlir/TableGen/OpClass.h" +#include "mlir/TableGen/Operator.h" + +#include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/StringExtras.h" +#include "llvm/Support/FormatVariadic.h" +#include "llvm/Support/Signals.h" +#include "llvm/TableGen/Error.h" +#include "llvm/TableGen/Record.h" +#include "llvm/TableGen/TableGenBackend.h" + +using namespace llvm; +using namespace mlir; +using namespace mlir::tblgen; + +using mlir::tblgen::Operator; + +// Fills a function with signature "void iterator_types()". +void iteratorTypesBody(const Record &r, llvm::raw_ostream &os) { + auto iteratorTypes = r.getValueAsListOfStrings("iterators_types"); + os << " return SmallVector{"; + interleaveComma(iteratorTypes, os, + [&os](StringRef str) { os << "\"" << str << "\""; }); + os << "};"; +} + +// Fils a function with signature "ArrayAttr iterator_maps()". +void indexingMapsBody(const Record &r, llvm::raw_ostream &os) { + auto iterators = r.getValueAsListOfStrings("iterators"); + unsigned numDims = iterators.size(); + /* return AffineMap::get({...}, ctx); */ + auto ins = r.getValueAsListOfDefs("input_indexing_maps"); + auto outs = r.getValueAsListOfDefs("output_indexing_maps"); + SmallVector maps(ins.begin(), ins.end()); + maps.append(outs.begin(), outs.end()); + + /* AffineExpr ...; */ + os << " MLIRContext *ctx = getOperation()->getContext();\n"; + interleaveComma(iterators, os << " AffineExpr "); + os << ";\n"; + /* bindDims(ctx, ...); */ + interleaveComma(iterators, os << " bindDims(ctx, "); + os << ");\n"; + os << " return SmallVector{"; + interleaveComma(maps, os, [&os, numDims](const Record *def) { + os << "\n AffineMap::get(" << numDims << ", 0, {"; + interleaveComma(def->getValueAsListOfStrings("indexing_maps"), os, + [&os, numDims](const StringRef &str) { + os << "\n simplifyAffineExpr(" << str << ", " + << numDims << ", 0)"; + }); + os << "})"; + }); + os << "};"; +} + +template +void emitMethodWithBody(StringRef traitName, StringRef opName, + StringRef resultType, StringRef funName, + StringRef operands, const Record &r, + llvm::raw_ostream &os, Lambda fun) { + os << "template<>\n"; + os << resultType << " " << traitName << "<" << opName << ">::" << funName + << "(" << operands << ") {\n"; + fun(r, os); + os << "\n}\n\n"; +} + +bool emitLinalgNamedOpsInterfaceFunctions(const llvm::RecordKeeper &records, + llvm::raw_ostream &os) { + llvm::emitSourceFileHeader("Operations for Linalg Named Structured Ops", os); + os << "#include \"mlir/Dialect/Linalg/IR/LinalgTraits.h\"\n\n"; + + auto defs = records.getAllDerivedDefinitions("LinalgNamedStructured_Op"); + for (const llvm::Record *r : defs) { + const auto *traitName = "NamedStructuredOpTraits"; + auto opName = r->getName(); + const auto *resultType = "llvm::Optional>"; + const auto *funName = "referenceIterators"; + emitMethodWithBody(traitName, opName, resultType, funName, /*operands=*/"", + *r, os, iteratorTypesBody); + + resultType = "llvm::Optional>"; + funName = "referenceIndexingMaps"; + emitMethodWithBody(traitName, opName, resultType, funName, /*operands=*/"", + *r, os, indexingMapsBody); + + if (r->getValueAsBit("hasLibraryImpl")) { + traitName = "NamedStructuredOpLibraryCallTraits"; + resultType = "std::string"; + funName = "getLibraryCallName"; + emitMethodWithBody( + traitName, opName, resultType, funName, /*operands=*/"", *r, os, + [](const Record &r, llvm::raw_ostream &os) { + os << " return generateLibraryCallName(getOperation());"; + }); + } + } + + return false; +} + +static mlir::GenRegistration + genLinalgNamedOps("gen-linalg-named-structured-ops-defs", + "Generate Linalg named StructuredOps", + emitLinalgNamedOpsInterfaceFunctions);