1 //===- CodegenStrategy.cpp - Linalg programmable codegen strategy ---------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements logic and helpers to expose Linalg transforms as
10 // composable rewrite patterns through a programmable CodegenStrategy object.
11 //
12 //===----------------------------------------------------------------------===//
13
14 #include "mlir/Dialect/Linalg/Transforms/CodegenStrategy.h"
15
16 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
17 #include "mlir/Dialect/Vector/VectorOps.h"
18 #include "mlir/Dialect/Vector/VectorTransforms.h"
19 #include "mlir/Pass/PassManager.h"
20 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
21 #include "mlir/Transforms/LoopUtils.h"
22 #include "mlir/Transforms/Passes.h"
23
24 using namespace mlir;
25 using namespace mlir::linalg;
26
27 #define DEBUG_TYPE "linalg-codegen-strategy"
28
transform(FuncOp func) const29 void mlir::linalg::CodegenStrategy::transform(FuncOp func) const {
30 MLIRContext *context = func.getContext();
31 // Emplace patterns one at a time while also maintaining a simple chained
32 // state transition.
33 unsigned stepCount = 0;
34 SmallVector<FrozenRewritePatternList, 4> stage1Patterns;
35 auto zeroState = Identifier::get(std::to_string(stepCount), context);
36 auto currentState = zeroState;
37 for (const std::unique_ptr<Transformation> &t : transformationSequence) {
38 auto nextState = Identifier::get(std::to_string(++stepCount), context);
39 auto marker = (currentState == zeroState)
40 ? linalg::LinalgMarker({}, nextState)
41 : linalg::LinalgMarker(currentState, nextState);
42 stage1Patterns.emplace_back(t->buildRewritePatterns(context, marker));
43 currentState = nextState;
44 }
45
46 OwningRewritePatternList stage2Patterns =
47 linalg::getLinalgTilingCanonicalizationPatterns(context);
48 stage2Patterns.insert<AffineMinSCFCanonicalizationPattern>(context);
49
50 auto stage3Transforms = [](Operation *op) {
51 // Some of these may be too aggressive as a stage 3 that is applied on each
52 // stage 1 application and may have to be split out to post staged patterns
53 // application (in which case they could just be passes, TBD).
54 PassManager pm(op->getContext());
55 pm.addPass(createLoopInvariantCodeMotionPass());
56 if (failed(pm.run(op->getParentOfType<ModuleOp>())))
57 llvm_unreachable("Unexpected failure in cleanup pass pipeline.");
58 promoteSingleIterationLoops(cast<FuncOp>(op));
59 hoistViewAllocOps(cast<FuncOp>(op));
60 hoistRedundantVectorTransfers(cast<FuncOp>(op));
61 return success();
62 };
63 linalg::applyStagedPatterns(func, stage1Patterns, std::move(stage2Patterns),
64 stage3Transforms);
65
66 //===--------------------------------------------------------------------===//
67 // Post staged patterns transforms
68 //===--------------------------------------------------------------------===//
69
70 ModuleOp module = func->getParentOfType<ModuleOp>();
71
72 // Programmatic splitting of slow/fast path vector transfers.
73 OwningRewritePatternList patterns;
74 patterns.insert<vector::VectorTransferFullPartialRewriter>(
75 context, vectorTransformsOptions);
76 applyPatternsAndFoldGreedily(module, std::move(patterns));
77
78 // Programmatic controlled lowering of vector.contract only.
79 OwningRewritePatternList vectorContractLoweringPatterns;
80 vectorContractLoweringPatterns
81 .insert<ContractionOpToOuterProductOpLowering,
82 ContractionOpToMatmulOpLowering, ContractionOpLowering>(
83 vectorTransformsOptions, context);
84 applyPatternsAndFoldGreedily(module,
85 std::move(vectorContractLoweringPatterns));
86
87 // Programmatic controlled lowering of vector.transfer only.
88 OwningRewritePatternList vectorToLoopsPatterns;
89 populateVectorToSCFConversionPatterns(vectorToLoopsPatterns, context,
90 vectorToSCFOptions);
91 applyPatternsAndFoldGreedily(module, std::move(vectorToLoopsPatterns));
92
93 // Ensure we drop the marker in the end.
94 module.walk([](LinalgOp op) {
95 op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
96 });
97 }
98