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 hoisting and padding.
72 struct LinalgStrategyPadPass
73     : public LinalgStrategyPadPassBase<LinalgStrategyPadPass> {
74 
75   LinalgStrategyPadPass() = default;
76 
77   LinalgStrategyPadPass(StringRef opName, LinalgPaddingOptions opt,
78                         LinalgTransformationFilter filt)
79       : options(opt), filter(filt) {
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 paddingPattern(funcOp.getContext());
89     if (!anchorOpName.empty()) {
90       paddingPattern.add<LinalgPaddingPattern>(
91           anchorOpName, funcOp.getContext(), options, filter);
92     } else {
93       paddingPattern.add<LinalgPaddingPattern>(funcOp.getContext(), options,
94                                                filter);
95     }
96     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(paddingPattern))))
97       signalPassFailure();
98   }
99 
100   LinalgPaddingOptions options;
101   LinalgTransformationFilter filter;
102 };
103 
104 /// Configurable pass to apply pattern-based linalg generalization.
105 struct LinalgStrategyGeneralizePass
106     : public LinalgStrategyGeneralizePassBase<LinalgStrategyGeneralizePass> {
107 
108   LinalgStrategyGeneralizePass() = default;
109 
110   LinalgStrategyGeneralizePass(StringRef opName,
111                                LinalgTransformationFilter filter)
112       : filter(filter) {
113     this->anchorOpName.setValue(opName.str());
114   }
115 
116   void runOnFunction() override {
117     auto funcOp = getFunction();
118     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
119       return;
120 
121     RewritePatternSet generalizationPattern(funcOp.getContext());
122     if (!anchorOpName.empty()) {
123       generalizationPattern.add<LinalgGeneralizationPattern>(
124           anchorOpName, funcOp.getContext(), filter);
125     } else {
126       generalizationPattern.add<LinalgGeneralizationPattern>(
127           funcOp.getContext(), filter);
128     }
129     if (failed(applyPatternsAndFoldGreedily(funcOp,
130                                             std::move(generalizationPattern))))
131       signalPassFailure();
132   }
133 
134   LinalgTransformationFilter filter;
135 };
136 
137 /// Configurable pass to apply pattern-based linalg generalization.
138 struct LinalgStrategyInterchangePass
139     : public LinalgStrategyInterchangePassBase<LinalgStrategyInterchangePass> {
140 
141   LinalgStrategyInterchangePass() = default;
142 
143   LinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
144                                 LinalgTransformationFilter filter)
145       : iteratorInterchange(iteratorInterchange.begin(),
146                             iteratorInterchange.end()),
147         filter(filter) {}
148 
149   void runOnFunction() override {
150     auto funcOp = getFunction();
151     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
152       return;
153 
154     SmallVector<unsigned> interchangeVector(iteratorInterchange.begin(),
155                                             iteratorInterchange.end());
156     RewritePatternSet interchangePattern(funcOp.getContext());
157     interchangePattern.add<GenericOpInterchangePattern>(
158         funcOp.getContext(), interchangeVector, filter);
159     if (failed(applyPatternsAndFoldGreedily(funcOp,
160                                             std::move(interchangePattern))))
161       signalPassFailure();
162   }
163 
164   SmallVector<int64_t> iteratorInterchange;
165   LinalgTransformationFilter filter;
166 };
167 
168 /// Configurable pass to apply pattern-based linalg promotion.
169 struct LinalgStrategyPromotePass
170     : public LinalgStrategyPromotePassBase<LinalgStrategyPromotePass> {
171 
172   LinalgStrategyPromotePass() = default;
173 
174   LinalgStrategyPromotePass(StringRef opName, LinalgPromotionOptions opt,
175                             LinalgTransformationFilter filt)
176       : options(opt), filter(filt) {
177     this->anchorOpName.setValue(opName.str());
178   }
179 
180   void runOnFunction() override {
181     auto funcOp = getFunction();
182     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
183       return;
184 
185     RewritePatternSet promotionPattern(funcOp.getContext());
186     if (!anchorOpName.empty()) {
187       promotionPattern.add<LinalgBasePromotionPattern>(
188           anchorOpName, funcOp.getContext(), options, filter);
189     } else {
190       promotionPattern.add<LinalgBasePromotionPattern>(funcOp.getContext(),
191                                                        filter, options);
192     }
193     (void)applyPatternsAndFoldGreedily(funcOp, std::move(promotionPattern));
194   }
195 
196   LinalgPromotionOptions options;
197   LinalgTransformationFilter filter;
198 };
199 
200 /// Configurable pass to apply pattern-based linalg vectorization.
201 struct LinalgStrategyVectorizePass
202     : public LinalgStrategyVectorizePassBase<LinalgStrategyVectorizePass> {
203 
204   LinalgStrategyVectorizePass() = default;
205 
206   LinalgStrategyVectorizePass(StringRef opName, LinalgVectorizationOptions opt,
207                               LinalgTransformationFilter filt)
208       : options(opt), filter(filt) {
209     this->anchorOpName.setValue(opName.str());
210   }
211 
212   void runOnFunction() override {
213     auto funcOp = getFunction();
214     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
215       return;
216 
217     RewritePatternSet vectorizationPatterns(funcOp.getContext());
218     if (!anchorOpName.empty()) {
219       vectorizationPatterns.add<LinalgVectorizationPattern>(
220           anchorOpName, funcOp.getContext(), options, filter);
221     } else {
222       vectorizationPatterns.add<LinalgVectorizationPattern>(funcOp.getContext(),
223                                                             filter, options);
224     }
225     vector::populateVectorTransferPermutationMapLoweringPatterns(
226         vectorizationPatterns);
227     vector::populateVectorReductionToContractPatterns(vectorizationPatterns);
228     vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
229                               linalg::LinalgCopyVTWForwardingPattern>(
230         funcOp.getContext(), /*benefit=*/2);
231     (void)applyPatternsAndFoldGreedily(funcOp,
232                                        std::move(vectorizationPatterns));
233   }
234 
235   LinalgVectorizationOptions options;
236   LinalgTransformationFilter filter;
237 };
238 
239 /// Configurable pass to enable the application of other pattern-based linalg
240 /// passes.
241 struct LinalgStrategyEnablePass
242     : public LinalgStrategyEnablePassBase<LinalgStrategyEnablePass> {
243 
244   LinalgStrategyEnablePass(LinalgEnablingOptions opt,
245                            LinalgTransformationFilter filt)
246       : options(opt), filter(filt) {}
247 
248   void runOnFunction() override {
249     auto funcOp = getFunction();
250     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
251       return;
252 
253     MLIRContext *context = funcOp.getContext();
254     RewritePatternSet patterns =
255         linalg::getLinalgTilingCanonicalizationPatterns(context);
256     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
257     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
258       return signalPassFailure();
259 
260     if (options.licm) {
261       if (funcOp
262               ->walk([&](LoopLikeOpInterface loopLike) {
263                 if (failed(moveLoopInvariantCode(loopLike)))
264                   return WalkResult::interrupt();
265                 return WalkResult::advance();
266               })
267               .wasInterrupted())
268         return signalPassFailure();
269     }
270 
271     promoteSingleIterationLoops(funcOp);
272     if (options.hoistRedundantVectorTransfers)
273       hoistRedundantVectorTransfers(funcOp);
274 
275     if (options.hoistRedundantVectorTransfersOnTensor)
276       hoistRedundantVectorTransfersOnTensor(funcOp);
277   }
278 
279   LinalgEnablingOptions options;
280   LinalgTransformationFilter filter;
281 };
282 
283 /// Configurable pass to lower vector operations.
284 struct LinalgStrategyLowerVectorsPass
285     : public LinalgStrategyLowerVectorsPassBase<
286           LinalgStrategyLowerVectorsPass> {
287 
288   LinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
289                                  LinalgTransformationFilter filt)
290       : options(opt), filter(filt) {}
291 
292   void runOnFunction() override {
293     auto funcOp = getFunction();
294     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
295       return;
296 
297     MLIRContext *context = funcOp.getContext();
298     RewritePatternSet patterns(context);
299     vector::populateVectorToVectorCanonicalizationPatterns(patterns);
300     // In a progressive lowering of vectors, this would be the 1st step.
301     if (options.contractionLowering) {
302       patterns.add<ContractionOpToOuterProductOpLowering,
303                    ContractionOpToMatmulOpLowering, ContractionOpLowering>(
304           options.vectorTransformOptions, context);
305       vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
306     }
307     // In a progressive lowering of vectors, this would be the 2nd step.
308     if (options.multiReductionLowering) {
309       vector::populateVectorMultiReductionLoweringPatterns(
310           patterns,
311           options.vectorTransformOptions.vectorMultiReductionLowering);
312     }
313     // In a progressive lowering of vectors, this would be the 3rd step.
314     if (options.transferPartialRewrite) {
315       patterns.add<vector::VectorTransferFullPartialRewriter>(
316           context, options.vectorTransformOptions);
317     }
318     // In a progressive lowering of vectors, this would be the 4th step.
319     if (options.transferLowering) {
320       vector::populateVectorTransferLoweringPatterns(patterns,
321                                                      options.maxTransferRank);
322     }
323     // In a progressive lowering of vectors, this would be the 5th step.
324     if (options.transferToSCFConversion) {
325       populateVectorToSCFConversionPatterns(
326           patterns, options.vectorTransferToSCFOptions.setTargetRank(
327                         options.maxTransferRank));
328     }
329     // In a progressive lowering of vectors, this would be the 6th step.
330     if (options.shapeCastLowering) {
331       vector::populateVectorShapeCastLoweringPatterns(patterns);
332     }
333     // In a progressive lowering of vectors, this would be the 7th step.
334     if (options.transposeLowering) {
335       vector::populateVectorTransposeLoweringPatterns(
336           patterns, options.vectorTransformOptions);
337     }
338     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
339   }
340 
341   LinalgVectorLoweringOptions options;
342   LinalgTransformationFilter filter;
343 };
344 
345 /// Configurable pass to lower vector operations.
346 struct LinalgStrategyRemoveMarkersPass
347     : public LinalgStrategyRemoveMarkersPassBase<
348           LinalgStrategyRemoveMarkersPass> {
349 
350   void runOnFunction() override {
351     auto funcOp = getFunction();
352     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
353       return;
354     funcOp.walk([](LinalgOp op) {
355       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
356     });
357   }
358 };
359 } // namespace
360 
361 /// Create a LinalgStrategyTilePass.
362 std::unique_ptr<OperationPass<FuncOp>>
363 mlir::createLinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
364                                    LinalgTransformationFilter filter) {
365   return std::make_unique<LinalgStrategyTilePass>(opName, opt, filter);
366 }
367 
368 /// Create a LinalgStrategyPadPass.
369 std::unique_ptr<OperationPass<FuncOp>>
370 mlir::createLinalgStrategyPadPass(StringRef opName, LinalgPaddingOptions opt,
371                                   LinalgTransformationFilter filter) {
372   return std::make_unique<LinalgStrategyPadPass>(opName, opt, filter);
373 }
374 
375 /// Create a LinalgStrategyPromotePass.
376 std::unique_ptr<OperationPass<FuncOp>>
377 mlir::createLinalgStrategyPromotePass(StringRef opName,
378                                       LinalgPromotionOptions opt,
379                                       LinalgTransformationFilter filter) {
380   return std::make_unique<LinalgStrategyPromotePass>(opName, opt, filter);
381 }
382 
383 /// Create a LinalgStrategyGeneralizePass.
384 std::unique_ptr<OperationPass<FuncOp>>
385 mlir::createLinalgStrategyGeneralizePass(StringRef opName,
386                                          LinalgTransformationFilter filter) {
387   return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
388 }
389 
390 /// Create a LinalgStrategyInterchangePass.
391 std::unique_ptr<OperationPass<FuncOp>>
392 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
393                                           LinalgTransformationFilter filter) {
394   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
395                                                          filter);
396 }
397 
398 /// Create a LinalgStrategyVectorizePass.
399 std::unique_ptr<OperationPass<FuncOp>>
400 mlir::createLinalgStrategyVectorizePass(StringRef opName,
401                                         LinalgVectorizationOptions opt,
402                                         LinalgTransformationFilter filter) {
403   return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter);
404 }
405 
406 /// Create a LinalgStrategyEnablePass.
407 std::unique_ptr<OperationPass<FuncOp>>
408 mlir::createLinalgStrategyEnablePass(LinalgEnablingOptions opt,
409                                      LinalgTransformationFilter filter) {
410   return std::make_unique<LinalgStrategyEnablePass>(opt, filter);
411 }
412 
413 /// Create a LinalgStrategyLowerVectorsPass.
414 std::unique_ptr<OperationPass<FuncOp>>
415 mlir::createLinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
416                                            LinalgTransformationFilter filter) {
417   return std::make_unique<LinalgStrategyLowerVectorsPass>(opt, filter);
418 }
419 
420 /// Create a LinalgStrategyRemoveMarkersPass.
421 std::unique_ptr<OperationPass<FuncOp>>
422 mlir::createLinalgStrategyRemoveMarkersPass() {
423   return std::make_unique<LinalgStrategyRemoveMarkersPass>();
424 }
425