From c818fa6729e02820c656431126d724d643001b70 Mon Sep 17 00:00:00 2001 From: Hanhan Wang Date: Thu, 28 Jan 2021 06:49:48 -0800 Subject: [PATCH] [mlir][Linalg] Replace SimplePad with PadTensor in tile-and-pad This revision creates a build method of PadTensorOp which can be mapped to SimplePad op. The verifier is updated to accept a static custom result type, which has the same semantic as SimplePadOp. Reviewed By: nicolasvasilache Differential Revision: https://reviews.llvm.org/D95555 --- mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td | 12 ++++ mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp | 78 +++++++++++++++++++++- mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp | 6 +- mlir/test/Dialect/Linalg/roundtrip.mlir | 19 ++++++ mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir | 9 ++- 5 files changed, 117 insertions(+), 7 deletions(-) diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td index 9ea1bc5..67c0615 100644 --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td @@ -199,6 +199,13 @@ def Linalg_PadTensorOp : Linalg_Op<"pad_tensor", static RankedTensorType inferResultType(RankedTensorType sourceType, ArrayRef staticLow, ArrayRef staticHigh); + + // Return a PadTensorOp that pads `source` to `type` size where the static + // sizes are assumed to be greater than the dynamic sizes. The op performs + // "high" padding (i.e. it adds trailing padding values until the desired + // size is met). + static linalg::PadTensorOp createPadHighOp( + Type type, Value source, Value pad, Location loc, OpBuilder & builder); }]; let builders = [ @@ -208,6 +215,11 @@ def Linalg_PadTensorOp : Linalg_Op<"pad_tensor", CArg<"ArrayRef", "{}">:$attrs)>, // Build a PadTensorOp with all dynamic entries. OpBuilderDAG<(ins "Value":$source, "ValueRange":$low, "ValueRange":$high, + CArg<"ArrayRef", "{}">:$attrs)>, + // Build a PadTensorOp with with mixed static and dynamic entries and custom + // result type. If the type passed is nullptr, it is inferred. + OpBuilderDAG<(ins "Type":$resultType, "Value":$source, + "ArrayRef":$low, "ArrayRef":$high, CArg<"ArrayRef", "{}">:$attrs)> ]; } diff --git a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp index 89a0b44..21c9108 100644 --- a/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp +++ b/mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp @@ -970,7 +970,11 @@ static LogicalResult verify(PadTensorOp op) { auto expectedType = PadTensorOp::inferResultType( sourceType, extractFromI64ArrayAttr(op.static_low()), extractFromI64ArrayAttr(op.static_high())); - if (resultType != expectedType) { + for (int i = 0, e = sourceType.getRank(); i < e; ++i) { + if (resultType.getDimSize(i) == expectedType.getDimSize(i)) + continue; + if (expectedType.isDynamicDim(i)) + continue; return op.emitError("specified type ") << resultType << " does not match the inferred type " << expectedType; @@ -1077,6 +1081,24 @@ static void print(OpAsmPrinter &p, PadTensorOp op) { p << " : " << op.source().getType() << " to " << op.getType(); } +/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if +/// it is a Value or into `staticVec` if it is an IntegerAttr. +/// In the case of a Value, a copy of the `sentinel` value is also pushed to +/// `staticVec`. This is useful to extract mixed static and dynamic entries that +/// come from an AttrSizedOperandSegments trait. +static void dispatchIndexOpFoldResult(OpFoldResult ofr, + SmallVectorImpl &dynamicVec, + SmallVectorImpl &staticVec, + int64_t sentinel) { + if (auto v = ofr.dyn_cast()) { + dynamicVec.push_back(v); + staticVec.push_back(sentinel); + return; + } + APInt apInt = ofr.dyn_cast().cast().getValue(); + staticVec.push_back(apInt.getSExtValue()); +} + void PadTensorOp::build(OpBuilder &b, OperationState &result, Value source, ArrayRef staticLow, ArrayRef staticHigh, ValueRange low, @@ -1097,6 +1119,60 @@ void PadTensorOp::build(OpBuilder &b, OperationState &result, Value source, build(b, result, source, staticVector, staticVector, low, high, attrs); } +void PadTensorOp::build(OpBuilder &b, OperationState &result, Type resultType, + Value source, ArrayRef low, + ArrayRef high, + ArrayRef attrs) { + assert(resultType.isa()); + auto sourceType = source.getType().cast(); + unsigned rank = sourceType.getRank(); + SmallVector dynamicLow, dynamicHigh; + SmallVector staticLow, staticHigh; + for (unsigned i = 0; i < rank; ++i) { + // staticLow and staticHigh have full information of the padding config. + // This will grow staticLow and staticHigh with 1 value. If the config is + // dynamic (ie not a constant), dynamicLow and dynamicHigh will grow with 1 + // value as well. + dispatchIndexOpFoldResult(low[i], dynamicLow, staticLow, + ShapedType::kDynamicSize); + dispatchIndexOpFoldResult(high[i], dynamicHigh, staticHigh, + ShapedType::kDynamicSize); + } + if (!resultType) { + resultType = + PadTensorOp::inferResultType(sourceType, staticLow, staticHigh); + } + build(b, result, resultType, source, dynamicLow, dynamicHigh, + b.getI64ArrayAttr(staticLow), b.getI64ArrayAttr(staticHigh)); +} + +PadTensorOp PadTensorOp::createPadHighOp(Type type, Value source, Value pad, + Location loc, OpBuilder &builder) { + SmallVector low, high; + auto rankedTensorType = type.cast(); + assert(rankedTensorType.hasStaticShape()); + int rank = rankedTensorType.getRank(); + for (int i = 0; i < rank; ++i) { + auto dimOp = builder.createOrFold(loc, source, i); + auto resultDimSize = builder.createOrFold( + loc, rankedTensorType.getDimSize(i)); + auto highValue = builder.createOrFold(loc, resultDimSize, dimOp); + high.push_back(highValue); + low.push_back(builder.createOrFold(loc, 0)); + } + auto padTensorOp = + builder.create(loc, type, source, low, high); + SmallVector blockArgTypes; + blockArgTypes.assign(rank, builder.getIndexType()); + auto ®ion = padTensorOp.region(); + // `builder.createBlock` changes the insertion point within the block. Create + // a guard to reset the insertion point of the builder after it is destroyed. + OpBuilder::InsertionGuard guard(builder); + builder.createBlock(®ion, region.end(), blockArgTypes); + builder.create(loc, pad); + return padTensorOp; +} + //===----------------------------------------------------------------------===// // ReshapeOp //===----------------------------------------------------------------------===// diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp index b4c94ae..7260bb4 100644 --- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp +++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp @@ -120,7 +120,7 @@ mlir::linalg::LinalgTilingOptions::setTileSizes(ArrayRef ts) { /// Return success if either: /// 1. The operand is already statically shaped, `result` is left unchanged. /// 2. The operand is (partially) dynamic, `result` is the result of a freshly -/// created SimplePadOp. +/// created PadTensorOp. /// Return failure if the operand cannot be padded to a static shape. static LogicalResult padOperandToSmallestStaticBoundingBox( PatternRewriter &rewriter, linalg::LinalgOp opToPad, Value operand, @@ -151,8 +151,8 @@ static LogicalResult padOperandToSmallestStaticBoundingBox( Value pad = options.paddingValueComputationFunction(rewriter, opToPad); auto staticTensorType = RankedTensorType::get(staticSizes, tensorType.getElementType()); - result = rewriter.create(opToPad->getLoc(), - staticTensorType, operand, pad); + result = linalg::PadTensorOp::createPadHighOp(staticTensorType, operand, pad, + opToPad->getLoc(), rewriter); return success(); } diff --git a/mlir/test/Dialect/Linalg/roundtrip.mlir b/mlir/test/Dialect/Linalg/roundtrip.mlir index 6dc0768..879dfa3 100644 --- a/mlir/test/Dialect/Linalg/roundtrip.mlir +++ b/mlir/test/Dialect/Linalg/roundtrip.mlir @@ -57,6 +57,25 @@ func @pad_asymmetrical(%arg0: tensor<2x3xf32>, %ub0: index, %ub1: index, // ----- +func @pad_to_static_size(%arg0: tensor, %ub0: index, %ub1: index, + %pad_value: f32) -> tensor<2x3xf32> { + %0 = linalg.pad_tensor %arg0 low[0, 0] high[%ub0, %ub1] { + ^bb0(%arg1: index, %arg2: index): + linalg.yield %pad_value : f32 + } : tensor to tensor<2x3xf32> + return %0 : tensor<2x3xf32> +} +// CHECK-LABEL: func @pad_to_static_size +// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]] +// CHECK-SAME: %[[UB0:[a-zA-Z0-9_]*]] +// CHECK-SAME: %[[UB1:[a-zA-Z0-9_]*]] +// CHECK: linalg.pad_tensor %[[ARG0]] +// CHECK-SAME: low[0, 0] +// CHECK-SAME: high[%[[UB0]], %[[UB1]]] +// CHECK: : tensor to tensor<2x3xf32> + +// ----- + func @range(%arg0: index, %arg1: index, %arg2: index) { %0 = linalg.range %arg0:%arg1:%arg2 : !linalg.range return diff --git a/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir b/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir index e412108..4eaf830 100644 --- a/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir +++ b/mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir @@ -18,9 +18,12 @@ func @matmul_tensors( // CHECK-NOT: linalg.matmul {{.*}} tensor // Padding injects static information. -// CHECK: %[[pA:.*]] = linalg.simple_pad %[[sTA]] pad %{{.*}} : tensor to tensor<2x4xf32> pad f32 -// CHECK: %[[pB:.*]] = linalg.simple_pad %[[sTB]] pad %{{.*}} : tensor to tensor<4x3xf32> pad f32 -// CHECK: %[[pC:.*]] = linalg.simple_pad %[[sTC]] pad %{{.*}} : tensor to tensor<2x3xf32> pad f32 +// CHECK: %[[pA:.*]] = linalg.pad_tensor %[[sTA]] low[0, 0] high[%{{.*}}, %{{.*}}] +// CHECK: : tensor to tensor<2x4xf32> +// CHECK: %[[pB:.*]] = linalg.pad_tensor %[[sTB]] low[0, 0] high[%{{.*}}, %{{.*}}] +// CHECK: : tensor to tensor<2x4xf32> +// CHECK: %[[pC:.*]] = linalg.pad_tensor %[[sTC]] low[0, 0] high[%{{.*}}, %{{.*}}] +// CHECK: : tensor to tensor<2x4xf32> // CHECK: %[[pD:.*]] = linalg.matmul ins(%[[pA]], %[[pB]] : tensor<2x4xf32>, tensor<4x3xf32>) // CHECK-SAME: outs(%[[pC]] : tensor<2x3xf32>) -> tensor<2x3xf32> // CHECK: %[[sTD:.*]] = subtensor %[[pD]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : tensor<2x3xf32> to tensor -- 2.7.4