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