Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[GlobalOpt] Add option to propagate transposes through conv #20339

Merged
merged 1 commit into from
Mar 21, 2025
Merged

[GlobalOpt] Add option to propagate transposes through conv #20339

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions compiler/src/iree/compiler/GlobalOptimization/Passes.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,8 @@ def PropagateLinalgTransposePass :
"Flag used for lit-testing bubbling patterns only. Not for general usage">,
Option<"enableAttentionVTranspose", "enable-attention-v-transpose", "bool",
/*default=*/"true", "Enable transposition of attention v operand">,
Option<"enableConvolutionPropagation", "enable-aggressive-propagation-through-conv", "bool",
/*default=*/"false", "enable propagation through convolutions">,
];
}

Expand Down
30 changes: 17 additions & 13 deletions compiler/src/iree/compiler/GlobalOptimization/PropagateLinalgTranspose.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,21 +136,23 @@ static void specializeGenericTransposeOp(RewriterBase &rewriter,
/// Returns the `op` if it is a linalg::GenericOp. If it is a named op and
/// `allowGeneralizing` is true, returns the generalized op. Otherwise, returns
/// failure.
/// TODO: Right now convolutions are not allowed due to fragility around
/// handling of convolutions.
/// TODO: Due to fragility around handling of convolutions, convolution
/// propagation is behind a flag.
static FailureOr<linalg::GenericOp>
getAllowedGenericOpOrGeneralizeNamedOp(RewriterBase &rewriter, Operation *op,
bool allowGeneralizing) {
bool allowGeneralizing, bool convProp) {
auto linalgOp = dyn_cast<linalg::LinalgOp>(op);
if (!linalgOp) {
return failure();
}

if (linalg::isaConvolutionOpInterface(linalgOp)) {
if (!convProp && linalg::isaConvolutionOpInterface(linalgOp)) {
return failure();
}

if (!isa<linalg::GenericOp>(linalgOp) &&
!(allowGeneralizing && linalg::isaContractionOpInterface(linalgOp))) {
!(allowGeneralizing && linalg::isaContractionOpInterface(linalgOp)) &&
!(convProp && linalg::isaConvolutionOpInterface(linalgOp))) {
return failure();
}

Expand Down Expand Up @@ -190,9 +192,9 @@ class FuseTransposeWithProducerLinalgOp
public:
using OpRewritePattern<linalg::TransposeOp>::OpRewritePattern;
FuseTransposeWithProducerLinalgOp(MLIRContext *ctx, bool aggressiveProp,
PatternBenefit b = 1)
bool convProp, PatternBenefit b = 1)
: OpRewritePattern<linalg::TransposeOp>(ctx, b),
allowGeneralizing(aggressiveProp) {}
allowGeneralizing(aggressiveProp), convProp(convProp) {}

LogicalResult matchAndRewrite(linalg::TransposeOp transposeOp,
PatternRewriter &rewriter) const override {
Expand All @@ -216,7 +218,7 @@ class FuseTransposeWithProducerLinalgOp

int64_t resultIndex = result.getResultNumber();
auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp(
rewriter, result.getOwner(), allowGeneralizing);
rewriter, result.getOwner(), allowGeneralizing, convProp);
if (failed(maybeGenericOp)) {
return rewriter.notifyMatchFailure(
transposeOp, "linalg op producer is not generic or contraction");
Expand Down Expand Up @@ -297,6 +299,7 @@ class FuseTransposeWithProducerLinalgOp

private:
bool allowGeneralizing = false;
bool convProp = false;
};

// Bubbles a transpose through a tensor.collapse_shape.
Expand Down Expand Up @@ -592,9 +595,9 @@ class FuseTransposeWithLinalgOpConsumer
public:
using OpInterfaceRewritePattern<linalg::LinalgOp>::OpInterfaceRewritePattern;
FuseTransposeWithLinalgOpConsumer(MLIRContext *ctx, bool aggressiveProp,
PatternBenefit b = 1)
bool convProp, PatternBenefit b = 1)
: OpInterfaceRewritePattern<linalg::LinalgOp>(ctx, b),
allowGeneralizing(aggressiveProp) {}
allowGeneralizing(aggressiveProp), convProp(convProp) {}

LogicalResult matchAndRewrite(linalg::LinalgOp linalgOp,
PatternRewriter &rewriter) const override {
Expand Down Expand Up @@ -623,7 +626,7 @@ class FuseTransposeWithLinalgOpConsumer
// to the results of the associated input's indexing map, and then forward
// the input to the transpose to the consumer generic.
auto maybeGenericOp = getAllowedGenericOpOrGeneralizeNamedOp(
rewriter, linalgOp, allowGeneralizing);
rewriter, linalgOp, allowGeneralizing, convProp);
if (failed(maybeGenericOp)) {
return failure();
}
Expand All @@ -649,6 +652,7 @@ class FuseTransposeWithLinalgOpConsumer

private:
bool allowGeneralizing = false;
bool convProp = false;
};

