1 //===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
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 logic for testing Linalg transformations.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/Affine/IR/AffineOps.h"
14 #include "mlir/Dialect/GPU/GPUDialect.h"
15 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
16 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
17 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
18 #include "mlir/Dialect/Linalg/Utils/Utils.h"
19 #include "mlir/Dialect/StandardOps/IR/Ops.h"
20 #include "mlir/Dialect/Vector/VectorOps.h"
21 #include "mlir/Pass/Pass.h"
22 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
23 
24 #include "llvm/ADT/SetVector.h"
25 
26 using namespace mlir;
27 using namespace mlir::linalg;
28 
29 namespace {
30 struct TestLinalgTransforms
31     : public PassWrapper<TestLinalgTransforms, FunctionPass> {
32   TestLinalgTransforms() = default;
33   TestLinalgTransforms(const TestLinalgTransforms &pass) {}
34 
35   void getDependentDialects(DialectRegistry &registry) const override {
36     // clang-format off
37     registry.insert<AffineDialect,
38                     memref::MemRefDialect,
39                     scf::SCFDialect,
40                     StandardOpsDialect,
41                     vector::VectorDialect,
42                     gpu::GPUDialect>();
43     // clang-format on
44   }
45   StringRef getArgument() const final {
46     return "test-linalg-transform-patterns";
47   }
48   StringRef getDescription() const final {
49     return "Test Linalg transformation patterns by applying them greedily.";
50   }
51 
52   void runOnFunction() override;
53 
54   Option<bool> testPatterns{*this, "test-patterns",
55                             llvm::cl::desc("Test a mixed set of patterns"),
56                             llvm::cl::init(false)};
57   Option<bool> testMatmulToVectorPatterns1dTiling{
58       *this, "test-matmul-to-vector-patterns-tile-1d",
59       llvm::cl::desc(
60           "Test a fused pass that applies patterns from matmul to vectors via "
61           "1-d tiling"),
62       llvm::cl::init(false)};
63   Option<bool> testMatmulToVectorPatterns2dTiling{
64       *this, "test-matmul-to-vector-patterns-tile-2d",
65       llvm::cl::desc(
66           "Test a fused pass that applies patterns from matmul to vectors via "
67           "2-d tiling"),
68       llvm::cl::init(false)};
69   Option<bool> testPromotionOptions{*this, "test-linalg-promotion-options",
70                                     llvm::cl::desc("Test promotion options"),
71                                     llvm::cl::init(false)};
72   Option<bool> testTileAndDistributionOptions{
73       *this, "test-tile-and-distribute-options",
74       llvm::cl::desc("Test tile and distribute options"),
75       llvm::cl::init(false)};
76   Option<bool> testVectorTransferForwardingPatterns{
77       *this, "test-vector-transfer-forwarding-patterns",
78       llvm::cl::desc(
79           "Test a fused pass that forwards linalg.copy to vector.transfer"),
80       llvm::cl::init(false)};
81   Option<bool> testGenericToVectorPattern{
82       *this, "test-linalg-to-vector-patterns",
83       llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
84                      "in vector.contract form"),
85       llvm::cl::init(false)};
86   Option<bool> testAffineMinSCFCanonicalizationPatterns{
87       *this, "test-affine-min-scf-canonicalization-patterns",
88       llvm::cl::desc("Test affine-min + scf canonicalization patterns."),
89       llvm::cl::init(false)};
90   Option<bool> testTileAndPadPattern{
91       *this, "test-tile-and-pad-pattern",
92       llvm::cl::desc("Test tile and pad pattern"), llvm::cl::init(false)};
93   Option<int> testHoistPadding{*this, "test-hoist-padding",
94                                llvm::cl::desc("Test hoist padding"),
95                                llvm::cl::init(0)};
96   Option<bool> testTransformPadTensor{
97       *this, "test-transform-pad-tensor",
98       llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
99       llvm::cl::init(false)};
100   ListOption<int64_t> tileSizesForPadding{
101       *this, "tile-sizes-for-padding",
102       llvm::cl::desc("Linalg tile sizes when tile+pad"), llvm::cl::ZeroOrMore,
103       llvm::cl::MiscFlags::CommaSeparated};
104   ListOption<unsigned> testInterchangePattern{
105       *this, "test-interchange-pattern", llvm::cl::MiscFlags::CommaSeparated,
106       llvm::cl::desc("Test the interchange pattern.")};
107 };
108 } // end anonymous namespace
109 
110 static void applyPatterns(FuncOp funcOp) {
111   MLIRContext *ctx = funcOp.getContext();
112   RewritePatternSet patterns(ctx);
113 
114   //===--------------------------------------------------------------------===//
115   // Linalg tiling patterns.
116   //===--------------------------------------------------------------------===//
117   patterns.add<LinalgTilingPattern<MatmulOp>>(
118       ctx, LinalgTilingOptions().setTileSizes({2000, 3000, 4000}),
119       LinalgTransformationFilter(Identifier::get("MEM", ctx),
120                                  Identifier::get("L3", ctx)));
121   patterns.add<LinalgTilingPattern<MatmulOp>>(
122       ctx, LinalgTilingOptions().setTileSizes({200, 300, 400}),
123       LinalgTransformationFilter(Identifier::get("L3", ctx),
124                                  Identifier::get("L2", ctx)));
125   patterns.add<LinalgTilingPattern<MatmulOp>>(
126       ctx, LinalgTilingOptions().setTileSizes({20, 30, 40}),
127       LinalgTransformationFilter(Identifier::get("L2", ctx),
128                                  Identifier::get("L1", ctx)));
129   patterns.add<LinalgTilingPattern<MatmulOp>>(
130       ctx, LinalgTilingOptions().setTileSizes({2, 3, 4}),
131       LinalgTransformationFilter(Identifier::get("L1", ctx),
132                                  Identifier::get("REG", ctx)));
133 
134   patterns.add<LinalgTilingPattern<MatvecOp>>(
135       ctx,
136       LinalgTilingOptions().setTileSizes({5, 6}).setLoopType(
137           LinalgTilingLoopType::ParallelLoops),
138       LinalgTransformationFilter(ArrayRef<Identifier>{},
139                                  Identifier::get("L1", ctx)));
140 
141   patterns.add<LinalgTilingPattern<DotOp>>(
142       ctx, LinalgTilingOptions().setTileSizes(8000),
143       LinalgTransformationFilter(
144           ArrayRef<Identifier>{Identifier::get("MEM", ctx),
145                                Identifier::get("L3", ctx),
146                                Identifier::get("L2", ctx)},
147           Identifier::get("REG", ctx)));
148 
149   //===--------------------------------------------------------------------===//
150   // Linalg tiling and permutation patterns.
151   //===--------------------------------------------------------------------===//
152   patterns.add<LinalgTilingPattern<MatmulOp>>(
153       ctx,
154       LinalgTilingOptions()
155           .setTileSizes({2000, 3000, 4000})
156           .setInterchange({1, 2, 0}),
157       LinalgTransformationFilter(Identifier::get("__with_perm__", ctx),
158                                  Identifier::get("L2__with_perm__", ctx)));
159   patterns.add<LinalgTilingPattern<MatmulOp>>(
160       ctx,
161       LinalgTilingOptions()
162           .setTileSizes({200, 300, 400})
163           .setInterchange({1, 0, 2}),
164       LinalgTransformationFilter(Identifier::get("L2__with_perm__", ctx),
165                                  Identifier::get("L1__with_perm__", ctx)));
166   patterns.add<LinalgTilingPattern<MatmulOp>>(
167       ctx, LinalgTilingOptions().setTileSizes({20, 30, 40}),
168       LinalgTransformationFilter(Identifier::get("L1__with_perm__", ctx),
169                                  Identifier::get("REG__with_perm__", ctx)));
170 
171   patterns.add<LinalgTilingPattern<MatvecOp>>(
172       ctx, LinalgTilingOptions().setTileSizes({5, 6}).setInterchange({1, 0}),
173       LinalgTransformationFilter(Identifier::get("__with_perm__", ctx),
174                                  Identifier::get("L1__with_perm__", ctx)));
175 
176   patterns.add<LinalgTilingPattern<MatmulOp>>(
177       ctx,
178       LinalgTilingOptions()
179           .setTileSizes({16, 8, 4})
180           .setInterchange({1, 2, 0})
181           .setLoopType(LinalgTilingLoopType::ParallelLoops),
182       LinalgTransformationFilter(
183           Identifier::get("par__with_perm__", ctx),
184           Identifier::get("after_par__with_perm__", ctx)));
185 
186   //===--------------------------------------------------------------------===//
187   // Linalg to loops patterns.
188   //===--------------------------------------------------------------------===//
189   patterns.add<LinalgLoweringPattern<DotOp>>(
190       ctx,
191       /*loweringType=*/LinalgLoweringType::Loops,
192       LinalgTransformationFilter(Identifier::get("REG", ctx)));
193 
194   //===--------------------------------------------------------------------===//
195   // Linalg distribution patterns.
196   //===--------------------------------------------------------------------===//
197   LinalgLoopDistributionOptions distributionOptions;
198 
199   //===--------------------------------------------------------------------===//
200   // Linalg to vector contraction patterns.
201   //===--------------------------------------------------------------------===//
202   patterns.add<LinalgVectorizationPattern>(
203       ctx, LinalgTransformationFilter(Identifier::get("VECTORIZE", ctx))
204                .addOpFilter<MatmulOp, FillOp, CopyOp, GenericOp>());
205 
206   //===--------------------------------------------------------------------===//
207   // Linalg generic interchange pattern.
208   //===--------------------------------------------------------------------===//
209   patterns.add<GenericOpInterchangePattern>(
210       ctx,
211       /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
212       LinalgTransformationFilter(ArrayRef<Identifier>{},
213                                  Identifier::get("PERMUTED", ctx)));
214 
215   //===--------------------------------------------------------------------===//
216   // Linalg subview operands promotion.
217   //===--------------------------------------------------------------------===//
218   patterns.add<LinalgPromotionPattern<MatmulOp>>(
219       ctx, LinalgPromotionOptions().setUseFullTileBuffersByDefault(true),
220       LinalgTransformationFilter(Identifier::get("_promote_views_", ctx),
221                                  Identifier::get("_views_promoted_", ctx)));
222   patterns.add<LinalgPromotionPattern<MatmulOp>>(
223       ctx,
224       LinalgPromotionOptions()
225           .setOperandsToPromote({0})
226           .setUseFullTileBuffersByDefault(true),
227       LinalgTransformationFilter(
228           Identifier::get("_promote_first_view_", ctx),
229           Identifier::get("_first_view_promoted_", ctx)));
230   patterns.add<LinalgPromotionPattern<FillOp>>(
231       ctx,
232       LinalgPromotionOptions()
233           .setOperandsToPromote({0})
234           .setUseFullTileBuffers({true})
235           .setAlignment(32),
236       LinalgTransformationFilter(
237           Identifier::get("_promote_views_aligned_", ctx),
238           Identifier::get("_views_aligned_promoted_", ctx)));
239 
240   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
241 
242   // Drop the marker.
243   funcOp.walk([](LinalgOp op) {
244     op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
245   });
246 }
247 
248 static void fillL1TilingAndMatmulToVectorPatterns(
249     FuncOp funcOp, StringRef startMarker,
250     SmallVectorImpl<RewritePatternSet> &patternsVector) {
251   MLIRContext *ctx = funcOp.getContext();
252   patternsVector.emplace_back(
253       ctx, std::make_unique<LinalgTilingPattern<MatmulOp>>(
254                ctx,
255                LinalgTilingOptions()
256                    .setTileSizes({8, 12, 16})
257                    .setInterchange({1, 0, 2}),
258                LinalgTransformationFilter(Identifier::get(startMarker, ctx),
259                                           Identifier::get("L1", ctx))));
260 
261   patternsVector.emplace_back(
262       ctx,
263       std::make_unique<LinalgPromotionPattern<MatmulOp>>(
264           ctx, LinalgPromotionOptions().setUseFullTileBuffersByDefault(true),
265           LinalgTransformationFilter(Identifier::get("L1", ctx),
266                                      Identifier::get("VEC", ctx))));
267 
268   patternsVector.emplace_back(
269       ctx, std::make_unique<LinalgVectorizationPattern>(
270                MatmulOp::getOperationName(), ctx, LinalgVectorizationOptions(),
271                LinalgTransformationFilter(Identifier::get("VEC", ctx))));
272   patternsVector.back().add<LinalgVectorizationPattern>(
273       ctx, LinalgTransformationFilter().addFilter(
274                [](Operation *op) { return success(isa<FillOp, CopyOp>(op)); }));
275 }
276 
277 //===----------------------------------------------------------------------===//
278 // Test promotion callbacks
279 //===----------------------------------------------------------------------===//
280 
281 // Allocation call back
282 static Optional<Value> allocCallBackFn(OpBuilder &b, memref::SubViewOp subView,
283                                        ArrayRef<Value> boundingSubViewSize,
284                                        DataLayout &layout) {
285   SmallVector<int64_t, 4> shape(boundingSubViewSize.size(), -1);
286   return b
287       .create<memref::AllocOp>(
288           subView.getLoc(),
289           MemRefType::get(shape, subView.getType().getElementType(),
290                           /*affineMapComposition =*/{}, 3),
291           boundingSubViewSize)
292       .getResult();
293 }
294 
295 // Deallocation callback
296 static LogicalResult deallocCallBackFn(OpBuilder &b, Value buffer) {
297   b.create<memref::DeallocOp>(buffer.getLoc(), buffer);
298   return success();
299 }
300 
301 // Copy in call back
302 static LogicalResult copyCallBackFn(OpBuilder &b, Value src, Value dst,
303                                     bool isOutput) {
304   auto floatType = src.getType().cast<MemRefType>().getElementType();
305   if (!floatType.isa<FloatType>())
306     return failure();
307   if (!isOutput)
308     b.create<FillOp>(
309         src.getLoc(), dst,
310         b.create<ConstantOp>(src.getLoc(), FloatAttr::get(floatType, 42.0)));
311   b.create<CopyOp>(src.getLoc(), src, dst);
312   return success();
313 }
314 
315 static void fillPromotionCallBackPatterns(MLIRContext *ctx,
316                                           RewritePatternSet &patterns) {
317   patterns.add<LinalgTilingPattern<MatmulOp>>(
318       ctx, LinalgTilingOptions().setTileSizes({16, 16, 16}),
319       LinalgTransformationFilter(Identifier::get("START", ctx),
320                                  Identifier::get("PROMOTE", ctx)));
321   patterns.add<LinalgPromotionPattern<MatmulOp>>(
322       ctx,
323       LinalgPromotionOptions()
324           .setOperandsToPromote({0, 2})
325           .setUseFullTileBuffers({false, false})
326           .setAllocationDeallocationFns(allocCallBackFn, deallocCallBackFn)
327           .setCopyInOutFns(
328               [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
329                 return copyCallBackFn(b, src, dst, false);
330               },
331               [](OpBuilder &b, Value src, Value dst) -> LogicalResult {
332                 return copyCallBackFn(b, src, dst, true);
333               }),
334       LinalgTransformationFilter(Identifier::get("PROMOTE", ctx)));
335 }
336 
337 template <typename IdOp, typename NProcsOp>
338 static SmallVector<ProcInfo, 2>
339 getGpuProcIds(OpBuilder &b, Location loc, ArrayRef<Range> parallelLoopRanges) {
340   size_t count = std::min<size_t>(3, parallelLoopRanges.size());
341   SmallVector<ProcInfo, 2> procInfo(count);
342   const char *xyz[] = {"x", "y", "z"};
343   Type indexType = b.getIndexType();
344   for (unsigned i = 0; i < count; ++i) {
345     procInfo[count - 1 - i] = {
346         b.create<IdOp>(loc, indexType, b.getStringAttr(xyz[i])),
347         b.create<NProcsOp>(loc, indexType, b.getStringAttr(xyz[i]))};
348   }
349   return procInfo;
350 }
351 
352 static void fillTileAndDistributePatterns(MLIRContext *context,
353                                           RewritePatternSet &patterns) {
354   {
355     LinalgLoopDistributionOptions cyclicNprocsEqNiters;
356     cyclicNprocsEqNiters.distributionMethod.resize(
357         2, DistributionMethod::CyclicNumProcsEqNumIters);
358     cyclicNprocsEqNiters.procInfo =
359         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
360     patterns.add<LinalgTilingPattern<MatmulOp>>(
361         context,
362         LinalgTilingOptions()
363             .setTileSizes({8, 8, 4})
364             .setLoopType(LinalgTilingLoopType::ParallelLoops)
365             .setDistributionOptions(cyclicNprocsEqNiters),
366         LinalgTransformationFilter(
367             Identifier::get("distribute1", context),
368             Identifier::get("after_distribute1", context)));
369   }
370 
371   {
372     LinalgLoopDistributionOptions cyclicNprocsGeNiters;
373     cyclicNprocsGeNiters.distributionMethod.resize(
374         2, DistributionMethod::CyclicNumProcsGeNumIters);
375     cyclicNprocsGeNiters.procInfo =
376         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
377     patterns.add<LinalgTilingPattern<MatmulOp>>(
378         context,
379         LinalgTilingOptions()
380             .setTileSizes({8, 8, 4})
381             .setLoopType(LinalgTilingLoopType::ParallelLoops)
382             .setDistributionOptions(cyclicNprocsGeNiters),
383         LinalgTransformationFilter(
384             Identifier::get("distribute2", context),
385             Identifier::get("after_distribute2", context)));
386   }
387 
388   {
389     LinalgLoopDistributionOptions cyclicNprocsDefault;
390     cyclicNprocsDefault.distributionMethod.resize(2,
391                                                   DistributionMethod::Cyclic);
392     cyclicNprocsDefault.procInfo =
393         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
394     patterns.add<LinalgTilingPattern<MatmulOp>>(
395         context,
396         LinalgTilingOptions()
397             .setTileSizes({8, 8, 4})
398             .setLoopType(LinalgTilingLoopType::ParallelLoops)
399             .setDistributionOptions(cyclicNprocsDefault),
400         LinalgTransformationFilter(
401             Identifier::get("distribute3", context),
402             Identifier::get("after_distribute3", context)));
403   }
404 
405   {
406     LinalgLoopDistributionOptions cyclicNprocsMixed1;
407     cyclicNprocsMixed1.distributionMethod = {
408         DistributionMethod::CyclicNumProcsEqNumIters,
409         DistributionMethod::CyclicNumProcsGeNumIters};
410     cyclicNprocsMixed1.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
411     patterns.add<LinalgTilingPattern<MatmulOp>>(
412         context,
413         LinalgTilingOptions()
414             .setTileSizes({8, 8, 4})
415             .setLoopType(LinalgTilingLoopType::ParallelLoops)
416             .setDistributionOptions(cyclicNprocsMixed1),
417         LinalgTransformationFilter(
418             Identifier::get("distribute4", context),
419             Identifier::get("after_distribute4", context)));
420   }
421 
422   {
423     LinalgLoopDistributionOptions cyclicNprocsMixed2;
424     cyclicNprocsMixed2.distributionMethod = {
425         DistributionMethod::CyclicNumProcsGeNumIters,
426         DistributionMethod::Cyclic};
427     cyclicNprocsMixed2.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
428     patterns.add<LinalgTilingPattern<MatmulOp>>(
429         context,
430         LinalgTilingOptions()
431             .setTileSizes({8, 8, 4})
432             .setLoopType(LinalgTilingLoopType::ParallelLoops)
433             .setDistributionOptions(cyclicNprocsMixed2),
434         LinalgTransformationFilter(
435             Identifier::get("distribute5", context),
436             Identifier::get("after_distribute5", context)));
437   }
438 
439   {
440     LinalgLoopDistributionOptions cyclicNprocsMixed3;
441     cyclicNprocsMixed3.distributionMethod = {
442         DistributionMethod::Cyclic,
443         DistributionMethod::CyclicNumProcsEqNumIters};
444     cyclicNprocsMixed3.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
445 
446     patterns.add<LinalgTilingPattern<MatmulOp>>(
447         context,
448         LinalgTilingOptions()
449             .setTileSizes({8, 8, 4})
450             .setLoopType(LinalgTilingLoopType::ParallelLoops)
451             .setDistributionOptions(cyclicNprocsMixed3),
452         LinalgTransformationFilter(
453             Identifier::get("distribute6", context),
454             Identifier::get("after_distribute6", context)));
455   }
456 
457   {
458     LinalgLoopDistributionOptions cyclicNprocsEqNiters;
459     cyclicNprocsEqNiters.distributionMethod.resize(2,
460                                                    DistributionMethod::Cyclic);
461     cyclicNprocsEqNiters.procInfo =
462         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
463     patterns.add<LinalgTilingPattern<MatmulOp>>(
464         context,
465         LinalgTilingOptions()
466             .setTileSizes({8, 8, 4})
467             .setLoopType(LinalgTilingLoopType::Loops)
468             .setDistributionOptions(cyclicNprocsEqNiters),
469         LinalgTransformationFilter(
470             Identifier::get("tensors_distribute1", context),
471             Identifier::get("tensors_after_distribute1", context)));
472   }
473 }
474 
475 static void
476 applyMatmulToVectorPatterns(FuncOp funcOp,
477                             bool testMatmulToVectorPatterns1dTiling,
478                             bool testMatmulToVectorPatterns2dTiling) {
479   MLIRContext *ctx = funcOp.getContext();
480   SmallVector<RewritePatternSet, 4> stage1Patterns;
481   if (testMatmulToVectorPatterns1dTiling) {
482     fillL1TilingAndMatmulToVectorPatterns(funcOp, Identifier::get("START", ctx),
483                                           stage1Patterns);
484   } else if (testMatmulToVectorPatterns2dTiling) {
485     stage1Patterns.emplace_back(
486         ctx, std::make_unique<LinalgTilingPattern<MatmulOp>>(
487                  ctx,
488                  LinalgTilingOptions()
489                      .setTileSizes({768, 264, 768})
490                      .setInterchange({1, 2, 0}),
491                  LinalgTransformationFilter(Identifier::get("START", ctx),
492                                             Identifier::get("L2", ctx))));
493     fillL1TilingAndMatmulToVectorPatterns(funcOp, Identifier::get("L2", ctx),
494                                           stage1Patterns);
495   }
496   SmallVector<FrozenRewritePatternSet, 4> frozenStage1Patterns;
497   llvm::move(stage1Patterns, std::back_inserter(frozenStage1Patterns));
498   FrozenRewritePatternSet stage2Patterns =
499       getLinalgTilingCanonicalizationPatterns(ctx);
500   (void)applyStagedPatterns(funcOp, frozenStage1Patterns,
501                             std::move(stage2Patterns));
502 }
503 
504 static void applyVectorTransferForwardingPatterns(FuncOp funcOp) {
505   RewritePatternSet forwardPattern(funcOp.getContext());
506   forwardPattern.add<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
507   forwardPattern.add<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
508   (void)applyPatternsAndFoldGreedily(funcOp, std::move(forwardPattern));
509 }
510 
511 static void applyLinalgToVectorPatterns(FuncOp funcOp) {
512   RewritePatternSet patterns(funcOp.getContext());
513   patterns.add<LinalgVectorizationPattern>(
514       funcOp.getContext(),
515       LinalgTransformationFilter()
516           .addOpFilter<ContractionOpInterface, FillOp, CopyOp, GenericOp>());
517   populatePadTensorOpVectorizationPatterns(patterns);
518   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
519 }
520 
521 static void applyPadTensorToGenericPatterns(FuncOp funcOp) {
522   RewritePatternSet patterns(funcOp.getContext());
523   patterns.add<PadTensorOpTransformationPattern>(funcOp.getContext());
524   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
525 }
526 
527 static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
528   RewritePatternSet foldPattern(funcOp.getContext());
529   foldPattern.add<AffineMinSCFCanonicalizationPattern>(funcOp.getContext());
530   FrozenRewritePatternSet frozenPatterns(std::move(foldPattern));
531 
532   // Explicitly walk and apply the pattern locally to avoid more general folding
533   // on the rest of the IR.
534   funcOp.walk([&frozenPatterns](AffineMinOp minOp) {
535     (void)applyOpPatternsAndFold(minOp, frozenPatterns);
536   });
537 }
538 
539 // For now, just assume it is the zero of type.
540 // In the future, it should be the zero of type + op.
541 static Value getNeutralOfLinalgOp(OpBuilder &b, OpOperand &op) {
542   auto t = getElementTypeOrSelf(op.get());
543   return b.create<ConstantOp>(op.getOwner()->getLoc(), t, b.getZeroAttr(t));
544 }
545 
546 static void applyTileAndPadPattern(FuncOp funcOp, ArrayRef<int64_t> tileSizes) {
547   MLIRContext *context = funcOp.getContext();
548   RewritePatternSet tilingPattern(context);
549   auto linalgTilingOptions =
550       linalg::LinalgTilingOptions()
551           .setTileSizes(tileSizes)
552           .setPaddingValueComputationFunction(getNeutralOfLinalgOp);
553   tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulI8I8I32Op>,
554                     linalg::LinalgTilingPattern<linalg::GenericOp>>(
555       context, linalgTilingOptions,
556       linalg::LinalgTransformationFilter(
557           Identifier::get("tile-and-pad", context)));
558   (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
559 }
560 
561 static void applyInterchangePattern(FuncOp funcOp,
562                                     ArrayRef<unsigned> interchangeVector) {
563   MLIRContext *context = funcOp.getContext();
564   RewritePatternSet interchangePattern(context);
565   interchangePattern.add<GenericOpInterchangePattern>(
566       context, interchangeVector,
567       LinalgTransformationFilter(ArrayRef<Identifier>{},
568                                  Identifier::get("interchange", context)));
569   (void)applyPatternsAndFoldGreedily(funcOp, std::move(interchangePattern));
570 }
571 
572 /// Apply transformations specified as patterns.
573 void TestLinalgTransforms::runOnFunction() {
574   auto lambda = [&](void *) {
575     getFunction().walk([](LinalgOp op) {
576       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
577     });
578   };
579   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
580 
581   if (testPromotionOptions) {
582     RewritePatternSet patterns(&getContext());
583     fillPromotionCallBackPatterns(&getContext(), patterns);
584     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
585     return;
586   }
587   if (testTileAndDistributionOptions) {
588     RewritePatternSet patterns(&getContext());
589     fillTileAndDistributePatterns(&getContext(), patterns);
590     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
591     return;
592   }
593   if (testPatterns)
594     return applyPatterns(getFunction());
595   if (testMatmulToVectorPatterns1dTiling || testMatmulToVectorPatterns2dTiling)
596     return applyMatmulToVectorPatterns(getFunction(),
597                                        testMatmulToVectorPatterns1dTiling,
598                                        testMatmulToVectorPatterns2dTiling);
599   if (testVectorTransferForwardingPatterns)
600     return applyVectorTransferForwardingPatterns(getFunction());
601   if (testGenericToVectorPattern)
602     return applyLinalgToVectorPatterns(getFunction());
603   if (testTransformPadTensor)
604     return applyPadTensorToGenericPatterns(getFunction());
605   if (testAffineMinSCFCanonicalizationPatterns)
606     return applyAffineMinSCFCanonicalizationPatterns(getFunction());
607   if (testTileAndPadPattern)
608     return applyTileAndPadPattern(getFunction(), tileSizesForPadding);
609   if (testHoistPadding) {
610     getFunction().walk([&](linalg::PadTensorOp padTensorOp) {
611       (void)linalg::hoistPaddingOnTensors(padTensorOp, testHoistPadding);
612     });
613   }
614   if (testInterchangePattern.hasValue())
615     return applyInterchangePattern(getFunction(), testInterchangePattern);
616 }
617 
618 namespace mlir {
619 namespace test {
620 void registerTestLinalgTransforms() {
621   PassRegistration<TestLinalgTransforms>();
622 }
623 } // namespace test
624 } // namespace mlir
625