Add additional options to the ApplyPatternsOp (#12519)
The following upstream options are added to the ApplyPatternsOp:
1. rank-reduction of linalg ops via reshapes
2. pack/unpack propagation
3. split bubble_expand from bubble_collapse as these patterns exhibit an
interference behavior
4. patterns for greedy fusion of linalg ops
All these are extensively tested upstream, this PR provides the plumbing
to be used with the transform dialect.
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
index eefb051..55b08e0 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.cpp
@@ -139,7 +139,9 @@
/// When touching something here, do not forget to update CommonExtensions.h.
///
ADD_PATTERN(additionalIreePatterns, getAdditionalIreePatternsAttrName)
- ADD_PATTERN(bubbleCollapseExpand, getBubbleCollapseExpandAttrName)
+ ADD_PATTERN(bubbleCollapse, getBubbleCollapseAttrName)
+ ADD_PATTERN(bubbleExpand, getBubbleExpandAttrName)
+ ADD_PATTERN(bubblePackUnPack, getBubblePackUnPackAttrName)
ADD_PATTERN(canonicalization, getCanonicalizationAttrName)
ADD_PATTERN(cse, getCseAttrName)
ADD_PATTERN(eraseUnnecessaryTensorOperands,
@@ -150,9 +152,13 @@
ADD_PATTERN(foldReassociativeReshapes, getFoldReassociativeReshapesAttrName)
ADD_PATTERN(foldTensorEmptyExtract, getFoldTensorEmptyExtractAttrName)
ADD_PATTERN(licm, getLicmAttrName)
+ ADD_PATTERN(linalgElementwiseGreedyFusion,
+ getLinalgElementwiseGreedyFusionAttrName)
ADD_PATTERN(lowerTransferOpPermutations,
getLowerTransferOpPermutationsAttrName)
ADD_PATTERN(rankReducingLinalg, getRankReducingLinalgAttrName)
+ ADD_PATTERN(rankReducingLinalgViaReshapes,
+ getRankReducingLinalgViaReshapesAttrName)
ADD_PATTERN(rankReducingVector, getRankReducingVectorAttrName)
ADD_PATTERN(swapPaddingElideConditional,
getSwapPaddingElideConditionalAttrName)
@@ -164,6 +170,31 @@
result.addTypes({pdl::OperationType::get(ctx)});
}
+static void addOperands(Operation *op, SetVector<Value> &operandSet) {
+ if (!op) return;
+ TypeSwitch<Operation *, void>(op)
+ .Case<linalg::LinalgOp>([&](linalg::LinalgOp linalgOp) {
+ SmallVector<Value> inputOperands{linalgOp.getDpsInputOperands()};
+ operandSet.insert(inputOperands.begin(), inputOperands.end());
+ })
+ .Default([&](Operation *operation) {
+ operandSet.insert(operation->operand_begin(), operation->operand_end());
+ });
+}
+
+template <int limit = 3>
+static bool setFusedOpOperandLimit(OpOperand *fusedOperand) {
+ Operation *producer = fusedOperand->get().getDefiningOp();
+ if (!producer) return false;
+ Operation *consumer = fusedOperand->getOwner();
+ SetVector<Value> fusedOpOperands;
+ if (producer->getNumResults() != 1) return false;
+ addOperands(consumer, fusedOpOperands);
+ fusedOpOperands.remove(producer->getResult(0));
+ addOperands(producer, fusedOpOperands);
+ return fusedOpOperands.size() <= limit;
+}
+
namespace {
/// Rewrite a tensor.generate as an arith.constant when possible.
struct GenerateToConstant : public OpRewritePattern<tensor::GenerateOp> {
@@ -232,6 +263,12 @@
linalg::populateFoldUnitExtentDimsViaSlicesPatterns(patterns);
}
+static void addRankReducingLinalgViaReshapesPatterns(
+ RewritePatternSet &patterns) {
+ populateReshapeToInterfaceTensorPatterns(patterns);
+ linalg::populateFoldUnitExtentDimsViaReshapesPatterns(patterns);
+}
+
static void addRankReducingVectorPatterns(RewritePatternSet &patterns) {
vector::populateCastAwayVectorLeadingOneDimPatterns(patterns);
}
@@ -309,10 +346,16 @@
MLIRContext *ctx = target->getContext();
RewritePatternSet patterns(ctx);
if (getAdditionalIreePatterns()) addAdditionalIreePatterns(patterns);
- if (getBubbleCollapseExpand()) {
+ if (getBubbleCollapse()) {
+ linalg::populateFoldReshapeOpsByCollapsingPatterns(
+ patterns, [](OpOperand *) { return true; });
+ }
+ if (getBubbleExpand()) {
linalg::populateFoldReshapeOpsByExpansionPatterns(
patterns, [](OpOperand *) { return true; });
}
+ if (getBubblePackUnPack())
+ linalg::populateDataLayoutPropagationPatterns(patterns);
if (getCanonicalization()) addAllRegisteredCanonicalizationPatterns(patterns);
if (getEraseUnnecessaryTensorOperands())
addEraseUnnecessaryTensorOperandsPatterns(patterns);
@@ -321,9 +364,14 @@
if (getFoldMemrefAliases()) addFoldMemrefAliasPatterns(patterns);
if (getFoldReassociativeReshapes()) addReassociativeReshapePatterns(patterns);
if (getFoldTensorEmptyExtract()) addFoldTensorEmptyExtract(patterns);
+ if (getLinalgElementwiseGreedyFusion())
+ linalg::populateElementwiseOpsFusionPatterns(patterns,
+ setFusedOpOperandLimit<3>);
if (getLowerTransferOpPermutations())
addLowerTransferOpPermutationsPatterns(patterns);
if (getRankReducingLinalg()) addRankReducingLinalgPatterns(patterns);
+ if (getRankReducingLinalgViaReshapes())
+ addRankReducingLinalgViaReshapesPatterns(patterns);
if (getRankReducingVector()) addRankReducingVectorPatterns(patterns);
if (getSwappingPatterns())
addSwappingPatterns(patterns, getSwapPaddingElideConditional());
@@ -568,7 +616,8 @@
return forallOp->emitError("mapping must be #gpu.block<x/y/z/>");
}
- // Step 1. Complete the blockMapping to a full mapping (with 1s) if necessary.
+ // Step 1. Complete the blockMapping to a full mapping (with 1s) if
+ // necessary.
SmallVector<Value> numBlocks =
llvm::to_vector(forallOp.getUpperBound(rewriter));
// Ensure we have 3 block sizes, one for each id.
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h
index c1064f6..797c3c0 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h
+++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensions.h
@@ -36,7 +36,9 @@
/// Selected patterns for ApplyPatternOp.
struct ApplyPatternsOpPatterns {
bool additionalIreePatterns = false;
- bool bubbleCollapseExpand = false;
+ bool bubbleCollapse = false;
+ bool bubbleExpand = false;
+ bool bubblePackUnPack = false;
bool canonicalization = false;
bool cse = false;
bool eraseUnnecessaryTensorOperands = false;
@@ -45,9 +47,11 @@
bool foldReassociativeReshapes = false;
bool foldTensorEmptyExtract = false;
bool licm = false;
+ bool linalgElementwiseGreedyFusion = false;
bool lowerTransferOpPermutations = false;
bool promoteForallCaptureToShared = false;
bool rankReducingLinalg = false;
+ bool rankReducingLinalgViaReshapes = false;
bool rankReducingVector = false;
bool swapPaddingElideConditional = false;
bool swappingPatterns = false;
diff --git a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td
index a55c29d..e588b60 100644
--- a/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td
+++ b/compiler/src/iree/compiler/Codegen/Common/TransformExtensions/CommonExtensionsOps.td
@@ -74,8 +74,14 @@
unspecified order:
- additional_iree_patterns: fancy patterns we shortcut into the system,
will need to be sliced out better in the future.
- - bubble_collapse_expand: bubble `expand_shape` up and `collapse_shape`
- down across Linalg ops.
+ - bubble_collapse: bubble `collapse_shape` down across Linalg ops. This
+ must be applied separately from `bubble_expand` patterns because of some
+ upstream pattern interference issue atm.
+ - bubble_expand: bubble `expand_shape` down across Linalg ops. This
+ must be applied separately from `bubble_collapse` patterns because of some
+ upstream pattern interference issue atm.
+ - bubble_pack_un_pack: bubble `pack` up and `unpack` down across Linalg
+ ops.
- canonicalization: adds all the canonicalization patterns of all
registered dialects and ops.
- cse: additionally apply common subexpression elimination. This must
@@ -97,10 +103,14 @@
iteration loop promotion. This is not a set of patterns per se but is still
very convenient to apply it close to canonicalization and other greedy
pattern applications.
+ - linalg_elementwise_greedy_fusion: add linalg elementwise ops fusion
+ patterns using a naive default heuristic.
- lower_transfer_op_permutations: Lower transfer ops to transfer ops
with minor identity permutations.
- rank_reducing_linalg: adds patterns that results in rank-reducing
- behavior on subset-based linalg operations.
+ behavior on subset-based linalg operations using insert/extract slices.
+ - rank_reducing_linalg_via_reshapes: adds patterns that results in rank-reducing
+ behavior on subset-based linalg operations using expand/collapse shape ops.
- rank_reducing_vector: adds patterns that results in rank-reducing
behavior on subset-based vector operations.
adopts the upstream version.
@@ -135,7 +145,9 @@
let arguments = (ins PDL_Operation:$target,
UnitAttr:$additional_iree_patterns,
- UnitAttr:$bubble_collapse_expand,
+ UnitAttr:$bubble_collapse,
+ UnitAttr:$bubble_expand,
+ UnitAttr:$bubble_pack_un_pack,
UnitAttr:$canonicalization,
UnitAttr:$cse,
UnitAttr:$erase_unnecessary_tensor_operands,
@@ -144,8 +156,10 @@
UnitAttr:$fold_reassociative_reshapes,
UnitAttr:$fold_tensor_empty_extract,
UnitAttr:$licm,
+ UnitAttr:$linalg_elementwise_greedy_fusion,
UnitAttr:$lower_transfer_op_permutations,
UnitAttr:$rank_reducing_linalg,
+ UnitAttr:$rank_reducing_linalg_via_reshapes,
UnitAttr:$rank_reducing_vector,
UnitAttr:$swap_padding_elide_conditional,
UnitAttr:$swapping_patterns,
@@ -183,8 +197,8 @@
It always return success.
}];
- let arguments = (ins Transform_ConcreteOpType<"func.func">:$target);
- let results = (outs Transform_ConcreteOpType<"func.func">:$result);
+ let arguments = (ins TransformHandleTypeInterface:$target);
+ let results = (outs TransformHandleTypeInterface:$result);
let assemblyFormat = "$target attr-dict `:` functional-type(operands, results)";
let cppNamespace = "mlir::iree_compiler::IREE::transform_dialect";
@@ -411,10 +425,10 @@
}];
let arguments = (
- ins Transform_ConcreteOpType<"scf.forall">:$forall_op,
+ ins TransformHandleTypeInterface:$forall_op,
DefaultValuedOptionalAttr<DenseI64ArrayAttr, "{}">:$share_operands
);
- let results = (outs Transform_ConcreteOpType<"scf.forall">:$result);
+ let results = (outs TransformHandleTypeInterface:$result);
let cppNamespace = "mlir::iree_compiler::IREE::transform_dialect";
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir b/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir
index 8811b76..6131ead 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/reductions_codegen_spec.mlir
@@ -18,7 +18,7 @@
( mapping = [#gpu.block<x>] )
%func = transform.structured.match ops{["func.func"]} in %arg0 : (!pdl.operation) -> !pdl.operation
- %func_1 = transform.iree.apply_patterns %func { bubble_collapse_expand }
+ %func_1 = transform.iree.apply_patterns %func { bubble_expand }
// Excessively eager canonicalization results in `fill`s being "fused" due to
// swapping with `extract_slice`, which confuses the fusion operation below.
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/transform_dialect_apply_pattern_op.mlir b/compiler/src/iree/compiler/Codegen/Common/test/transform_dialect_apply_pattern_op.mlir
index b74075a..c3eebc0 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/transform_dialect_apply_pattern_op.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/transform_dialect_apply_pattern_op.mlir
@@ -87,5 +87,5 @@
transform.sequence failures(propagate) {
^bb1(%arg1: !pdl.operation):
%0 = transform.structured.match ops{["func.func"]} in %arg1 : (!pdl.operation) -> !pdl.operation
- transform.iree.apply_patterns %0 { bubble_collapse_expand }
+ transform.iree.apply_patterns %0 { bubble_expand }
}
diff --git a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp
index 3417d9c..dbddae9 100644
--- a/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp
+++ b/compiler/src/iree/compiler/Codegen/TransformDialectStrategies/Common/Common.cpp
@@ -341,7 +341,7 @@
/// produced as parent of reduction splitting if necessary for fusion of the
/// leading elementwise operation.
// TODO: consider passing a problem-specific struct to control information.
-static ReductionSplitResult createExpansionBubbleUp(
+static ReductionSplitResult createBubbleExpand(
ImplicitLocOpBuilder &b, Value variantH,
SplitReductionOp splitReductionTransformOp, bool hasLeadingEltwise,
bool hasTrailingEltwise) {
@@ -355,7 +355,7 @@
auto funcH = b.create<MatchOp>(variantH, func::FuncOp::getOperationName());
ApplyPatternsOpPatterns configuration;
- configuration.bubbleCollapseExpand = true;
+ configuration.bubbleExpand = true;
b.create<ApplyPatternsOp>(funcH, configuration);
std::tie(result.originalFillH, result.splitFillH) =
matchAndUnpack<2>(b, variantH, linalg::FillOp::getOperationName());
diff --git a/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir b/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir
index 5d13de7..e39aed5 100644
--- a/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir
+++ b/tests/transform_dialect/cuda/eltwise_reduction_codegen_spec.mlir
@@ -25,7 +25,7 @@
// able to preserve the handles.
// ===========================================================================
%func = transform.structured.match ops{["func.func"]} in %variant_op : (!pdl.operation) -> !pdl.operation
- transform.iree.apply_patterns %func { bubble_collapse_expand }
+ transform.iree.apply_patterns %func { bubble_expand }
%fills = transform.structured.match ops{["linalg.fill"]} in %variant_op : (!pdl.operation) -> !pdl.operation
%fill_2, %more_parallel_fill_2 = transform.split_handles %fills in [2]
: (!pdl.operation) -> (!pdl.operation, !pdl.operation)
diff --git a/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir b/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir
index 69af78b..bd5115c 100644
--- a/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir
+++ b/tests/transform_dialect/cuda/eltwise_reduction_eltwise_codegen_spec.mlir
@@ -27,7 +27,7 @@
// able to preserve the handles.
// ===========================================================================
%func = transform.structured.match ops{["func.func"]} in %variant_op : (!pdl.operation) -> !pdl.operation
- transform.iree.apply_patterns %func { bubble_collapse_expand }
+ transform.iree.apply_patterns %func { bubble_expand }
%fills = transform.structured.match ops{["linalg.fill"]} in %variant_op : (!pdl.operation) -> !pdl.operation
%fill_2, %more_parallel_fill_2 = transform.split_handles %fills in [2]
: (!pdl.operation) -> (!pdl.operation, !pdl.operation)
