[Mlir-commits] [mlir] [MLIR][Linalg] Add boundary fill option to pad fusion pass (PR #223792)
Chibuoyim Ogbonna
llvmlistbot at llvm.org
Wed Sep 16 01:33:19 PDT 2026
https://github.com/bruteforceboy updated https://github.com/llvm/llvm-project/pull/223792
>From 62cb4878b8308fafbc458f9071237976857c2557 Mon Sep 17 00:00:00 2001
From: bruteforceboy <chibuoyim.faith.ogbonna at huawei.com>
Date: Tue, 15 Sep 2026 16:12:00 +0100
Subject: [PATCH] [MLIR][Linalg] Add boundary fill option to pad fusion pass
---
.../Dialect/Linalg/Transforms/Transforms.h | 7 +-
.../FusePadOpWithLinalgProducer.cpp | 101 ++++-
.../Linalg/pad_fusion_fill_boundary.mlir | 377 ++++++++++++++++++
.../pad_fusion_fill_boundary_bufferize.mlir | 46 +++
.../test/lib/Dialect/Linalg/TestPadFusion.cpp | 13 +-
5 files changed, 527 insertions(+), 17 deletions(-)
create mode 100644 mlir/test/Dialect/Linalg/pad_fusion_fill_boundary.mlir
create mode 100644 mlir/test/Dialect/Linalg/pad_fusion_fill_boundary_bufferize.mlir
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index cbc338c9f1b321..62401e25c2c9b0 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -2062,8 +2062,13 @@ void populateFoldAddIntoDestPatterns(RewritePatternSet &patterns);
/// Pattern to fuse a `tensor.pad` operation with the producer of its source,
/// if the producer is a `linalg` operation with all parallel iterator types.
+///
+/// If `fillBoundaryOnly` is set, static pads whose producer overwrites the
+/// interior fill only the padding instead of the whole destination.
+/// Otherwise, including when boundary-only filling is inapplicable, fill the
+/// whole destination.
void populateFuseTensorPadWithProducerLinalgOpPatterns(
- RewritePatternSet &patterns);
+ RewritePatternSet &patterns, bool fillBoundaryOnly = false);
/// Patterns to simplify depthwise convolutions.
void populateSimplifyDepthwiseConvPatterns(RewritePatternSet &patterns);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp b/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp
index 0ca89040f34538..154964c0e782a9 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/FusePadOpWithLinalgProducer.cpp
@@ -15,9 +15,66 @@
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
using namespace mlir;
+/// Fill the padding in disjoint slabs. After filling both sides of a dimension,
+/// restrict it to the interior so subsequent slabs do not overlap.
+static Value fillPaddedBoundary(RewriterBase &rewriter, Location loc,
+ Value dest, Value padValue,
+ ArrayRef<int64_t> lowPad,
+ ArrayRef<int64_t> sourceShape) {
+ ArrayRef<int64_t> paddedShape =
+ cast<RankedTensorType>(dest.getType()).getShape();
+ SmallVector<int64_t> offsets(paddedShape.size(), 0);
+ SmallVector<int64_t> sizes(paddedShape);
+ SmallVector<OpFoldResult> strides(paddedShape.size(),
+ rewriter.getIndexAttr(1));
+
+ auto fillSlab = [&](unsigned dim, int64_t offset, int64_t size) {
+ offsets[dim] = offset;
+ sizes[dim] = size;
+ if (llvm::is_contained(sizes, 0))
+ return;
+ auto sliceOffsets = getAsIndexOpFoldResult(rewriter.getContext(), offsets);
+ auto sliceSizes = getAsIndexOpFoldResult(rewriter.getContext(), sizes);
+ auto slice = tensor::ExtractSliceOp::create(
+ rewriter, loc, dest, sliceOffsets, sliceSizes, strides);
+ auto filled =
+ linalg::FillOp::create(rewriter, loc, padValue, slice.getResult());
+ dest =
+ tensor::InsertSliceOp::create(rewriter, loc, filled.getResult(0), dest,
+ sliceOffsets, sliceSizes, strides);
+ };
+
+ for (auto [dim, sourceSize] : llvm::enumerate(sourceShape)) {
+ int64_t highOffset = lowPad[dim] + sourceSize;
+ fillSlab(dim, 0, lowPad[dim]);
+ fillSlab(dim, highOffset, paddedShape[dim] - highOffset);
+ offsets[dim] = lowPad[dim];
+ sizes[dim] = sourceSize;
+ }
+ return dest;
+}
+
+static bool canLeaveInteriorUninitialized(linalg::GenericOp linalgOp,
+ OpOperand *initOperand,
+ tensor::PadOp padOp) {
+ if (!cast<RankedTensorType>(padOp.getSource().getType()).hasStaticShape() ||
+ !padOp.getResultType().hasStaticShape() ||
+ llvm::any_of(padOp.getStaticLow(), ShapedType::isDynamic))
+ return false;
+
+ if (linalgOp.payloadUsesValueFromOperand(initOperand))
+ return false;
+
+ // A projected permutation is not enough: it may drop a loop dimension, e.g.
+ // `(d0, d1) -> (d0)`, and if that dimension has zero extent the producer runs
+ // no iterations at all and leaves the interior uninitialized.
+ return linalgOp.getMatchingIndexingMap(initOperand).isPermutation();
+}
+
namespace {
/// A sequence of operations
@@ -38,7 +95,10 @@ namespace {
///
/// if the `linalg.generic` has all parallel iterator types.
struct FusePadOp : OpRewritePattern<tensor::PadOp> {
- using OpRewritePattern<tensor::PadOp>::OpRewritePattern;
+ FusePadOp(MLIRContext *context, bool fillBoundaryOnly,
+ PatternBenefit benefit = 1)
+ : OpRewritePattern<tensor::PadOp>(context, benefit),
+ fillBoundaryOnly(fillBoundaryOnly) {}
LogicalResult matchAndRewrite(tensor::PadOp padOp,
PatternRewriter &rewriter) const override {
@@ -75,22 +135,30 @@ struct FusePadOp : OpRewritePattern<tensor::PadOp> {
auto emptyTensor = tensor::EmptyOp::create(rewriter, loc, resultSizes,
padResultType.getElementType());
- // Fill the tensor with the pad value.
- // TODO: There is an option to fill only the boundaries. For now just
- // filling the whole tensor.
- auto fillTensor = linalg::FillOp::create(rewriter, loc, padValue,
- emptyTensor.getResult());
+ unsigned resultNumber = cast<OpResult>(source).getResultNumber();
+ auto sourceType = cast<RankedTensorType>(source.getType());
+
+ Value fillTensor;
+ if (fillBoundaryOnly &&
+ canLeaveInteriorUninitialized(
+ linalgOp, linalgOp.getDpsInitOperand(resultNumber), padOp)) {
+ fillTensor =
+ fillPaddedBoundary(rewriter, loc, emptyTensor.getResult(), padValue,
+ padOp.getStaticLow(), sourceType.getShape());
+ } else {
+ fillTensor = linalg::FillOp::create(rewriter, loc, padValue,
+ emptyTensor.getResult())
+ .getResult(0);
+ }
// Construct a slice of the fill result that is to be replaced with the
// result of the generic op. The low pad values are the offsets, the size of
// the source is the size of the slice.
// TODO: This insert/extract could be potentially made a utility method.
- unsigned resultNumber = cast<OpResult>(source).getResultNumber();
SmallVector<OpFoldResult> offsets = padOp.getMixedLowPad();
SmallVector<OpFoldResult> sizes;
sizes.reserve(offsets.size());
- for (const auto &shape :
- llvm::enumerate(cast<RankedTensorType>(source.getType()).getShape())) {
+ for (const auto &shape : llvm::enumerate(sourceType.getShape())) {
if (ShapedType::isDynamic(shape.value())) {
sizes.push_back(
tensor::DimOp::create(rewriter, loc, source, shape.index())
@@ -100,8 +168,8 @@ struct FusePadOp : OpRewritePattern<tensor::PadOp> {
}
}
SmallVector<OpFoldResult> strides(offsets.size(), rewriter.getIndexAttr(1));
- auto slice = tensor::ExtractSliceOp::create(
- rewriter, loc, fillTensor.getResult(0), offsets, sizes, strides);
+ auto slice = tensor::ExtractSliceOp::create(rewriter, loc, fillTensor,
+ offsets, sizes, strides);
// Clone the generic op.
auto clonedOp =
@@ -110,14 +178,17 @@ struct FusePadOp : OpRewritePattern<tensor::PadOp> {
// Insert it back into the result of the fill.
rewriter.replaceOpWithNewOp<tensor::InsertSliceOp>(
- padOp, clonedOp.getResult(resultNumber), fillTensor.getResult(0),
- offsets, sizes, strides);
+ padOp, clonedOp.getResult(resultNumber), fillTensor, offsets, sizes,
+ strides);
return success();
}
+
+private:
+ bool fillBoundaryOnly;
};
} // namespace
void mlir::linalg::populateFuseTensorPadWithProducerLinalgOpPatterns(
- RewritePatternSet &patterns) {
- patterns.add<FusePadOp>(patterns.getContext());
+ RewritePatternSet &patterns, bool fillBoundaryOnly) {
+ patterns.add<FusePadOp>(patterns.getContext(), fillBoundaryOnly);
}
diff --git a/mlir/test/Dialect/Linalg/pad_fusion_fill_boundary.mlir b/mlir/test/Dialect/Linalg/pad_fusion_fill_boundary.mlir
new file mode 100644
index 00000000000000..3a829d6654f8e9
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/pad_fusion_fill_boundary.mlir
@@ -0,0 +1,377 @@
+// RUN: mlir-opt -test-linalg-pad-fusion=fill-boundary-only -split-input-file %s | FileCheck %s
+// RUN: mlir-opt -test-linalg-pad-fusion -split-input-file %s | FileCheck %s --check-prefix=DEFAULT
+
+func.func @fill_boundary_2d(
+ %arg0 : tensor<4x3xf32>, %arg1 : f32) -> tensor<7x6xf32> {
+ %init = tensor.empty() : tensor<4x3xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x3xf32>) outs(%init : tensor<4x3xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.mulf %arg2, %arg2 : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x3xf32>
+ %1 = tensor.pad %0 low [1, 2] high [2, 1] {
+ ^bb0(%arg2: index, %arg3 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<4x3xf32> to tensor<7x6xf32>
+ return %1 : tensor<7x6xf32>
+}
+
+// CHECK: func @fill_boundary_2d
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<4x3xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<7x6xf32>
+// CHECK: %[[TOP_S:.+]] = tensor.extract_slice %[[EMPTY]][0, 0] [1, 6] [1, 1]
+// CHECK: %[[TOP_F:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[TOP_S]]
+// CHECK: %[[TOP:.+]] = tensor.insert_slice %[[TOP_F]] into %[[EMPTY]][0, 0] [1, 6] [1, 1]
+// CHECK: %[[BOT_S:.+]] = tensor.extract_slice %[[TOP]][5, 0] [2, 6] [1, 1]
+// CHECK: %[[BOT_F:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[BOT_S]]
+// CHECK: %[[BOT:.+]] = tensor.insert_slice %[[BOT_F]] into %[[TOP]][5, 0] [2, 6] [1, 1]
+// CHECK: %[[LEFT_S:.+]] = tensor.extract_slice %[[BOT]][1, 0] [4, 2] [1, 1]
+// CHECK: %[[LEFT_F:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[LEFT_S]]
+// CHECK: %[[LEFT:.+]] = tensor.insert_slice %[[LEFT_F]] into %[[BOT]][1, 0] [4, 2] [1, 1]
+// CHECK: %[[RIGHT_S:.+]] = tensor.extract_slice %[[LEFT]][1, 5] [4, 1] [1, 1]
+// CHECK: %[[RIGHT_F:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[RIGHT_S]]
+// CHECK: %[[BOUNDARY:.+]] = tensor.insert_slice %[[RIGHT_F]] into %[[LEFT]][1, 5] [4, 1] [1, 1]
+// CHECK-NOT: linalg.fill
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[BOUNDARY]][1, 2] [4, 3] [1, 1]
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: outs(%[[SLICE]] :
+// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[GENERIC]] into %[[BOUNDARY]][1, 2] [4, 3] [1, 1]
+// CHECK: return %[[RESULT]]
+
+// DEFAULT: func @fill_boundary_2d
+// DEFAULT: %[[EMPTY:.+]] = tensor.empty() : tensor<7x6xf32>
+// DEFAULT: %[[FILL:.+]] = linalg.fill ins(%{{.+}} : f32) outs(%[[EMPTY]] : tensor<7x6xf32>)
+// DEFAULT: tensor.extract_slice %[[FILL]][1, 2] [4, 3] [1, 1]
+
+// -----
+
+func.func @fill_boundary_3d(%arg0 : tensor<2x2x2xf32>, %arg1 : f32) -> tensor<3x3x5xf32> {
+ %init = tensor.empty() : tensor<2x2x2xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>,
+ affine_map<(d0, d1, d2) -> (d0, d1, d2)>],
+ iterator_types = ["parallel", "parallel", "parallel"]}
+ ins(%arg0 : tensor<2x2x2xf32>) outs(%init : tensor<2x2x2xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.mulf %arg2, %arg2 : f32
+ linalg.yield %1 : f32
+ } -> tensor<2x2x2xf32>
+ %1 = tensor.pad %0 low [1, 0, 2] high [0, 1, 1] {
+ ^bb0(%arg2: index, %arg3 : index, %arg4 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<2x2x2xf32> to tensor<3x3x5xf32>
+ return %1 : tensor<3x3x5xf32>
+}
+
+// CHECK: func @fill_boundary_3d
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<2x2x2xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<3x3x5xf32>
+// CHECK: %[[S0:.+]] = tensor.extract_slice %[[EMPTY]][0, 0, 0] [1, 3, 5] [1, 1, 1]
+// CHECK: %[[F0:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[S0]]
+// CHECK: %[[B0:.+]] = tensor.insert_slice %[[F0]] into %[[EMPTY]][0, 0, 0] [1, 3, 5] [1, 1, 1]
+// CHECK: %[[S1:.+]] = tensor.extract_slice %[[B0]][1, 2, 0] [2, 1, 5] [1, 1, 1]
+// CHECK: %[[F1:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[S1]]
+// CHECK: %[[B1:.+]] = tensor.insert_slice %[[F1]] into %[[B0]][1, 2, 0] [2, 1, 5] [1, 1, 1]
+// CHECK: %[[S2:.+]] = tensor.extract_slice %[[B1]][1, 0, 0] [2, 2, 2] [1, 1, 1]
+// CHECK: %[[F2:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[S2]]
+// CHECK: %[[B2:.+]] = tensor.insert_slice %[[F2]] into %[[B1]][1, 0, 0] [2, 2, 2] [1, 1, 1]
+// CHECK: %[[S3:.+]] = tensor.extract_slice %[[B2]][1, 0, 4] [2, 2, 1] [1, 1, 1]
+// CHECK: %[[F3:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[S3]]
+// CHECK: %[[BOUNDARY:.+]] = tensor.insert_slice %[[F3]] into %[[B2]][1, 0, 4] [2, 2, 1] [1, 1, 1]
+// CHECK-NOT: linalg.fill
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[BOUNDARY]][1, 0, 2] [2, 2, 2] [1, 1, 1]
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: outs(%[[SLICE]] :
+// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[GENERIC]] into %[[BOUNDARY]][1, 0, 2] [2, 2, 2] [1, 1, 1]
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @fill_boundary_transposed_init_map(
+ %arg0 : tensor<4x3xf32>, %arg1 : f32) -> tensor<5x6xf32> {
+ %init = tensor.empty() : tensor<3x4xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d1, d0)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x3xf32>) outs(%init : tensor<3x4xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.mulf %arg2, %arg2 : f32
+ linalg.yield %1 : f32
+ } -> tensor<3x4xf32>
+ %1 = tensor.pad %0 low [1, 2] high [1, 0] {
+ ^bb0(%arg2: index, %arg3 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<3x4xf32> to tensor<5x6xf32>
+ return %1 : tensor<5x6xf32>
+}
+
+// CHECK: func @fill_boundary_transposed_init_map
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<4x3xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<5x6xf32>
+// CHECK: %[[TOP_S:.+]] = tensor.extract_slice %[[EMPTY]][0, 0] [1, 6] [1, 1]
+// CHECK: %[[TOP_F:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[TOP_S]]
+// CHECK: %[[TOP:.+]] = tensor.insert_slice %[[TOP_F]] into %[[EMPTY]][0, 0] [1, 6] [1, 1]
+// CHECK: %[[BOT_S:.+]] = tensor.extract_slice %[[TOP]][4, 0] [1, 6] [1, 1]
+// CHECK: %[[BOT_F:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[BOT_S]]
+// CHECK: %[[BOT:.+]] = tensor.insert_slice %[[BOT_F]] into %[[TOP]][4, 0] [1, 6] [1, 1]
+// CHECK: %[[LEFT_S:.+]] = tensor.extract_slice %[[BOT]][1, 0] [3, 2] [1, 1]
+// CHECK: %[[LEFT_F:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[LEFT_S]]
+// CHECK: %[[BOUNDARY:.+]] = tensor.insert_slice %[[LEFT_F]] into %[[BOT]][1, 0] [3, 2] [1, 1]
+// CHECK-NOT: linalg.fill
+// CHECK: %[[SLICE:.+]] = tensor.extract_slice %[[BOUNDARY]][1, 2] [3, 4] [1, 1]
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: outs(%[[SLICE]] :
+// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[GENERIC]] into %[[BOUNDARY]][1, 2] [3, 4] [1, 1]
+// CHECK: return %[[RESULT]]
+
+// -----
+
+func.func @no_fill_boundary_producer_reads_init(
+ %arg0 : tensor<4x3xf32>, %arg1 : f32) -> tensor<7x6xf32> {
+ %init = tensor.empty() : tensor<4x3xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x3xf32>) outs(%init : tensor<4x3xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.addf %arg2, %arg3 : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x3xf32>
+ %1 = tensor.pad %0 low [1, 2] high [2, 1] {
+ ^bb0(%arg2: index, %arg3 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<4x3xf32> to tensor<7x6xf32>
+ return %1 : tensor<7x6xf32>
+}
+
+// CHECK: func @no_fill_boundary_producer_reads_init
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<4x3xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<7x6xf32>
+// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[EMPTY]] : tensor<7x6xf32>)
+// CHECK-NOT: linalg.fill
+// CHECK: tensor.extract_slice %[[FILL]][1, 2] [4, 3] [1, 1]
+
+// -----
+
+func.func @no_fill_boundary_projected_init_map(
+ %arg0 : tensor<4x?xf32>, %arg1 : f32) -> tensor<6xf32> {
+ %init = tensor.empty() : tensor<4xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x?xf32>) outs(%init : tensor<4xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ linalg.yield %arg2 : f32
+ } -> tensor<4xf32>
+ %1 = tensor.pad %0 low [1] high [1] {
+ ^bb0(%arg2: index):
+ tensor.yield %arg1 : f32
+ } : tensor<4xf32> to tensor<6xf32>
+ return %1 : tensor<6xf32>
+}
+
+// CHECK: func @no_fill_boundary_projected_init_map
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<4x?xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<6xf32>
+// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[EMPTY]] : tensor<6xf32>)
+// CHECK-NOT: linalg.fill
+// CHECK: tensor.extract_slice %[[FILL]][1] [4] [1]
+
+// -----
+
+func.func @no_fill_boundary_dynamic_source(
+ %arg0 : tensor<?x3xf32>, %arg1 : f32, %arg2 : index) -> tensor<?x6xf32> {
+ %init = tensor.empty(%arg2) : tensor<?x3xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<?x3xf32>) outs(%init : tensor<?x3xf32>) {
+ ^bb0(%arg3 : f32, %arg4 : f32):
+ %1 = arith.mulf %arg3, %arg3 : f32
+ linalg.yield %1 : f32
+ } -> tensor<?x3xf32>
+ %1 = tensor.pad %0 low [1, 2] high [1, 1] {
+ ^bb0(%arg3: index, %arg4 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<?x3xf32> to tensor<?x6xf32>
+ return %1 : tensor<?x6xf32>
+}
+
+// CHECK: func @no_fill_boundary_dynamic_source
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x3xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty(%{{.+}}) : tensor<?x6xf32>
+// CHECK: %[[FILL:.+]] = linalg.fill ins(%[[ARG1]] : f32) outs(%[[EMPTY]] : tensor<?x6xf32>)
+// CHECK-NOT: linalg.fill
+// CHECK: tensor.extract_slice %[[FILL]][1, 2]
+
+// -----
+
+// `nofold` keeps the pad op from being folded away before the pattern runs.
+func.func @fill_boundary_zero_padding(
+ %arg0 : tensor<4x3xf32>, %arg1 : f32) -> tensor<4x3xf32> {
+ %init = tensor.empty() : tensor<4x3xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x3xf32>) outs(%init : tensor<4x3xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.mulf %arg2, %arg2 : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x3xf32>
+ %1 = tensor.pad %0 nofold low [0, 0] high [0, 0] {
+ ^bb0(%arg2: index, %arg3 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<4x3xf32> to tensor<4x3xf32>
+ return %1 : tensor<4x3xf32>
+}
+
+// CHECK: func @fill_boundary_zero_padding
+// CHECK-NOT: linalg.fill
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<4x3xf32>
+// CHECK: %[[GENERIC:.+]] = linalg.generic
+// CHECK-SAME: outs(%[[EMPTY]] :
+// CHECK-NOT: linalg.fill
+// CHECK: return %[[GENERIC]]
+
+// -----
+
+func.func @fill_boundary_second_result(
+ %input : tensor<4xf32>, %init0 : tensor<4xf32>, %pad : f32) -> tensor<7xf32> {
+ %init1 = tensor.empty() : tensor<4xf32>
+ %results:2 = linalg.generic {
+ indexing_maps = [affine_map<(d0) -> (d0)>,
+ affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+ iterator_types = ["parallel"]}
+ ins(%input : tensor<4xf32>) outs(%init0, %init1 : tensor<4xf32>, tensor<4xf32>) {
+ ^bb0(%in : f32, %out0 : f32, %out1 : f32):
+ %sum = arith.addf %in, %out0 : f32
+ linalg.yield %sum, %in : f32, f32
+ } -> (tensor<4xf32>, tensor<4xf32>)
+ %padded = tensor.pad %results#1 low [1] high [2] {
+ ^bb0(%i : index):
+ tensor.yield %pad : f32
+ } : tensor<4xf32> to tensor<7xf32>
+ return %padded : tensor<7xf32>
+}
+
+// CHECK-LABEL: func @fill_boundary_second_result(
+// CHECK-SAME: %[[INPUT:.*]]: tensor<4xf32>, %[[INIT0:.*]]: tensor<4xf32>, %[[PAD:.*]]: f32
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<7xf32>
+// CHECK-NEXT: %[[LOW:.*]] = tensor.extract_slice %[[EMPTY]][0] [1] [1]
+// CHECK-NEXT: %[[LOW_FILL:.*]] = linalg.fill ins(%[[PAD]] : f32) outs(%[[LOW]]
+// CHECK-NEXT: %[[LOW_INSERT:.*]] = tensor.insert_slice %[[LOW_FILL]] into %[[EMPTY]][0] [1] [1]
+// CHECK-NEXT: %[[HIGH:.*]] = tensor.extract_slice %[[LOW_INSERT]][5] [2] [1]
+// CHECK-NEXT: %[[HIGH_FILL:.*]] = linalg.fill ins(%[[PAD]] : f32) outs(%[[HIGH]]
+// CHECK-NEXT: %[[BOUNDARY:.*]] = tensor.insert_slice %[[HIGH_FILL]] into %[[LOW_INSERT]][5] [2] [1]
+// CHECK-NEXT: %[[INTERIOR:.*]] = tensor.extract_slice %[[BOUNDARY]][1] [4] [1]
+// CHECK-NEXT: %[[RESULTS:.*]]:2 = linalg.generic
+// CHECK-SAME: outs(%[[INIT0]], %[[INTERIOR]] : tensor<4xf32>, tensor<4xf32>)
+// CHECK: %[[RESULT:.*]] = tensor.insert_slice %[[RESULTS]]#1 into %[[BOUNDARY]][1] [4] [1]
+// CHECK-NEXT: return %[[RESULT]]
+
+
+// -----
+
+func.func @fill_boundary_empty_source_first_dim(
+ %arg0 : tensor<0x3xf32>, %arg1 : f32) -> tensor<3x6xf32> {
+ %init = tensor.empty() : tensor<0x3xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<0x3xf32>) outs(%init : tensor<0x3xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.mulf %arg2, %arg2 : f32
+ linalg.yield %1 : f32
+ } -> tensor<0x3xf32>
+ %1 = tensor.pad %0 low [1, 2] high [2, 1] {
+ ^bb0(%arg2: index, %arg3 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<0x3xf32> to tensor<3x6xf32>
+ return %1 : tensor<3x6xf32>
+}
+
+// CHECK: func @fill_boundary_empty_source_first_dim
+// CHECK: tensor.empty() : tensor<3x6xf32>
+// CHECK: tensor.extract_slice %{{.+}}[0, 0] [1, 6] [1, 1]
+// CHECK: linalg.fill
+// CHECK: tensor.extract_slice %{{.+}}[1, 0] [2, 6] [1, 1]
+// CHECK: linalg.fill
+// CHECK-NOT: linalg.fill
+// CHECK: return
+
+// -----
+
+func.func @fill_boundary_empty_source_later_dim(
+ %arg0 : tensor<3x0xf32>, %arg1 : f32) -> tensor<6x3xf32> {
+ %init = tensor.empty() : tensor<3x0xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<3x0xf32>) outs(%init : tensor<3x0xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.mulf %arg2, %arg2 : f32
+ linalg.yield %1 : f32
+ } -> tensor<3x0xf32>
+ %1 = tensor.pad %0 low [1, 2] high [2, 1] {
+ ^bb0(%arg2: index, %arg3 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<3x0xf32> to tensor<6x3xf32>
+ return %1 : tensor<6x3xf32>
+}
+
+// CHECK: func @fill_boundary_empty_source_later_dim
+// CHECK: tensor.empty() : tensor<6x3xf32>
+// CHECK: tensor.extract_slice %{{.+}}[0, 0] [1, 3] [1, 1]
+// CHECK: linalg.fill
+// CHECK: tensor.extract_slice %{{.+}}[4, 0] [2, 3] [1, 1]
+// CHECK: linalg.fill
+// CHECK: tensor.extract_slice %{{.+}}[1, 0] [3, 2] [1, 1]
+// CHECK: linalg.fill
+// CHECK: tensor.extract_slice %{{.+}}[1, 2] [3, 1] [1, 1]
+// CHECK: linalg.fill
+// CHECK-NOT: linalg.fill
+// CHECK: return
+
+// -----
+
+func.func @fill_boundary_dynamic_high_pad(
+ %arg0 : tensor<4x3xf32>, %arg1 : f32, %arg2 : index) -> tensor<7x6xf32> {
+ %init = tensor.empty() : tensor<4x3xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x3xf32>) outs(%init : tensor<4x3xf32>) {
+ ^bb0(%arg3 : f32, %arg4 : f32):
+ %1 = arith.mulf %arg3, %arg3 : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x3xf32>
+ %1 = tensor.pad %0 low [1, 2] high [%arg2, 1] {
+ ^bb0(%arg3: index, %arg4 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<4x3xf32> to tensor<7x6xf32>
+ return %1 : tensor<7x6xf32>
+}
+
+// CHECK: func @fill_boundary_dynamic_high_pad
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[EMPTY:.+]] = tensor.empty() : tensor<7x6xf32>
+// CHECK: tensor.extract_slice %[[EMPTY]][0, 0] [1, 6] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32)
+// CHECK: tensor.extract_slice %{{.+}}[5, 0] [2, 6] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32)
+// CHECK: tensor.extract_slice %{{.+}}[1, 0] [4, 2] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32)
+// CHECK: tensor.extract_slice %{{.+}}[1, 5] [4, 1] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32)
+// CHECK-NOT: linalg.fill
+// CHECK: tensor.extract_slice %{{.+}}[1, 2] [4, 3] [1, 1]
+// CHECK: linalg.generic
diff --git a/mlir/test/Dialect/Linalg/pad_fusion_fill_boundary_bufferize.mlir b/mlir/test/Dialect/Linalg/pad_fusion_fill_boundary_bufferize.mlir
new file mode 100644
index 00000000000000..651a153a6733a4
--- /dev/null
+++ b/mlir/test/Dialect/Linalg/pad_fusion_fill_boundary_bufferize.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-opt %s -test-linalg-pad-fusion=fill-boundary-only \
+// RUN: -one-shot-bufferize="bufferize-function-boundaries" -cse -canonicalize \
+// RUN: > %t
+// RUN: FileCheck %s < %t
+// RUN: FileCheck %s --check-prefix=ONEALLOC < %t
+// RUN: FileCheck %s --check-prefix=NOCOPY < %t
+
+// Separate prefixes: a CHECK-NOT would only cover the gap between two checks.
+func.func @pad_fusion_boundary_bufferizes_in_place(
+ %arg0 : tensor<4x3xf32>, %arg1 : f32) -> tensor<7x6xf32> {
+ %init = tensor.empty() : tensor<4x3xf32>
+ %0 = linalg.generic {
+ indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],
+ iterator_types = ["parallel", "parallel"]}
+ ins(%arg0 : tensor<4x3xf32>) outs(%init : tensor<4x3xf32>) {
+ ^bb0(%arg2 : f32, %arg3 : f32):
+ %1 = arith.mulf %arg2, %arg2 : f32
+ linalg.yield %1 : f32
+ } -> tensor<4x3xf32>
+ %1 = tensor.pad %0 low [1, 2] high [2, 1] {
+ ^bb0(%arg2: index, %arg3 : index):
+ tensor.yield %arg1 : f32
+ } : tensor<4x3xf32> to tensor<7x6xf32>
+ return %1 : tensor<7x6xf32>
+}
+
+// CHECK: func @pad_fusion_boundary_bufferizes_in_place
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: memref<4x3xf32
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: f32
+// CHECK: %[[ALLOC:.+]] = memref.alloc() {{.*}} : memref<7x6xf32>
+// CHECK: %[[TOP:.+]] = memref.subview %[[ALLOC]][0, 0] [1, 6] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32) outs(%[[TOP]]
+// CHECK: %[[BOT:.+]] = memref.subview %[[ALLOC]][5, 0] [2, 6] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32) outs(%[[BOT]]
+// CHECK: %[[LEFT:.+]] = memref.subview %[[ALLOC]][1, 0] [4, 2] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32) outs(%[[LEFT]]
+// CHECK: %[[RIGHT:.+]] = memref.subview %[[ALLOC]][1, 5] [4, 1] [1, 1]
+// CHECK: linalg.fill ins(%[[ARG1]] : f32) outs(%[[RIGHT]]
+// CHECK: %[[INTERIOR:.+]] = memref.subview %[[ALLOC]][1, 2] [4, 3] [1, 1]
+// CHECK: linalg.generic
+// CHECK-SAME: outs(%[[INTERIOR]]
+// CHECK: return %[[ALLOC]]
+
+// ONEALLOC-COUNT-1: memref.alloc
+// ONEALLOC-NOT: memref.alloc
+// NOCOPY-NOT: memref.copy
diff --git a/mlir/test/lib/Dialect/Linalg/TestPadFusion.cpp b/mlir/test/lib/Dialect/Linalg/TestPadFusion.cpp
index e60a2e105da833..a5f659d63035c2 100644
--- a/mlir/test/lib/Dialect/Linalg/TestPadFusion.cpp
+++ b/mlir/test/lib/Dialect/Linalg/TestPadFusion.cpp
@@ -24,6 +24,16 @@ struct TestPadFusionPass
: public PassWrapper<TestPadFusionPass, OperationPass<>> {
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestPadFusionPass)
+ TestPadFusionPass() = default;
+ TestPadFusionPass(const TestPadFusionPass &pass) : PassWrapper(pass) {}
+
+ Option<bool> fillBoundaryOnly{
+ *this, "fill-boundary-only",
+ llvm::cl::desc("Fill only the padded boundary instead of the whole "
+ "result, for static pads whose producer overwrites the "
+ "interior"),
+ llvm::cl::init(false)};
+
void getDependentDialects(DialectRegistry ®istry) const override {
registry.insert<affine::AffineDialect, linalg::LinalgDialect,
tensor::TensorDialect>();
@@ -35,7 +45,8 @@ struct TestPadFusionPass
void runOnOperation() override {
MLIRContext *context = &getContext();
RewritePatternSet patterns(context);
- linalg::populateFuseTensorPadWithProducerLinalgOpPatterns(patterns);
+ linalg::populateFuseTensorPadWithProducerLinalgOpPatterns(patterns,
+ fillBoundaryOnly);
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
return signalPassFailure();
}
More information about the Mlir-commits
mailing list