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/VectorTransforms.h"
26 #include "mlir/IR/AffineExpr.h"
27 #include "mlir/IR/AffineMap.h"
28 #include "mlir/Support/LLVM.h"
29 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
30 #include "mlir/Transforms/LoopUtils.h"
31 #include "mlir/Transforms/Utils.h"
32 
33 using namespace mlir;
34 using namespace mlir::vector;
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::populateVectorReductionToContractPatterns(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.licm) {
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.hoistRedundantVectorTransfers)
240       hoistRedundantVectorTransfers(funcOp);
241 
242     if (options.hoistRedundantVectorTransfersOnTensor)
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     vector::populateVectorToVectorCanonicalizationPatterns(patterns);
267     if (options.transferLowering) {
268       vector::populateVectorTransferLoweringPatterns(patterns,
269                                                      options.maxTransferRank);
270     }
271     if (options.transposeLowering) {
272       vector::populateVectorTransposeLoweringPatterns(
273           patterns, options.vectorTransformOptions);
274     }
275     if (options.multiReductionLowering) {
276       vector::populateVectorMultiReductionLoweringPatterns(
277           patterns,
278           options.vectorTransformOptions.vectorMultiReductionLowering);
279     }
280     if (options.contractionLowering) {
281       patterns.add<ContractionOpToOuterProductOpLowering,
282                    ContractionOpToMatmulOpLowering, ContractionOpLowering>(
283           options.vectorTransformOptions, context);
284       vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
285     }
286     if (options.transferPartialRewrite) {
287       patterns.add<vector::VectorTransferFullPartialRewriter>(
288           context, options.vectorTransformOptions);
289     }
290     if (options.transferToSCFConversion) {
291       populateVectorToSCFConversionPatterns(patterns,
292                                             options.vectorTransferToSCFOptions);
293     }
294     vector::populateVectorShapeCastLoweringPatterns(patterns);
295     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
296   }
297 
298   LinalgVectorLoweringOptions options;
299   LinalgTransformationFilter filter;
300 };
301 
302 /// Configurable pass to lower vector operations.
303 struct LinalgStrategyRemoveMarkersPass
304     : public LinalgStrategyRemoveMarkersPassBase<
305           LinalgStrategyRemoveMarkersPass> {
306 
307   void runOnFunction() override {
308     auto funcOp = getFunction();
309     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
310       return;
311     funcOp.walk([](LinalgOp op) {
312       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
313     });
314   }
315 };
316 } // namespace
317 
318 /// Create a LinalgStrategyTilePass.
319 std::unique_ptr<OperationPass<FuncOp>>
320 mlir::createLinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
321                                    LinalgTransformationFilter filter) {
322   return std::make_unique<LinalgStrategyTilePass>(opName, opt, filter);
323 }
324 
325 /// Create a LinalgStrategyPromotePass.
326 std::unique_ptr<OperationPass<FuncOp>>
327 mlir::createLinalgStrategyPromotePass(StringRef opName,
328                                       LinalgPromotionOptions opt,
329                                       LinalgTransformationFilter filter) {
330   return std::make_unique<LinalgStrategyPromotePass>(opName, opt, filter);
331 }
332 
333 /// Create a LinalgStrategyGeneralizePass.
334 std::unique_ptr<OperationPass<FuncOp>>
335 mlir::createLinalgStrategyGeneralizePass(StringRef opName,
336                                          LinalgTransformationFilter filter) {
337   return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
338 }
339 
340 /// Create a LinalgStrategyInterchangePass.
341 std::unique_ptr<OperationPass<FuncOp>>
342 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
343                                           LinalgTransformationFilter filter) {
344   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
345                                                          filter);
346 }
347 
348 /// Create a LinalgStrategyVectorizePass.
349 std::unique_ptr<OperationPass<FuncOp>>
350 mlir::createLinalgStrategyVectorizePass(StringRef opName,
351                                         LinalgVectorizationOptions opt,
352                                         LinalgTransformationFilter filter) {
353   return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter);
354 }
355 
356 /// Create a LinalgStrategyEnablePass.
357 std::unique_ptr<OperationPass<FuncOp>>
358 mlir::createLinalgStrategyEnablePass(LinalgEnablingOptions opt,
359                                      LinalgTransformationFilter filter) {
360   return std::make_unique<LinalgStrategyEnablePass>(opt, filter);
361 }
362 
363 /// Create a LinalgStrategyLowerVectorsPass.
364 std::unique_ptr<OperationPass<FuncOp>>
365 mlir::createLinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
366                                            LinalgTransformationFilter filter) {
367   return std::make_unique<LinalgStrategyLowerVectorsPass>(opt, filter);
368 }
369 
370 /// Create a LinalgStrategyRemoveMarkersPass.
371 std::unique_ptr<OperationPass<FuncOp>>
372 mlir::createLinalgStrategyRemoveMarkersPass() {
373   return std::make_unique<LinalgStrategyRemoveMarkersPass>();
374 }
375