[Mlir-commits] [mlir] 98fff51 - [mlir][linalg] Lower PadTensorOp to InitTensorOp + FillOp + SubTensorInitOp
Matthias Springer
llvmlistbot at llvm.org
Sun Jun 13 22:26:31 PDT 2021
Author: Matthias Springer
Date: 2021-06-14T14:21:08+09:00
New Revision: 98fff5153a81fc501a69a4c184de5236fc806b9f
URL: https://github.com/llvm/llvm-project/commit/98fff5153a81fc501a69a4c184de5236fc806b9f
DIFF: https://github.com/llvm/llvm-project/commit/98fff5153a81fc501a69a4c184de5236fc806b9f.diff
LOG: [mlir][linalg] Lower PadTensorOp to InitTensorOp + FillOp + SubTensorInitOp
Currently limited to constant pad values. Any combination of dynamic/static tensor sizes and padding sizes is supported.
Differential Revision: https://reviews.llvm.org/D103679
Added:
Modified:
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
mlir/test/Dialect/Linalg/vectorization.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 3e8fb3813ff2..f52bd10da122 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -671,6 +671,75 @@ static SmallVector<Value> ofrToIndexValues(OpBuilder &builder, Location loc,
return result;
}
+/// Rewrite a PadTensorOp into a sequence of InitTensorOp, FillOp and
+/// SubTensorInsertOp. For now, only constant padding values are supported.
+/// Note: This rewrite is not yet a vectorization, but some of the generated ops
+/// may be vectorized down the line (e.g., FillOp).
+/// TODO: If there is enough static shape information, generate TransferReadOps
+/// and TransferWriteOps instead of SubTensorInsertOp.
+struct GenericPadTensorOpVectorizationPattern
+ : public OpRewritePattern<PadTensorOp> {
+ using OpRewritePattern<PadTensorOp>::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(PadTensorOp padOp,
+ PatternRewriter &rewriter) const final {
+ // Given an OpFoldResult, return an index-typed value.
+ auto getIdxValue = [&](OpFoldResult ofr) {
+ if (auto val = ofr.dyn_cast<Value>())
+ return val;
+ return rewriter.create<ConstantIndexOp>(
+ padOp.getLoc(), getIntFromAttr(ofr.get<Attribute>())).getResult();
+ };
+
+ // Pad value must be a constant.
+ auto padValue = padOp.getConstantPaddingValue();
+ if (!padValue) return failure();
+
+ auto resultType = padOp.getResultType();
+ // Compute size of InitTensorOp. Any combination of static/dynamic is
+ // supported.
+ SmallVector<Value> dynSizes;
+ SmallVector<int64_t> staticSizes;
+ for (unsigned dim = 0; dim < resultType.getRank(); ++dim) {
+ if (resultType.isDynamicDim(dim)) {
+ auto srcSize = rewriter.createOrFold<memref::DimOp>(
+ padOp.getLoc(), padOp.source(), dim);
+ // Add low and high padding value.
+ auto plusLow = rewriter.createOrFold<AddIOp>(
+ padOp.getLoc(), srcSize, getIdxValue(padOp.getMixedLowPad()[dim]));
+ auto plusHigh = rewriter.createOrFold<AddIOp>(
+ padOp.getLoc(), plusLow, getIdxValue(padOp.getMixedHighPad()[dim]));
+ dynSizes.push_back(plusHigh);
+ }
+ staticSizes.push_back(resultType.getDimSize(dim));
+ }
+
+ Value init = rewriter.create<InitTensorOp>(
+ padOp.getLoc(), dynSizes, staticSizes, resultType.getElementType());
+ Value fill =
+ rewriter.create<FillOp>(padOp.getLoc(), init, padValue).result();
+
+ auto sourceType = padOp.getSourceType();
+ // Compute size of source of PadTensorOp.
+ SmallVector<OpFoldResult> srcSizes;
+ for (unsigned dim = 0; dim < sourceType.getRank(); ++dim) {
+ if (sourceType.isDynamicDim(dim)) {
+ srcSizes.push_back(rewriter.createOrFold<memref::DimOp>(
+ padOp.getLoc(), padOp.source(), dim));
+ } else {
+ srcSizes.push_back(rewriter.getIndexAttr(sourceType.getDimSize(dim)));
+ }
+ }
+ // Strides of SubTensorInsertOp are all 1.
+ SmallVector<OpFoldResult> strides(sourceType.getRank(),
+ rewriter.getIndexAttr(1));
+ rewriter.replaceOpWithNewOp<SubTensorInsertOp>(
+ padOp, padOp.source(), fill, padOp.getMixedLowPad(), srcSizes, strides);
+
+ return success();
+ }
+};
+
/// Base pattern for rewriting PadTensorOps whose result is consumed by a given
/// operation type OpTy.
template <typename OpTy>
@@ -949,14 +1018,13 @@ struct PadTensorOpVectorizationWithSubTensorInsertPattern
void mlir::linalg::populatePadTensorOpVectorizationPatterns(
RewritePatternSet &patterns, PatternBenefit baseBenefit) {
- // TODO: Canonicalizer handles simple cases where low = 0 and high = 0, but a
- // generic vectorization pattern is still missing.
-
+ patterns.add<GenericPadTensorOpVectorizationPattern>(
+ patterns.getContext(), baseBenefit);
// Try these specialized patterns first before resorting to the generic one.
patterns.add<PadTensorOpVectorizationWithTransferReadPattern,
PadTensorOpVectorizationWithTransferWritePattern,
PadTensorOpVectorizationWithSubTensorInsertPattern>(
- patterns.getContext(), baseBenefit);
+ patterns.getContext(), baseBenefit.getBenefit() + 1);
}
// TODO: cleanup all the convolution vectorization patterns.
diff --git a/mlir/test/Dialect/Linalg/vectorization.mlir b/mlir/test/Dialect/Linalg/vectorization.mlir
index 43dd39602b35..a3223309cc01 100644
--- a/mlir/test/Dialect/Linalg/vectorization.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization.mlir
@@ -512,21 +512,44 @@ func @matmul_i8_i8_i32(%a: memref<4x6xi8>, %b: memref<6x12xi8>, %c: memref<4x12x
// -----
-// CHECK-LABEL: func @pad_static_high_padding
-// CHECK: linalg.pad_tensor
-func @pad_static_high_padding(%arg0: tensor<?x?x?xf32>, %pad_value: f32) -> tensor<2x3x4xf32> {
- %0 = linalg.pad_tensor %arg0 low[0, 0, 0] high[0, 1, 0] {
+// CHECK-LABEL: func @pad_static(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<2x?x2xf32>, %[[PAD:.*]]: f32
+// CHECK-NOT: linalg.pad_tensor
+// CHECK-DAG: %[[C1:.*]] = constant 1 : index
+// CHECK-DAG: %[[INIT:.*]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32>
+// CHECK-DAG: %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x3x4xf32>
+// CHECK: %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]]{{.*}} : vector<2x3x4xf32>, tensor<2x3x4xf32>
+// CHECK-DAG: %[[DIM1:.*]] = memref.dim %[[ARG0]], %[[C1]]
+// CHECK: %[[RESULT:.*]] = subtensor_insert %[[ARG0]] into %2[0, 0, 2] [2, %[[DIM1]], 2] [1, 1, 1] : tensor<2x?x2xf32> into tensor<2x3x4xf32>
+// CHECK: return %[[RESULT]]
+func @pad_static(%arg0: tensor<2x?x2xf32>, %pad_value: f32) -> tensor<2x3x4xf32> {
+ %0 = linalg.pad_tensor %arg0 low[0, 0, 2] high[0, 1, 0] {
^bb0(%arg1: index, %arg2: index, %arg3: index):
linalg.yield %pad_value : f32
- } : tensor<?x?x?xf32> to tensor<2x3x4xf32>
+ } : tensor<2x?x2xf32> to tensor<2x3x4xf32>
return %0 : tensor<2x3x4xf32>
}
// -----
-// CHECK-LABEL: func @pad_dynamic
-// CHECK: linalg.pad_tensor
-func @pad_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
+// CHECK-LABEL: func @pad_static_dynamic(
+// CHECK-SAME: %[[SRC:.*]]: tensor<1x2x2x?xf32>, %[[LOW:.*]]: index, %[[HIGH:.*]]: index
+// CHECK-NOT: linalg.pad_tensor
+// CHECK-DAG: %[[C2:.*]] = constant 2 : index
+// CHECK-DAG: %[[C3:.*]] = constant 3 : index
+// CHECK-DAG: %[[C5:.*]] = constant 5 : index
+// CHECK: %[[V0:.*]] = addi %[[LOW]], %[[C2]] : index
+// CHECK: %[[V1:.*]] = addi %[[V0]], %[[C3]] : index
+// CHECK: %[[V2:.*]] = addi %[[HIGH]], %[[C5]] : index
+// CHECK: %[[DIM3:.*]] = memref.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32>
+// CHECK: %[[V4:.*]] = addi %[[DIM3]], %[[C3]] : index
+// CHECK: %[[V5:.*]] = addi %[[V4]], %[[C2]] : index
+// CHECK: %[[INIT:.*]] = linalg.init_tensor [6, %[[V1]], %[[V2]], %[[V5]]] : tensor<6x?x?x?xf32>
+// CHECK: %[[FILL:.*]] = linalg.fill(%[[INIT]], %{{.*}}) : tensor<6x?x?x?xf32>, f32 -> tensor<6x?x?x?xf32>
+// CHECK: %[[SRCDIM:.*]] = memref.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32>
+// CHECK: %[[RESULT:.*]] = subtensor_insert %[[SRC]] into %[[FILL]][2, %[[LOW]], 3, 3] [1, 2, 2, %[[SRCDIM]]] [1, 1, 1, 1] : tensor<1x2x2x?xf32> into tensor<6x?x?x?xf32>
+// CHECK: return %[[RESULT]]
+func @pad_static_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index,
%pad_value: f32) -> tensor<6x?x?x?xf32> {
%0 = linalg.pad_tensor %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] {
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
More information about the Mlir-commits
mailing list