[mlir][linalg] Extract GeneralizePadOpPattern into a standalone transformation
#117329
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
…nsformation Currently, `GeneralizePadOpPattern` is grouped under `populatePadOpVectorizationPatterns`. However, as noted in llvm#111349, this transformation "decomposes" rather than "vectorizes" `tensor.pad`. As such, it functions as: * a vectorization _pre-processing_ transformation, not * a vectorization transformation itself. To clarify its purpose, this PR turns `GeneralizePadOpPattern` into a standalone transformation by: * introducing a dedicated `populateDecomposePadPatterns` method, * adding a `apply_patterns.linalg.decompose_pad` Transform Dialect Op, and * removing it from `populatePadOpVectorizationPatterns`. In addition, to better reflect its role, it is renamed as "decomposition" rather then "generalization". That's to better reflect its role. This is in line with the recent renaming of similar ops, i.e. tensor.pack/tensor.unpack Ops in llvm#116439.
banach-space
requested review from
ftynse,
nicolasvasilache,
dcaballe,
hanhanW,
Groverkss and
rengolin
as code owners
|
@llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-linalg Author: Andrzej Warzyński (banach-space) ChangesCurrently,
To clarify its purpose, this PR turns
In addition, to better reflect its role, it is renamed as Full diff: https://github.com/llvm/llvm-project/pull/117329.diff 9 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
index e3084530bd11b5..dc10f3a1c58ae3 100644
--- a/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
+++ b/mlir/include/mlir/Dialect/Linalg/TransformOps/LinalgTransformOps.td
@@ -52,6 +52,17 @@ def ApplyDecomposeTensorPackUnpackPatternsOp
let assemblyFormat = "attr-dict";
}
+def ApplyDecomposeTensorPadPatternsOp
+ : Op<Transform_Dialect, "apply_patterns.linalg.decompose_pad",
+ [DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
+ let description = [{
+ Collect patterns to decompose tensor.pad into e.g. tensor::EmptyOp,
+ linalg::FillOp and tensor::InsertSliceOp.
+ }];
+
+ let assemblyFormat = "attr-dict";
+}
+
def ApplyFoldUnitExtentDimsViaReshapesPatternsOp : Op<Transform_Dialect,
"apply_patterns.linalg.fold_unit_extent_dims_via_reshapes",
[DeclareOpInterfaceMethods<PatternDescriptorOpInterface>]> {
diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 51967f83fee377..3c160d55a38e75 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -1503,8 +1503,8 @@ using OptimizeCopyFn =
/// Rewrite a tensor::PadOp into a sequence of EmptyOp, FillOp and
/// InsertSliceOp. For now, only constant padding values are supported.
-struct GeneralizePadOpPattern : public OpRewritePattern<tensor::PadOp> {
- GeneralizePadOpPattern(MLIRContext *context, PatternBenefit benefit = 1)
+struct DecomposePadOpPattern : public OpRewritePattern<tensor::PadOp> {
+ DecomposePadOpPattern(MLIRContext *context, PatternBenefit benefit = 1)
: OpRewritePattern<tensor::PadOp>(context, benefit) {}
LogicalResult matchAndRewrite(tensor::PadOp padOp,
PatternRewriter &rewriter) const override;
@@ -1688,6 +1688,10 @@ void populateDecomposeConvolutionPatterns(RewritePatternSet &patterns,
/// outer dims to be unit.
void populateDecomposePackUnpackPatterns(RewritePatternSet &patterns);
+/// Populates patterns to decompose tensor.pad into e.g.
+/// tensor.empty, linalg.fill, tensor.insert_slice.
+void populateDecomposePadPatterns(RewritePatternSet &patterns);
+
/// Populates patterns to transform linalg.conv_2d_xxx operations into
/// linalg.generic (for img2col packing) and linalg.matmul.
/// \see rewriteInIm2Col for more details.
diff --git a/mlir/lib/Conversion/TensorToLinalg/TensorToLinalg.cpp b/mlir/lib/Conversion/TensorToLinalg/TensorToLinalg.cpp
index 5bb79d4bc84e2b..b0ca0ca13d0624 100644
--- a/mlir/lib/Conversion/TensorToLinalg/TensorToLinalg.cpp
+++ b/mlir/lib/Conversion/TensorToLinalg/TensorToLinalg.cpp
@@ -25,5 +25,7 @@ using namespace mlir;
//===----------------------------------------------------------------------===//
void mlir::populateTensorToLinalgPatterns(RewritePatternSet &patterns) {
- patterns.add<mlir::linalg::GeneralizePadOpPattern>(patterns.getContext());
+ // TODO: Add the remaining patterns, e.g. to decompose Pack/Unpack Ops.
+ // Alternatively, delete this file.
+ patterns.add<mlir::linalg::DecomposePadOpPattern>(patterns.getContext());
}
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index ada80deacfdbfe..e08be7d2ebd6ae 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -234,6 +234,11 @@ void transform::ApplyDecomposeTensorPackUnpackPatternsOp::populatePatterns(
linalg::populateDecomposePackUnpackPatterns(patterns);
}
+void transform::ApplyDecomposeTensorPadPatternsOp::populatePatterns(
+ RewritePatternSet &patterns) {
+ linalg::populateDecomposePadPatterns(patterns);
+}
+
void transform::ApplyFoldUnitExtentDimsViaReshapesPatternsOp::populatePatterns(
RewritePatternSet &patterns) {
linalg::ControlDropUnitDims options;
@@ -3491,8 +3496,12 @@ transform::VectorizeChildrenAndApplyPatternsOp::applyToOne(
// Add misc. vectorization patterns (e.g. for tensor.insert_slice)
linalg::populateInsertSliceVectorizationPatterns(patterns);
- if (getVectorizePadding())
+ if (getVectorizePadding()) {
linalg::populatePadOpVectorizationPatterns(patterns);
+ // This creates an alternative path for lowering tensor.pad - by
+ // decomposing it into e.g. linalg.fill.
+ linalg::populateDecomposePadPatterns(patterns);
+ }
vector::populateVectorStepLoweringPatterns(patterns);
TrackingListener listener(state, *this);
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
index d92543d7264625..c3e176299317ef 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Transforms.cpp
@@ -921,7 +921,7 @@ LogicalResult mlir::linalg::CopyVectorizationPattern::matchAndRewrite(
/// Filling `dest` using FillOp constant padding value if possible.
/// Otherwise, generate a tensor::GenerateOp.
-Value GeneralizePadOpPattern::createFillOrGenerateOp(
+Value DecomposePadOpPattern::createFillOrGenerateOp(
RewriterBase &rewriter, tensor::PadOp padOp, Value dest,
const SmallVector<Value> &dynSizes) const {
auto padValue = padOp.getConstantPaddingValue();
@@ -938,8 +938,8 @@ Value GeneralizePadOpPattern::createFillOrGenerateOp(
}
LogicalResult
-GeneralizePadOpPattern::matchAndRewrite(tensor::PadOp padOp,
- PatternRewriter &rewriter) const {
+DecomposePadOpPattern::matchAndRewrite(tensor::PadOp padOp,
+ PatternRewriter &rewriter) const {
// Given an OpFoldResult, return an index-typed value.
auto getIdxValue = [&](OpFoldResult ofr) {
if (auto val = llvm::dyn_cast_if_present<Value>(ofr))
@@ -1623,3 +1623,7 @@ void linalg::populateDecomposePackUnpackPatterns(RewritePatternSet &patterns) {
// TODO: Add and test patterns for tensor.unpack
patterns.add<DecomposeOuterUnitDimsPackOpPattern>(patterns.getContext());
}
+
+void linalg::populateDecomposePadPatterns(RewritePatternSet &patterns) {
+ patterns.add<DecomposePadOpPattern>(patterns.getContext());
+}
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index 23b46a2ee55f8d..06bb6c0fb1cac9 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2770,12 +2770,6 @@ void mlir::linalg::populateInsertSliceVectorizationPatterns(
void mlir::linalg::populatePadOpVectorizationPatterns(
RewritePatternSet &patterns, PatternBenefit baseBenefit) {
- // TODO: The following pattern implements "decomposition" and
- // optional "vectorization". Seperate "decomposition" into a sepereate
- // pre-processing pattern group.
- patterns.add<GeneralizePadOpPattern>(patterns.getContext(), baseBenefit);
-
- // Try these specialized patterns first before resorting to the generic one.
patterns.add<PadOpVectorizationWithTransferReadPattern,
PadOpVectorizationWithTransferWritePattern,
PadOpVectorizationWithInsertSlicePattern>(
diff --git a/mlir/test/Dialect/Linalg/generalize-pad-tensor.mlir b/mlir/test/Dialect/Linalg/decompose-pad-tensor.mlir
similarity index 98%
rename from mlir/test/Dialect/Linalg/generalize-pad-tensor.mlir
rename to mlir/test/Dialect/Linalg/decompose-pad-tensor.mlir
index 2beab31b613d54..184361dfb30dfd 100644
--- a/mlir/test/Dialect/Linalg/generalize-pad-tensor.mlir
+++ b/mlir/test/Dialect/Linalg/decompose-pad-tensor.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-generalize-pad-tensor" %s | FileCheck %s
+// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-decompose-pad-tensor" %s | FileCheck %s
// CHECK-LABEL: func @generalize_pad_tensor_static_shape(
// CHECK-SAME: %[[IN:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
diff --git a/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir b/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir
index 640de85cc5f12e..41e480648177f5 100644
--- a/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir
+++ b/mlir/test/Dialect/Linalg/vectorization-pad-patterns.mlir
@@ -202,6 +202,8 @@ module attributes {transform.with_named_sequence} {
%func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op {
+ // TODO: Split into two tests, one for each pattern
+ transform.apply_patterns.linalg.decompose_pad
transform.apply_patterns.linalg.pad_vectorization
} : !transform.op<"func.func">
transform.yield
@@ -236,6 +238,8 @@ module attributes {transform.with_named_sequence} {
%func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op {
+ // TODO: Split into two tests, one for each pattern
+ transform.apply_patterns.linalg.decompose_pad
transform.apply_patterns.linalg.pad_vectorization
} : !transform.op<"func.func">
transform.yield
@@ -270,6 +274,8 @@ module attributes {transform.with_named_sequence} {
%func_op = transform.structured.match ops{["func.func"]} in %arg1 : (!transform.any_op) -> !transform.op<"func.func">
transform.apply_patterns to %func_op {
+ // TODO: Split into two tests, one for each pattern
+ transform.apply_patterns.linalg.decompose_pad
transform.apply_patterns.linalg.pad_vectorization
} : !transform.op<"func.func">
transform.yield
diff --git a/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp b/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
index c65e68eaf31f09..25aec75c3c14ad 100644
--- a/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
+++ b/mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp
@@ -70,8 +70,8 @@ struct TestLinalgTransforms
llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
"in vector.contract form"),
llvm::cl::init(false)};
- Option<bool> testGeneralizePadTensor{
- *this, "test-generalize-pad-tensor",
+ Option<bool> testDecomposePadTensor{
+ *this, "test-decompose-pad-tensor",
llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
llvm::cl::init(false)};
Option<bool> testDecomposeTensorPackOp{
@@ -166,9 +166,9 @@ static void applyLinalgToVectorPatterns(func::FuncOp funcOp) {
(void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}
-static void applyGeneralizePadTensorPatterns(func::FuncOp funcOp) {
+static void applyDecomposePadPatterns(func::FuncOp funcOp) {
RewritePatternSet patterns(funcOp.getContext());
- patterns.add<GeneralizePadOpPattern>(funcOp.getContext());
+ patterns.add<DecomposePadOpPattern>(funcOp.getContext());
(void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
}
@@ -235,8 +235,8 @@ void TestLinalgTransforms::runOnOperation() {
return applyVectorTransferForwardingPatterns(getOperation());
if (testGenericToVectorPattern)
return applyLinalgToVectorPatterns(getOperation());
- if (testGeneralizePadTensor)
- return applyGeneralizePadTensorPatterns(getOperation());
+ if (testDecomposePadTensor)
+ return applyDecomposePadPatterns(getOperation());
if (testDecomposeTensorPackOp)
return applyDecomposeTensorPackPatterns(getOperation());
if (testDecomposeTensorUnPackOp)
|
|
The reasoning here looks good to me, haven't looked at the PR yet. Side comment: this is another "layering violation", where |
|
Thanks for taking a look!
Note that your RFC covers Now, given that this is merely moving things around within (effectively) Linalg and making sure that we don't mix "decomposition" and "vectorization", I am hoping that this is non-controversial 😅 |
Indeed, didn't mean to detract from your goals, just a side note. (and yes, I know you were on to it from the beginning 😄). |
banach-space
added a commit
to banach-space/llvm-project
that referenced
this pull request
Nov 25, 2024
[mlir][linalg][nfc] Update pack-dynamic-inner-tile.mlir Builds on: * llvm#117329: Extract GeneralizePadOpPattern into a standalone transformation. * llvm#116373: Update pack-dynamic-inner-tile.mlir. This update adds vectorization to the "pack-dynamic-inner-tile.mlir" pipeline. The pipeline first decomposes `tensor.pack` into `tensor.pad` and then into `linalg.fill` (llvm#117329). Next, `linalg.fill` is vectorized, with vector sizes matching the inner tile sizes of the original `tensor.pack`. ••NOTE:** Depends on llvm#117329 - please only review the top commit!
hanhanW
approved these changes
Nov 25, 2024
banach-space
added a commit
to banach-space/llvm-project
that referenced
this pull request
Nov 26, 2024
[mlir][linalg][nfc] Update pack-dynamic-inner-tile.mlir Builds on: * llvm#117329: Extract GeneralizePadOpPattern into a standalone transformation. * llvm#116373: Update pack-dynamic-inner-tile.mlir. This update adds vectorization to the "pack-dynamic-inner-tile.mlir" pipeline. The pipeline first decomposes `tensor.pack` into `tensor.pad` and then into `linalg.fill` (llvm#117329). Next, `linalg.fill` is vectorized, with vector sizes matching the inner tile sizes of the original `tensor.pack`. ••NOTE:** Depends on llvm#117329 - please only review the top commit!
banach-space
added a commit
that referenced
this pull request
Nov 26, 2024
Builds on: * #117329: "Extract GeneralizePadOpPattern into a standalone transformation". * #116373: "Update pack-dynamic-inner-tile.mlir". This update adds vectorization to the "pack-dynamic-inner-tile.mlir" pipeline. The pipeline first decomposes `tensor.pack` into `tensor.pad` and then into `linalg.fill` (#117329). Next, `linalg.fill` is vectorized, with vector sizes matching the inner tile sizes of the original `tensor.pack`.
banach-space
added a commit
to banach-space/llvm-project
that referenced
this pull request
Jan 13, 2025
…tern * Relocates two tests for `PadOpVectorizationWithTransferWritePattern` in "vectorization-pad-patterns.mlir" to group them with other tests for the same pattern. * Adds a note clarifying that these are negative tests and explains the reasoning behind them. * Removes `transform.apply_patterns.linalg.decompose_pad` from the TD sequences as it's no longer needed (*). This is essentially a small clean-up in preparation for upcoming changes. (*) `transform.apply_patterns.linalg.decompose_pad` was split off from `transform.apply_patterns.linalg.pad_vectorization` in llvm#117329. "vectorization-pad-patterns.mlir" is meant to test the latter, not the former.
banach-space
added a commit
to banach-space/llvm-project
that referenced
this pull request
Jan 14, 2025
…tern * Relocates two tests for `PadOpVectorizationWithTransferWritePattern` in "vectorization-pad-patterns.mlir" to group them with other tests for the same pattern. * Adds a note clarifying that these are negative tests and explains the reasoning behind them. * Removes `transform.apply_patterns.linalg.decompose_pad` from the TD sequences as it's no longer needed (*). This is essentially a small clean-up in preparation for upcoming changes. (*) `transform.apply_patterns.linalg.decompose_pad` was split off from `transform.apply_patterns.linalg.pad_vectorization` in llvm#117329. "vectorization-pad-patterns.mlir" is meant to test the latter, not the former.
banach-space
added a commit
that referenced
this pull request
Jan 14, 2025
…tern (#122721) * Relocates two tests for `PadOpVectorizationWithTransferWritePattern` in "vectorization-pad-patterns.mlir" to group them with other tests for the same pattern. * Adds a note clarifying that these are negative tests and explains the reasoning behind them. * Removes `transform.apply_patterns.linalg.decompose_pad` from the TD sequences as it's no longer needed (*). This is essentially a small clean-up in preparation for upcoming changes. (*) `transform.apply_patterns.linalg.decompose_pad` was split off from `transform.apply_patterns.linalg.pad_vectorization` in #117329. "vectorization-pad-patterns.mlir" is meant to test the latter, not the former.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Currently,
GeneralizePadOpPatternis grouped underpopulatePadOpVectorizationPatterns. However, as noted in #111349, thistransformation "decomposes" rather than "vectorizes"
tensor.pad. Assuch, it functions as:
To clarify its purpose, this PR turns
GeneralizePadOpPatterninto astandalone transformation by:
populateDecomposePadPatternsmethod,apply_patterns.linalg.decompose_padTransform Dialect Op,populatePadOpVectorizationPatterns.In addition, to better reflect its role, it is renamed as "decomposition"
rather then "generalization". This is in line with the recent renaming
of similar ops, i.e. tensor.pack/tensor.unpack Ops in #116439.