From 3747eb9c85b3393aa00ad12e9e7ef31ffec8bd4c Mon Sep 17 00:00:00 2001 From: Nicolas Vasilache Date: Mon, 25 Jan 2021 09:11:31 +0000 Subject: [mlir][Linalg] Add a padding option to Linalg tiling This revision allows the base Linalg tiling pattern to optionally require padding to a constant bounding shape. When requested, a simple analysis is performed, similar to buffer promotion. A temporary `linalg.simple_pad` op is added to model padding for the purpose of connecting the dots. This will be replaced by a more fleshed out `linalg.pad_tensor` op when it is available. In the meantime, this temporary op serves the purpose of exhibiting the necessary properties required from a more fleshed out pad op, to compose with transformations properly. Reviewed By: ftynse Differential Revision: https://reviews.llvm.org/D95149 --- mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td | 32 ++++++++++++++++++++++ .../mlir/Dialect/Linalg/Transforms/Transforms.h | 20 ++++++++++++++ mlir/include/mlir/Interfaces/ViewLikeInterface.h | 1 + mlir/include/mlir/Interfaces/ViewLikeInterface.td | 25 +++++++++++++++++ 4 files changed, 78 insertions(+) (limited to 'mlir/include') diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td index ae9f81d043f5..9ea1bc5a3587 100644 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td @@ -475,6 +475,38 @@ def Linalg_SliceOp : Linalg_Op<"slice", [ let hasFolder = 1; } +def Linalg_SimplePadOp : Linalg_Op<"simple_pad", [NoSideEffect]> { + let summary = "TODO: replace with pad_tensors when ready."; + + let description = [{ + `linalg.simple_pad` is a tmp placeholder for padding and packing on tensors. + Its semantics are to pad a partially dynamic tensor to a fully static tensor + where the static sizes are assumed to be greater than the dynamic sizes. The + op perforrms "high" padding (i.e. it adds trailing padding values until the + desired size is met). + }]; + + let arguments = (ins AnyRankedTensor:$tensor, AnyType:$padding); + let results = (outs AnyRankedTensor:$result); + + // TODO: verify all static result, some dynamic input, static shapes match, + // element types match, ranks match etc. Use pad_tensors when ready but for + // now just let it ne fully specified by traits. + let verifier = ?; + + let extraClassDeclaration = [{ + RankedTensorType getSourceType() { + return tensor().getType().cast(); } + RankedTensorType getResultType() { + return getResult().getType().cast(); } + }]; + + let assemblyFormat = [{ + $tensor `pad` $padding attr-dict `:` + type($tensor) `to` type($result) `pad` type($padding) + }]; +} + def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>, Arguments<(ins Variadic:$values)> { let summary = "Linalg yield operation"; diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h index 611ab6867372..f359992e5ff1 100644 --- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h +++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h @@ -345,6 +345,9 @@ enum class LinalgTilingLoopType { using TileSizeComputationFunction = std::function(OpBuilder &, Operation *)>; +using PaddingValueComputationFunction = + std::function; + struct LinalgTilingOptions { /// Computation function that returns the tile sizes for each operation. /// Delayed construction of constant tile sizes should occur to interoperate @@ -393,6 +396,18 @@ struct LinalgTilingOptions { distribution = std::move(distributionOptions); return *this; } + + /// Computation function that returns a padding value to use when padding to + /// force static sizes. When `paddingValueComputationFunction` is set, padding + /// operations are introduced, that guarantee the underlying op is statically + /// shaped and can thus be vectorized. + PaddingValueComputationFunction paddingValueComputationFunction = nullptr; + + LinalgTilingOptions & + setPaddingValueComputationFunction(PaddingValueComputationFunction fun) { + paddingValueComputationFunction = std::move(fun); + return *this; + } }; /// Canonicalization patterns relevant to apply after tiling patterns. These are @@ -403,6 +418,11 @@ getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx); void populateLinalgTilingCanonicalizationPatterns( OwningRewritePatternList &patterns, MLIRContext *ctx); +/// Base pattern that applied the tiling transformation specified by `options`. +/// Abort and return failure in 2 cases: +/// 1. if the tiling specification is invalid and tiling fails to occur. +/// 2. if tiling occurs but `options.paddingValueComputationFunction` is set +/// and some operand shape cannot be bounded statically. struct LinalgBaseTilingPattern : public RewritePattern { // Entry point to match any LinalgOp OpInterface. LinalgBaseTilingPattern(LinalgTilingOptions options, diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.h b/mlir/include/mlir/Interfaces/ViewLikeInterface.h index d5f44c3e63da..2b3a054338ab 100644 --- a/mlir/include/mlir/Interfaces/ViewLikeInterface.h +++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.h @@ -14,6 +14,7 @@ #define MLIR_INTERFACES_VIEWLIKEINTERFACE_H_ #include "mlir/IR/Builders.h" +#include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypes.h" #include "mlir/IR/OpImplementation.h" diff --git a/mlir/include/mlir/Interfaces/ViewLikeInterface.td b/mlir/include/mlir/Interfaces/ViewLikeInterface.td index 62c371b2f97d..6c72b47f2ac3 100644 --- a/mlir/include/mlir/Interfaces/ViewLikeInterface.td +++ b/mlir/include/mlir/Interfaces/ViewLikeInterface.td @@ -108,6 +108,28 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface return $_op.sizes(); }] >, + InterfaceMethod< + /*desc=*/[{ + Return a vector of all the static or dynamic sizes of the op. + }], + /*retTy=*/"SmallVector", + /*methodName=*/"getMixedSizes", + /*args=*/(ins), + /*methodBody=*/"", + /*defaultImplementation=*/[{ + SmallVector res; + std::array ranks = $_op.getArrayAttrRanks(); + unsigned numDynamic = 0; + unsigned count = ranks[getOffsetOperandGroupPosition()]; + for (unsigned idx = 0; idx < count; ++idx) { + if (isDynamicSize(idx)) + res.push_back($_op.sizes()[numDynamic++]); + else + res.push_back($_op.static_sizes()[idx]); + } + return res; + }] + >, InterfaceMethod< /*desc=*/[{ Return the dynamic stride operands. @@ -359,6 +381,9 @@ def OffsetSizeAndStrideOpInterface : OpInterface<"OffsetSizeAndStrideOpInterface ]; let extraClassDeclaration = [{ + static unsigned getOffsetOperandGroupPosition() { return 0; } + static unsigned getSizeOperandGroupPosition() { return 1; } + static unsigned getStrideOperandGroupPosition() { return 2; } static StringRef getStaticOffsetsAttrName() { return "static_offsets"; } -- cgit v1.2.3