static bool isIndexingMapAffectedByTransposeMap(
Expand Down Expand Up @@ -1107,7 +1111,7 @@ void PropagateLinalgTransposePass::runOnOperation() {
bubblingPatterns, bubbleTransposeControlFn);
}
bubblingPatterns.insert<FuseTransposeWithProducerLinalgOp>(
context, enableAggressivePropagation);
context, enableAggressivePropagation, enableConvolutionPropagation);
bubblingPatterns.insert<BubbleTransposeThroughCollapseShape>(context);
bubblingPatterns.add<BubbleTransposeThroughUnaryElementwiseDpsInit>(
context, /*benefit=*/2);
Expand Down Expand Up @@ -1165,7 +1169,7 @@ void PropagateLinalgTransposePass::runOnOperation() {
sinkingPatterns.insert<SinkTransposeThroughExtractSlice>(context);
sinkingPatterns.insert<SinkTransposeThroughExpandShape>(context);
sinkingPatterns.insert<FuseTransposeWithLinalgOpConsumer>(
context, enableAggressivePropagation);
context, enableAggressivePropagation, enableConvolutionPropagation);
sinkingPatterns.insert<ComposeTransposes>(context);
populateNamedOpSinkingPatterns(context, sinkingPatterns);
populateCommonCanonicalizationPatterns(context, sinkingPatterns);
Expand Down
9 changes: 9 additions & 0 deletions compiler/src/iree/compiler/GlobalOptimization/test/propagate_linalg_transpose.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{enable-aggressive-propagation=true}))" --split-input-file %s | FileCheck %s --check-prefix=APROP
// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{test-sinking-only=true}))" --split-input-file %s | FileCheck %s --check-prefix=SINK
// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{test-bubbling-only=true}))" --split-input-file %s | FileCheck %s --check-prefix=BUBBLE
// RUN: iree-opt --pass-pipeline="builtin.module(util.func(iree-global-opt-propagate-linalg-transpose{enable-aggressive-propagation-through-conv=true}))" --split-input-file %s | FileCheck %s --check-prefix=CONV

util.func public @specialize_transpose_op(%arg0 : tensor<1x2x3xf32>,
%empty : tensor<3x2x1xf32>) -> tensor<3x2x1xf32> {
Expand Down Expand Up @@ -254,6 +255,10 @@ util.func public @do_not_propagate_to_conv(%transposed_lhs: tensor<18x2x18x8xf32
// APROP-LABEL: util.func public @do_not_propagate_to_conv
// APROP: linalg.conv_2d_nhwc_hwcf

// CONV-LABEL: util.func public @do_not_propagate_to_conv
// CONV-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor<18x2x18x8xf32>
// CONV: linalg.generic {{.*}} ins(%[[ARG0]]

// -----

#map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>
Expand All @@ -278,6 +283,10 @@ module {
// APROP: linalg.transpose
// APROP: linalg.generic

// CONV-LABEL: util.func public @do_not_propagate_to_conv_generic
// CONV-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor<18x2x18x8xf32>
// CONV: linalg.generic {{.*}} ins(%[[ARG0]]

// -----

util.func public @sink_through_expand_shape(%arg0: tensor<?x?x?xf32>) -> tensor<32x?x16x?x?xf32> {
Expand Down
Loading