1 //===- DynamicPass.cpp - Implementation of a dynamic configurable pass ----===//
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 a configurable pass that can apply patterns liberally
10 // and be plugged in a pass pipeline.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #include "PassDetail.h"
15 #include "mlir/Analysis/SliceAnalysis.h"
16 #include "mlir/Dialect/Affine/IR/AffineOps.h"
17 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
18 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
19 #include "mlir/Dialect/Linalg/Passes.h"
20 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
21 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
22 #include "mlir/Dialect/Linalg/Utils/Utils.h"
23 #include "mlir/Dialect/SCF/Transforms.h"
24 #include "mlir/Dialect/Tensor/IR/Tensor.h"
25 #include "mlir/Dialect/Vector/VectorOps.h"
26 #include "mlir/Dialect/Vector/VectorTransforms.h"
27 #include "mlir/IR/AffineExpr.h"
28 #include "mlir/IR/AffineMap.h"
29 #include "mlir/Support/LLVM.h"
30 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
31 #include "mlir/Transforms/LoopUtils.h"
32 #include "mlir/Transforms/Utils.h"
33 
34 using namespace mlir;
35 using namespace linalg;
36 
37 namespace {
38 
39 /// Configurable pass to apply pattern-based linalg tiling.
40 struct LinalgStrategyTilePass
41     : public LinalgStrategyTilePassBase<LinalgStrategyTilePass> {
42 
43   LinalgStrategyTilePass() = default;
44 
45   LinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
46                          LinalgTransformationFilter filt)
47       : options(opt), filter(filt) {
48     this->anchorOpName.setValue(opName.str());
49   }
50 
51   void runOnFunction() override {
52     auto funcOp = getFunction();
53     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
54       return;
55 
56     RewritePatternSet tilingPattern(funcOp.getContext());
57     if (!anchorOpName.empty()) {
58       tilingPattern.add<LinalgGenericTilingPattern>(
59           anchorOpName, funcOp.getContext(), options, filter);
60     } else {
61       tilingPattern.add<LinalgGenericTilingPattern>(funcOp.getContext(), filter,
62                                                     options);
63     }
64     (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
65   }
66 
67   LinalgTilingOptions options;
68   LinalgTransformationFilter filter;
69 };
70 
71 /// Configurable pass to apply pattern-based linalg generalization.
72 struct LinalgStrategyGeneralizePass
73     : public LinalgStrategyGeneralizePassBase<LinalgStrategyGeneralizePass> {
74 
75   LinalgStrategyGeneralizePass() = default;
76 
77   LinalgStrategyGeneralizePass(StringRef opName,
78                                LinalgTransformationFilter filter)
79       : filter(filter) {
80     this->anchorOpName.setValue(opName.str());
81   }
82 
83   void runOnFunction() override {
84     auto funcOp = getFunction();
85     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
86       return;
87 
88     RewritePatternSet generalizationPattern(funcOp.getContext());
89     if (!anchorOpName.empty()) {
90       generalizationPattern.add<LinalgGeneralizationPattern>(
91           anchorOpName, funcOp.getContext(), filter);
92     } else {
93       generalizationPattern.add<LinalgGeneralizationPattern>(
94           funcOp.getContext(), filter);
95     }
96     if (failed(applyPatternsAndFoldGreedily(funcOp,
97                                             std::move(generalizationPattern))))
98       signalPassFailure();
99   }
100 
101   LinalgTransformationFilter filter;
102 };
103 
104 /// Configurable pass to apply pattern-based linalg generalization.
105 struct LinalgStrategyInterchangePass
106     : public LinalgStrategyInterchangePassBase<LinalgStrategyInterchangePass> {
107 
108   LinalgStrategyInterchangePass() = default;
109 
110   LinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
111                                 LinalgTransformationFilter filter)
112       : iteratorInterchange(iteratorInterchange.begin(),
113                             iteratorInterchange.end()),
114         filter(filter) {}
115 
116   void runOnFunction() override {
117     auto funcOp = getFunction();
118     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
119       return;
120 
121     SmallVector<unsigned> interchangeVector(iteratorInterchange.begin(),
122                                             iteratorInterchange.end());
123     RewritePatternSet interchangePattern(funcOp.getContext());
124     interchangePattern.add<GenericOpInterchangePattern>(
125         funcOp.getContext(), interchangeVector, filter);
126     if (failed(applyPatternsAndFoldGreedily(funcOp,
127                                             std::move(interchangePattern))))
128       signalPassFailure();
129   }
130 
131   SmallVector<int64_t> iteratorInterchange;
132   LinalgTransformationFilter filter;
133 };
134 
135 /// Configurable pass to apply pattern-based linalg promotion.
136 struct LinalgStrategyPromotePass
137     : public LinalgStrategyPromotePassBase<LinalgStrategyPromotePass> {
138 
139   LinalgStrategyPromotePass() = default;
140 
141   LinalgStrategyPromotePass(StringRef opName, LinalgPromotionOptions opt,
142                             LinalgTransformationFilter filt)
143       : options(opt), filter(filt) {
144     this->anchorOpName.setValue(opName.str());
145   }
146 
147   void runOnFunction() override {
148     auto funcOp = getFunction();
149     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
150       return;
151 
152     RewritePatternSet promotionPattern(funcOp.getContext());
153     if (!anchorOpName.empty()) {
154       promotionPattern.add<LinalgBasePromotionPattern>(
155           anchorOpName, funcOp.getContext(), options, filter);
156     } else {
157       promotionPattern.add<LinalgBasePromotionPattern>(funcOp.getContext(),
158                                                        filter, options);
159     }
160     (void)applyPatternsAndFoldGreedily(funcOp, std::move(promotionPattern));
161   }
162 
163   LinalgPromotionOptions options;
164   LinalgTransformationFilter filter;
165 };
166 
167 /// Configurable pass to apply pattern-based linalg vectorization.
168 struct LinalgStrategyVectorizePass
169     : public LinalgStrategyVectorizePassBase<LinalgStrategyVectorizePass> {
170 
171   LinalgStrategyVectorizePass() = default;
172 
173   LinalgStrategyVectorizePass(StringRef opName, LinalgVectorizationOptions opt,
174                               LinalgTransformationFilter filt)
175       : options(opt), filter(filt) {
176     this->anchorOpName.setValue(opName.str());
177   }
178 
179   void runOnFunction() override {
180     auto funcOp = getFunction();
181     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
182       return;
183 
184     RewritePatternSet vectorizationPatterns(funcOp.getContext());
185     if (!anchorOpName.empty()) {
186       vectorizationPatterns.add<LinalgVectorizationPattern>(
187           anchorOpName, funcOp.getContext(), options, filter);
188     } else {
189       vectorizationPatterns.add<LinalgVectorizationPattern>(funcOp.getContext(),
190                                                             filter, options);
191     }
192     vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
193                               linalg::LinalgCopyVTWForwardingPattern>(
194         funcOp.getContext(), /*benefit=*/2);
195     (void)applyPatternsAndFoldGreedily(funcOp,
196                                        std::move(vectorizationPatterns));
197   }
198 
199   LinalgVectorizationOptions options;
200   LinalgTransformationFilter filter;
201 };
202 
203 /// Configurable pass to enable the application of other pattern-based linalg
204 /// passes.
205 struct LinalgStrategyEnablePass
206     : public LinalgStrategyEnablePassBase<LinalgStrategyEnablePass> {
207 
208   LinalgStrategyEnablePass(LinalgEnablingOptions opt,
209                            LinalgTransformationFilter filt)
210       : options(opt), filter(filt) {}
211 
212   void runOnFunction() override {
213     auto funcOp = getFunction();
214     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
215       return;
216 
217     MLIRContext *context = funcOp.getContext();
218     RewritePatternSet patterns =
219         linalg::getLinalgTilingCanonicalizationPatterns(context);
220     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
221     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
222       return signalPassFailure();
223 
224     if (options.enableLICM) {
225       if (funcOp
226               ->walk([&](LoopLikeOpInterface loopLike) {
227                 if (failed(moveLoopInvariantCode(loopLike)))
228                   return WalkResult::interrupt();
229                 return WalkResult::advance();
230               })
231               .wasInterrupted())
232         return signalPassFailure();
233     }
234 
235     promoteSingleIterationLoops(funcOp);
236     if (options.enableHoistRedundantVectorTransfers)
237       hoistRedundantVectorTransfers(funcOp);
238 
239     if (options.enableHoistRedundantVectorTransfersOnTensor)
240       hoistRedundantVectorTransfersOnTensor(funcOp);
241   }
242 
243   LinalgEnablingOptions options;
244   LinalgTransformationFilter filter;
245 };
246 
247 /// Configurable pass to lower vector operations.
248 struct LinalgStrategyLowerVectorsPass
249     : public LinalgStrategyLowerVectorsPassBase<
250           LinalgStrategyLowerVectorsPass> {
251 
252   LinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
253                                  LinalgTransformationFilter filt)
254       : options(opt), filter(filt) {}
255 
256   void runOnFunction() override {
257     auto funcOp = getFunction();
258     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
259       return;
260 
261     MLIRContext *context = funcOp.getContext();
262     RewritePatternSet patterns(context);
263     if (options.enableVectorTransferLowering) {
264       vector::populateVectorTransferLoweringPatterns(patterns,
265                                                      options.maxTransferRank);
266     }
267     if (options.enableVectorTransferPartialRewrite) {
268       patterns.add<vector::VectorTransferFullPartialRewriter>(
269           context, options.vectorTransformOptions);
270     }
271     if (options.enableVectorContractLowering) {
272       patterns.add<ContractionOpToOuterProductOpLowering,
273                    ContractionOpToMatmulOpLowering, ContractionOpLowering>(
274           options.vectorTransformOptions, context);
275       vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
276     }
277     if (options.enableVectorToSCFConversion) {
278       populateVectorToSCFConversionPatterns(patterns,
279                                             options.vectorTransferToSCFOptions);
280     }
281     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
282   }
283 
284   LinalgVectorLoweringOptions options;
285   LinalgTransformationFilter filter;
286 };
287 } // namespace
288 
289 /// Create a LinalgStrategyTilePass.
290 std::unique_ptr<OperationPass<FuncOp>>
291 mlir::createLinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
292                                    LinalgTransformationFilter filter) {
293   return std::make_unique<LinalgStrategyTilePass>(opName, opt, filter);
294 }
295 
296 /// Create a LinalgStrategyPromotePass.
297 std::unique_ptr<OperationPass<FuncOp>>
298 mlir::createLinalgStrategyPromotePass(StringRef opName,
299                                       LinalgPromotionOptions opt,
300                                       LinalgTransformationFilter filter) {
301   return std::make_unique<LinalgStrategyPromotePass>(opName, opt, filter);
302 }
303 
304 /// Create a LinalgStrategyGeneralizePass.
305 std::unique_ptr<OperationPass<FuncOp>>
306 mlir::createLinalgStrategyGeneralizePass(StringRef opName,
307                                          LinalgTransformationFilter filter) {
308   return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
309 }
310 
311 /// Create a LinalgStrategyInterchangePass.
312 std::unique_ptr<OperationPass<FuncOp>>
313 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
314                                           LinalgTransformationFilter filter) {
315   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
316                                                          filter);
317 }
318 
319 /// Create a LinalgStrategyVectorizePass.
320 std::unique_ptr<OperationPass<FuncOp>>
321 mlir::createLinalgStrategyVectorizePass(StringRef opName,
322                                         LinalgVectorizationOptions opt,
323                                         LinalgTransformationFilter filter) {
324   return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter);
325 }
326 
327 /// Create a LinalgStrategyEnablePass.
328 std::unique_ptr<OperationPass<FuncOp>>
329 mlir::createLinalgStrategyEnablePass(LinalgEnablingOptions opt,
330                                      LinalgTransformationFilter filter) {
331   return std::make_unique<LinalgStrategyEnablePass>(opt, filter);
332 }
333 
334 /// Create a LinalgStrategyLowerVectorsPass.
335 std::unique_ptr<OperationPass<FuncOp>>
336 mlir::createLinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
337                                            LinalgTransformationFilter filter) {
338   return std::make_unique<LinalgStrategyLowerVectorsPass>(opt, filter);
339 }
340