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     vector::populateVectorTransferPermutationMapLoweringPatterns(
193         vectorizationPatterns);
194     vector::populateVetorReductionToContractPatterns(vectorizationPatterns);
195     vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
196                               linalg::LinalgCopyVTWForwardingPattern>(
197         funcOp.getContext(), /*benefit=*/2);
198     (void)applyPatternsAndFoldGreedily(funcOp,
199                                        std::move(vectorizationPatterns));
200   }
201 
202   LinalgVectorizationOptions options;
203   LinalgTransformationFilter filter;
204 };
205 
206 /// Configurable pass to enable the application of other pattern-based linalg
207 /// passes.
208 struct LinalgStrategyEnablePass
209     : public LinalgStrategyEnablePassBase<LinalgStrategyEnablePass> {
210 
211   LinalgStrategyEnablePass(LinalgEnablingOptions opt,
212                            LinalgTransformationFilter filt)
213       : options(opt), filter(filt) {}
214 
215   void runOnFunction() override {
216     auto funcOp = getFunction();
217     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
218       return;
219 
220     MLIRContext *context = funcOp.getContext();
221     RewritePatternSet patterns =
222         linalg::getLinalgTilingCanonicalizationPatterns(context);
223     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
224     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
225       return signalPassFailure();
226 
227     if (options.enableLICM) {
228       if (funcOp
229               ->walk([&](LoopLikeOpInterface loopLike) {
230                 if (failed(moveLoopInvariantCode(loopLike)))
231                   return WalkResult::interrupt();
232                 return WalkResult::advance();
233               })
234               .wasInterrupted())
235         return signalPassFailure();
236     }
237 
238     promoteSingleIterationLoops(funcOp);
239     if (options.enableHoistRedundantVectorTransfers)
240       hoistRedundantVectorTransfers(funcOp);
241 
242     if (options.enableHoistRedundantVectorTransfersOnTensor)
243       hoistRedundantVectorTransfersOnTensor(funcOp);
244   }
245 
246   LinalgEnablingOptions options;
247   LinalgTransformationFilter filter;
248 };
249 
250 /// Configurable pass to lower vector operations.
251 struct LinalgStrategyLowerVectorsPass
252     : public LinalgStrategyLowerVectorsPassBase<
253           LinalgStrategyLowerVectorsPass> {
254 
255   LinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
256                                  LinalgTransformationFilter filt)
257       : options(opt), filter(filt) {}
258 
259   void runOnFunction() override {
260     auto funcOp = getFunction();
261     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
262       return;
263 
264     MLIRContext *context = funcOp.getContext();
265     RewritePatternSet patterns(context);
266     if (options.enableVectorTransferLowering) {
267       vector::populateVectorTransferLoweringPatterns(patterns,
268                                                      options.maxTransferRank);
269     }
270     if (options.enableVectorTransferPartialRewrite) {
271       patterns.add<vector::VectorTransferFullPartialRewriter>(
272           context, options.vectorTransformOptions);
273     }
274     if (options.enableVectorContractLowering) {
275       patterns.add<ContractionOpToOuterProductOpLowering,
276                    ContractionOpToMatmulOpLowering, ContractionOpLowering>(
277           options.vectorTransformOptions, context);
278       vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
279     }
280     if (options.enableVectorToSCFConversion) {
281       populateVectorToSCFConversionPatterns(patterns,
282                                             options.vectorTransferToSCFOptions);
283     }
284     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
285   }
286 
287   LinalgVectorLoweringOptions options;
288   LinalgTransformationFilter filter;
289 };
290 } // namespace
291 
292 /// Create a LinalgStrategyTilePass.
293 std::unique_ptr<OperationPass<FuncOp>>
294 mlir::createLinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
295                                    LinalgTransformationFilter filter) {
296   return std::make_unique<LinalgStrategyTilePass>(opName, opt, filter);
297 }
298 
299 /// Create a LinalgStrategyPromotePass.
300 std::unique_ptr<OperationPass<FuncOp>>
301 mlir::createLinalgStrategyPromotePass(StringRef opName,
302                                       LinalgPromotionOptions opt,
303                                       LinalgTransformationFilter filter) {
304   return std::make_unique<LinalgStrategyPromotePass>(opName, opt, filter);
305 }
306 
307 /// Create a LinalgStrategyGeneralizePass.
308 std::unique_ptr<OperationPass<FuncOp>>
309 mlir::createLinalgStrategyGeneralizePass(StringRef opName,
310                                          LinalgTransformationFilter filter) {
311   return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
312 }
313 
314 /// Create a LinalgStrategyInterchangePass.
315 std::unique_ptr<OperationPass<FuncOp>>
316 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
317                                           LinalgTransformationFilter filter) {
318   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
319                                                          filter);
320 }
321 
322 /// Create a LinalgStrategyVectorizePass.
323 std::unique_ptr<OperationPass<FuncOp>>
324 mlir::createLinalgStrategyVectorizePass(StringRef opName,
325                                         LinalgVectorizationOptions opt,
326                                         LinalgTransformationFilter filter) {
327   return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter);
328 }
329 
330 /// Create a LinalgStrategyEnablePass.
331 std::unique_ptr<OperationPass<FuncOp>>
332 mlir::createLinalgStrategyEnablePass(LinalgEnablingOptions opt,
333                                      LinalgTransformationFilter filter) {
334   return std::make_unique<LinalgStrategyEnablePass>(opt, filter);
335 }
336 
337 /// Create a LinalgStrategyLowerVectorsPass.
338 std::unique_ptr<OperationPass<FuncOp>>
339 mlir::createLinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
340                                            LinalgTransformationFilter filter) {
341   return std::make_unique<LinalgStrategyLowerVectorsPass>(opt, filter);
342 }
343