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 tiling and fusion.
40 struct LinalgStrategyTileAndFusePass
41     : public LinalgStrategyTileAndFusePassBase<LinalgStrategyTileAndFusePass> {
42 
43   LinalgStrategyTileAndFusePass() = default;
44 
45   LinalgStrategyTileAndFusePass(StringRef opName,
46                                 LinalgTilingAndFusionOptions opt,
47                                 LinalgTransformationFilter filt)
48       : options(opt), filter(filt) {
49     this->anchorOpName.setValue(opName.str());
50   }
51 
52   void runOnFunction() override {
53     auto funcOp = getFunction();
54     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
55       return;
56 
57     RewritePatternSet tilingAndFusionPattern(funcOp.getContext());
58     if (!anchorOpName.empty()) {
59       tilingAndFusionPattern.add<LinalgTileAndFuseTensorOpsPattern>(
60           anchorOpName, funcOp.getContext(), options, filter);
61     } else {
62       tilingAndFusionPattern.add<LinalgTileAndFuseTensorOpsPattern>(
63           funcOp.getContext(), options, filter);
64     }
65     // Search the root operation using bottom up traversal.
66     GreedyRewriteConfig grc;
67     grc.useTopDownTraversal = false;
68     (void)applyPatternsAndFoldGreedily(funcOp,
69                                        std::move(tilingAndFusionPattern), grc);
70   }
71 
72   LinalgTilingAndFusionOptions options;
73   LinalgTransformationFilter filter;
74 };
75 
76 /// Configurable pass to apply pattern-based linalg tiling.
77 struct LinalgStrategyTilePass
78     : public LinalgStrategyTilePassBase<LinalgStrategyTilePass> {
79 
80   LinalgStrategyTilePass() = default;
81 
82   LinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
83                          LinalgTransformationFilter filt)
84       : options(opt), filter(filt) {
85     this->anchorOpName.setValue(opName.str());
86   }
87 
88   void runOnFunction() override {
89     auto funcOp = getFunction();
90     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
91       return;
92 
93     RewritePatternSet tilingPattern(funcOp.getContext());
94     if (!anchorOpName.empty()) {
95       tilingPattern.add<LinalgGenericTilingPattern>(
96           anchorOpName, funcOp.getContext(), options, filter);
97     } else {
98       tilingPattern.add<LinalgGenericTilingPattern>(funcOp.getContext(), filter,
99                                                     options);
100     }
101     (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
102   }
103 
104   LinalgTilingOptions options;
105   LinalgTransformationFilter filter;
106 };
107 
108 /// Configurable pass to apply hoisting and padding.
109 struct LinalgStrategyPadPass
110     : public LinalgStrategyPadPassBase<LinalgStrategyPadPass> {
111 
112   LinalgStrategyPadPass() = default;
113 
114   LinalgStrategyPadPass(StringRef opName, LinalgPaddingOptions opt,
115                         LinalgTransformationFilter filt)
116       : options(opt), filter(filt) {
117     this->anchorOpName.setValue(opName.str());
118   }
119 
120   void runOnFunction() override {
121     auto funcOp = getFunction();
122     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
123       return;
124 
125     RewritePatternSet paddingPattern(funcOp.getContext());
126     if (!anchorOpName.empty()) {
127       paddingPattern.add<LinalgPaddingPattern>(
128           anchorOpName, funcOp.getContext(), options, filter);
129     } else {
130       paddingPattern.add<LinalgPaddingPattern>(funcOp.getContext(), options,
131                                                filter);
132     }
133     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(paddingPattern))))
134       signalPassFailure();
135   }
136 
137   LinalgPaddingOptions options;
138   LinalgTransformationFilter filter;
139 };
140 
141 /// Configurable pass to apply pattern-based linalg generalization.
142 struct LinalgStrategyGeneralizePass
143     : public LinalgStrategyGeneralizePassBase<LinalgStrategyGeneralizePass> {
144 
145   LinalgStrategyGeneralizePass() = default;
146 
147   LinalgStrategyGeneralizePass(StringRef opName,
148                                LinalgTransformationFilter filter)
149       : filter(filter) {
150     this->anchorOpName.setValue(opName.str());
151   }
152 
153   void runOnFunction() override {
154     auto funcOp = getFunction();
155     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
156       return;
157 
158     RewritePatternSet generalizationPattern(funcOp.getContext());
159     if (!anchorOpName.empty()) {
160       generalizationPattern.add<LinalgGeneralizationPattern>(
161           anchorOpName, funcOp.getContext(), filter);
162     } else {
163       generalizationPattern.add<LinalgGeneralizationPattern>(
164           funcOp.getContext(), filter);
165     }
166     if (failed(applyPatternsAndFoldGreedily(funcOp,
167                                             std::move(generalizationPattern))))
168       signalPassFailure();
169   }
170 
171   LinalgTransformationFilter filter;
172 };
173 
174 /// Configurable pass to apply lowering of coarser-grained named linalg ops into
175 /// finer-grained named versions.
176 struct LinalgStrategyDecomposePass
177     : public LinalgStrategyDecomposePassBase<LinalgStrategyDecomposePass> {
178 
179   LinalgStrategyDecomposePass() = default;
180 
181   void runOnFunction() override {
182     auto funcOp = getFunction();
183     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
184       return;
185     RewritePatternSet decompositionPattern(funcOp.getContext());
186     populateDecomposeConvolutionPatterns(decompositionPattern);
187     if (failed(applyPatternsAndFoldGreedily(funcOp,
188                                             std::move(decompositionPattern))))
189       signalPassFailure();
190   }
191 };
192 
193 /// Configurable pass to apply pattern-based linalg generalization.
194 struct LinalgStrategyInterchangePass
195     : public LinalgStrategyInterchangePassBase<LinalgStrategyInterchangePass> {
196 
197   LinalgStrategyInterchangePass() = default;
198 
199   LinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
200                                 LinalgTransformationFilter filter)
201       : iteratorInterchange(iteratorInterchange.begin(),
202                             iteratorInterchange.end()),
203         filter(filter) {}
204 
205   void runOnFunction() override {
206     auto funcOp = getFunction();
207     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
208       return;
209 
210     SmallVector<unsigned> interchangeVector(iteratorInterchange.begin(),
211                                             iteratorInterchange.end());
212     RewritePatternSet interchangePattern(funcOp.getContext());
213     interchangePattern.add<GenericOpInterchangePattern>(
214         funcOp.getContext(), interchangeVector, filter);
215     if (failed(applyPatternsAndFoldGreedily(funcOp,
216                                             std::move(interchangePattern))))
217       signalPassFailure();
218   }
219 
220   SmallVector<int64_t> iteratorInterchange;
221   LinalgTransformationFilter filter;
222 };
223 
224 /// Configurable pass to apply pattern-based linalg promotion.
225 struct LinalgStrategyPromotePass
226     : public LinalgStrategyPromotePassBase<LinalgStrategyPromotePass> {
227 
228   LinalgStrategyPromotePass() = default;
229 
230   LinalgStrategyPromotePass(StringRef opName, LinalgPromotionOptions opt,
231                             LinalgTransformationFilter filt)
232       : options(opt), filter(filt) {
233     this->anchorOpName.setValue(opName.str());
234   }
235 
236   void runOnFunction() override {
237     auto funcOp = getFunction();
238     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
239       return;
240 
241     RewritePatternSet promotionPattern(funcOp.getContext());
242     if (!anchorOpName.empty()) {
243       promotionPattern.add<LinalgBasePromotionPattern>(
244           anchorOpName, funcOp.getContext(), options, filter);
245     } else {
246       promotionPattern.add<LinalgBasePromotionPattern>(funcOp.getContext(),
247                                                        filter, options);
248     }
249     (void)applyPatternsAndFoldGreedily(funcOp, std::move(promotionPattern));
250   }
251 
252   LinalgPromotionOptions options;
253   LinalgTransformationFilter filter;
254 };
255 
256 /// Configurable pass to apply pattern-based linalg vectorization.
257 struct LinalgStrategyVectorizePass
258     : public LinalgStrategyVectorizePassBase<LinalgStrategyVectorizePass> {
259 
260   LinalgStrategyVectorizePass() = default;
261 
262   LinalgStrategyVectorizePass(StringRef opName, LinalgVectorizationOptions opt,
263                               LinalgTransformationFilter filt)
264       : options(opt), filter(filt) {
265     this->anchorOpName.setValue(opName.str());
266   }
267 
268   void runOnFunction() override {
269     auto funcOp = getFunction();
270     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
271       return;
272 
273     RewritePatternSet vectorizationPatterns(funcOp.getContext());
274     if (!anchorOpName.empty()) {
275       vectorizationPatterns.add<LinalgVectorizationPattern>(
276           anchorOpName, funcOp.getContext(), options, filter);
277     } else {
278       vectorizationPatterns.add<LinalgVectorizationPattern>(funcOp.getContext(),
279                                                             filter, options);
280     }
281     vector::populateVectorTransferPermutationMapLoweringPatterns(
282         vectorizationPatterns);
283     vector::populateVectorReductionToContractPatterns(vectorizationPatterns);
284     vectorizationPatterns.add<linalg::LinalgCopyVTRForwardingPattern,
285                               linalg::LinalgCopyVTWForwardingPattern>(
286         funcOp.getContext(), /*benefit=*/2);
287     (void)applyPatternsAndFoldGreedily(funcOp,
288                                        std::move(vectorizationPatterns));
289   }
290 
291   LinalgVectorizationOptions options;
292   LinalgTransformationFilter filter;
293 };
294 
295 /// Configurable pass to enable the application of other pattern-based linalg
296 /// passes.
297 struct LinalgStrategyEnablePass
298     : public LinalgStrategyEnablePassBase<LinalgStrategyEnablePass> {
299 
300   LinalgStrategyEnablePass(LinalgEnablingOptions opt,
301                            LinalgTransformationFilter filt)
302       : options(opt), filter(filt) {}
303 
304   void runOnFunction() override {
305     auto funcOp = getFunction();
306     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
307       return;
308 
309     MLIRContext *context = funcOp.getContext();
310     RewritePatternSet patterns =
311         linalg::getLinalgTilingCanonicalizationPatterns(context);
312     scf::populateSCFForLoopCanonicalizationPatterns(patterns);
313     if (failed(applyPatternsAndFoldGreedily(funcOp, std::move(patterns))))
314       return signalPassFailure();
315 
316     if (options.licm) {
317       if (funcOp
318               ->walk([&](LoopLikeOpInterface loopLike) {
319                 if (failed(moveLoopInvariantCode(loopLike)))
320                   return WalkResult::interrupt();
321                 return WalkResult::advance();
322               })
323               .wasInterrupted())
324         return signalPassFailure();
325     }
326 
327     promoteSingleIterationLoops(funcOp);
328     if (options.hoistRedundantVectorTransfers)
329       hoistRedundantVectorTransfers(funcOp);
330 
331     if (options.hoistRedundantVectorTransfersOnTensor)
332       hoistRedundantVectorTransfersOnTensor(funcOp);
333   }
334 
335   LinalgEnablingOptions options;
336   LinalgTransformationFilter filter;
337 };
338 
339 /// Configurable pass to lower vector operations.
340 struct LinalgStrategyLowerVectorsPass
341     : public LinalgStrategyLowerVectorsPassBase<
342           LinalgStrategyLowerVectorsPass> {
343 
344   LinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
345                                  LinalgTransformationFilter filt)
346       : options(opt), filter(filt) {}
347 
348   void runOnFunction() override {
349     auto funcOp = getFunction();
350     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
351       return;
352 
353     MLIRContext *context = funcOp.getContext();
354     RewritePatternSet patterns(context);
355     vector::populateVectorToVectorCanonicalizationPatterns(patterns);
356     // In a progressive lowering of vectors, this would be the 1st step.
357     if (options.contractionLowering) {
358       patterns.add<ContractionOpToOuterProductOpLowering,
359                    ContractionOpToMatmulOpLowering, ContractionOpLowering>(
360           options.vectorTransformOptions, context);
361       vector::populateVectorTransferPermutationMapLoweringPatterns(patterns);
362     }
363     // In a progressive lowering of vectors, this would be the 2nd step.
364     if (options.multiReductionLowering) {
365       vector::populateVectorMultiReductionLoweringPatterns(
366           patterns,
367           options.vectorTransformOptions.vectorMultiReductionLowering);
368     }
369     // In a progressive lowering of vectors, this would be the 3rd step.
370     if (options.transferPartialRewrite) {
371       patterns.add<vector::VectorTransferFullPartialRewriter>(
372           context, options.vectorTransformOptions);
373     }
374     // In a progressive lowering of vectors, this would be the 4th step.
375     if (options.transferLowering) {
376       vector::populateVectorTransferLoweringPatterns(patterns,
377                                                      options.maxTransferRank);
378     }
379     // In a progressive lowering of vectors, this would be the 5th step.
380     if (options.transferToSCFConversion) {
381       populateVectorToSCFConversionPatterns(
382           patterns, options.vectorTransferToSCFOptions.setTargetRank(
383                         options.maxTransferRank));
384     }
385     // In a progressive lowering of vectors, this would be the 6th step.
386     if (options.shapeCastLowering) {
387       vector::populateVectorShapeCastLoweringPatterns(patterns);
388     }
389     // In a progressive lowering of vectors, this would be the 7th step.
390     if (options.transposeLowering) {
391       vector::populateVectorTransposeLoweringPatterns(
392           patterns, options.vectorTransformOptions);
393       if (options.avx2Lowering)
394         x86vector::avx2::populateSpecializedTransposeLoweringPatterns(
395             patterns, options.avx2LoweringOptions, /*benefit=*/10);
396     }
397     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
398   }
399 
400   LinalgVectorLoweringOptions options;
401   LinalgTransformationFilter filter;
402 };
403 
404 /// Configurable pass to lower vector operations.
405 struct LinalgStrategyRemoveMarkersPass
406     : public LinalgStrategyRemoveMarkersPassBase<
407           LinalgStrategyRemoveMarkersPass> {
408 
409   void runOnFunction() override {
410     auto funcOp = getFunction();
411     if (!anchorFuncName.empty() && funcOp.getName() != anchorFuncName)
412       return;
413     funcOp.walk([](LinalgOp op) {
414       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
415     });
416   }
417 };
418 } // namespace
419 
420 /// Create a LinalgStrategyTileAndFusePass.
421 std::unique_ptr<OperationPass<FuncOp>>
422 mlir::createLinalgStrategyTileAndFusePass(StringRef opName,
423                                           LinalgTilingAndFusionOptions options,
424                                           LinalgTransformationFilter filter) {
425   return std::make_unique<LinalgStrategyTileAndFusePass>(opName, options,
426                                                          filter);
427 }
428 
429 /// Create a LinalgStrategyTilePass.
430 std::unique_ptr<OperationPass<FuncOp>>
431 mlir::createLinalgStrategyTilePass(StringRef opName, LinalgTilingOptions opt,
432                                    LinalgTransformationFilter filter) {
433   return std::make_unique<LinalgStrategyTilePass>(opName, opt, filter);
434 }
435 
436 /// Create a LinalgStrategyPadPass.
437 std::unique_ptr<OperationPass<FuncOp>>
438 mlir::createLinalgStrategyPadPass(StringRef opName, LinalgPaddingOptions opt,
439                                   LinalgTransformationFilter filter) {
440   return std::make_unique<LinalgStrategyPadPass>(opName, opt, filter);
441 }
442 
443 /// Create a LinalgStrategyPromotePass.
444 std::unique_ptr<OperationPass<FuncOp>>
445 mlir::createLinalgStrategyPromotePass(StringRef opName,
446                                       LinalgPromotionOptions opt,
447                                       LinalgTransformationFilter filter) {
448   return std::make_unique<LinalgStrategyPromotePass>(opName, opt, filter);
449 }
450 
451 /// Create a LinalgStrategyGeneralizePass.
452 std::unique_ptr<OperationPass<FuncOp>>
453 mlir::createLinalgStrategyGeneralizePass(StringRef opName,
454                                          LinalgTransformationFilter filter) {
455   return std::make_unique<LinalgStrategyGeneralizePass>(opName, filter);
456 }
457 /// Create a LinalgStrategyDecomposePass.
458 // TODO: atm this is applied to all supported ops. If/when we need finer control
459 // this should be exposed with an opName + filter and a proper pattern.
460 std::unique_ptr<OperationPass<FuncOp>>
461 mlir::createLinalgStrategyDecomposePass() {
462   return std::make_unique<LinalgStrategyDecomposePass>();
463 }
464 
465 /// Create a LinalgStrategyInterchangePass.
466 std::unique_ptr<OperationPass<FuncOp>>
467 mlir::createLinalgStrategyInterchangePass(ArrayRef<int64_t> iteratorInterchange,
468                                           LinalgTransformationFilter filter) {
469   return std::make_unique<LinalgStrategyInterchangePass>(iteratorInterchange,
470                                                          filter);
471 }
472 
473 /// Create a LinalgStrategyVectorizePass.
474 std::unique_ptr<OperationPass<FuncOp>>
475 mlir::createLinalgStrategyVectorizePass(StringRef opName,
476                                         LinalgVectorizationOptions opt,
477                                         LinalgTransformationFilter filter) {
478   return std::make_unique<LinalgStrategyVectorizePass>(opName, opt, filter);
479 }
480 
481 /// Create a LinalgStrategyEnablePass.
482 std::unique_ptr<OperationPass<FuncOp>>
483 mlir::createLinalgStrategyEnablePass(LinalgEnablingOptions opt,
484                                      LinalgTransformationFilter filter) {
485   return std::make_unique<LinalgStrategyEnablePass>(opt, filter);
486 }
487 
488 /// Create a LinalgStrategyLowerVectorsPass.
489 std::unique_ptr<OperationPass<FuncOp>>
490 mlir::createLinalgStrategyLowerVectorsPass(LinalgVectorLoweringOptions opt,
491                                            LinalgTransformationFilter filter) {
492   return std::make_unique<LinalgStrategyLowerVectorsPass>(opt, filter);
493 }
494 
495 /// Create a LinalgStrategyRemoveMarkersPass.
496 std::unique_ptr<OperationPass<FuncOp>>
497 mlir::createLinalgStrategyRemoveMarkersPass() {
498   return std::make_unique<LinalgStrategyRemoveMarkersPass>();
499 }
500