13fef2d26SRiver Riddle //===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
23fef2d26SRiver Riddle //
33fef2d26SRiver Riddle // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
43fef2d26SRiver Riddle // See https://llvm.org/LICENSE.txt for license information.
53fef2d26SRiver Riddle // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
63fef2d26SRiver Riddle //
73fef2d26SRiver Riddle //===----------------------------------------------------------------------===//
83fef2d26SRiver Riddle //
93fef2d26SRiver Riddle // This file implements logic for testing Linalg transformations.
103fef2d26SRiver Riddle //
113fef2d26SRiver Riddle //===----------------------------------------------------------------------===//
123fef2d26SRiver Riddle 
133fef2d26SRiver Riddle #include "mlir/Dialect/Affine/IR/AffineOps.h"
14a54f4eaeSMogball #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
15178f9bd6SNicolas Vasilache #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
1636550692SRiver Riddle #include "mlir/Dialect/Func/IR/FuncOps.h"
17d7ef488bSMogball #include "mlir/Dialect/GPU/IR/GPUDialect.h"
18b7f2c108Sgysit #include "mlir/Dialect/Linalg/IR/Linalg.h"
19800694a6SNicolas Vasilache #include "mlir/Dialect/Linalg/Passes.h"
20b74493ecSNicolas Vasilache #include "mlir/Dialect/Linalg/Transforms/HoistPadding.h"
213fef2d26SRiver Riddle #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
223fef2d26SRiver Riddle #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
233fef2d26SRiver Riddle #include "mlir/Dialect/Linalg/Utils/Utils.h"
2499ef9eebSMatthias Springer #include "mlir/Dialect/Vector/IR/VectorOps.h"
25800694a6SNicolas Vasilache #include "mlir/Pass/PassManager.h"
263fef2d26SRiver Riddle #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
273fef2d26SRiver Riddle 
283fef2d26SRiver Riddle #include "llvm/ADT/SetVector.h"
29c57c4f88SMatthias Springer #include "llvm/ADT/SmallVector.h"
303fef2d26SRiver Riddle 
313fef2d26SRiver Riddle using namespace mlir;
323fef2d26SRiver Riddle using namespace mlir::linalg;
333fef2d26SRiver Riddle 
343fef2d26SRiver Riddle namespace {
353fef2d26SRiver Riddle struct TestLinalgTransforms
3658ceae95SRiver Riddle     : public PassWrapper<TestLinalgTransforms, OperationPass<func::FuncOp>> {
375e50dd04SRiver Riddle   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgTransforms)
385e50dd04SRiver Riddle 
393fef2d26SRiver Riddle   TestLinalgTransforms() = default;
TestLinalgTransforms__anon52d1bf190111::TestLinalgTransforms403bab9d4eSMehdi Amini   TestLinalgTransforms(const TestLinalgTransforms &pass) : PassWrapper(pass) {}
413fef2d26SRiver Riddle 
getDependentDialects__anon52d1bf190111::TestLinalgTransforms423fef2d26SRiver Riddle   void getDependentDialects(DialectRegistry &registry) const override {
433fef2d26SRiver Riddle     // clang-format off
443fef2d26SRiver Riddle     registry.insert<AffineDialect,
45178f9bd6SNicolas Vasilache                     bufferization::BufferizationDialect,
463fef2d26SRiver Riddle                     memref::MemRefDialect,
473fef2d26SRiver Riddle                     scf::SCFDialect,
48fd0c6f53SAlexander Belyaev                     linalg::LinalgDialect,
493fef2d26SRiver Riddle                     vector::VectorDialect,
503fef2d26SRiver Riddle                     gpu::GPUDialect>();
513fef2d26SRiver Riddle     // clang-format on
523fef2d26SRiver Riddle   }
getArgument__anon52d1bf190111::TestLinalgTransforms53b5e22e6dSMehdi Amini   StringRef getArgument() const final {
54b5e22e6dSMehdi Amini     return "test-linalg-transform-patterns";
55b5e22e6dSMehdi Amini   }
getDescription__anon52d1bf190111::TestLinalgTransforms56b5e22e6dSMehdi Amini   StringRef getDescription() const final {
57b5e22e6dSMehdi Amini     return "Test Linalg transformation patterns by applying them greedily.";
58b5e22e6dSMehdi Amini   }
593fef2d26SRiver Riddle 
6041574554SRiver Riddle   void runOnOperation() override;
613fef2d26SRiver Riddle 
623fef2d26SRiver Riddle   Option<bool> testPatterns{*this, "test-patterns",
633fef2d26SRiver Riddle                             llvm::cl::desc("Test a mixed set of patterns"),
643fef2d26SRiver Riddle                             llvm::cl::init(false)};
653fef2d26SRiver Riddle   Option<bool> testTileAndDistributionOptions{
663fef2d26SRiver Riddle       *this, "test-tile-and-distribute-options",
673fef2d26SRiver Riddle       llvm::cl::desc("Test tile and distribute options"),
683fef2d26SRiver Riddle       llvm::cl::init(false)};
69748bf4bbSHanhan Wang   Option<bool> testTileFuseAndDistributionOptions{
70748bf4bbSHanhan Wang       *this, "test-tile-fuse-and-distribute-options",
71748bf4bbSHanhan Wang       llvm::cl::desc("Test tile, fuse and distribute options"),
72748bf4bbSHanhan Wang       llvm::cl::init(false)};
733fef2d26SRiver Riddle   Option<bool> testVectorTransferForwardingPatterns{
743fef2d26SRiver Riddle       *this, "test-vector-transfer-forwarding-patterns",
753fef2d26SRiver Riddle       llvm::cl::desc(
76ebc81537SAlexander Belyaev           "Test a fused pass that forwards memref.copy to vector.transfer"),
773fef2d26SRiver Riddle       llvm::cl::init(false)};
783fef2d26SRiver Riddle   Option<bool> testGenericToVectorPattern{
793fef2d26SRiver Riddle       *this, "test-linalg-to-vector-patterns",
803fef2d26SRiver Riddle       llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
813fef2d26SRiver Riddle                      "in vector.contract form"),
823fef2d26SRiver Riddle       llvm::cl::init(false)};
838faf35c0SMatthias Springer   Option<bool> testTilePattern{*this, "test-tile-pattern",
848faf35c0SMatthias Springer                                llvm::cl::desc("Test tile pattern"),
858faf35c0SMatthias Springer                                llvm::cl::init(false)};
86fb1def9cSMatthias Springer   Option<bool> testTileScalarizeDynamicDims{
87fb1def9cSMatthias Springer       *this, "test-tile-scalarize-dynamic-dims",
88fb1def9cSMatthias Springer       llvm::cl::desc("Test tiling of dynamic dims by 1"),
89fb1def9cSMatthias Springer       llvm::cl::init(false)};
900804a88eSNicolas Agostini   Option<bool> testTransformPadTensor{
910804a88eSNicolas Agostini       *this, "test-transform-pad-tensor",
920804a88eSNicolas Agostini       llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
930804a88eSNicolas Agostini       llvm::cl::init(false)};
9435df2f6fSYi Zhang   Option<bool> testGeneralizePadTensor{
9535df2f6fSYi Zhang       *this, "test-generalize-pad-tensor",
9635df2f6fSYi Zhang       llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
9735df2f6fSYi Zhang       llvm::cl::init(false)};
9824199f53SMatthias Springer   Option<bool> testSwapSubTensorPadTensor{
9924199f53SMatthias Springer       *this, "test-swap-subtensor-padtensor",
1001ad9b266Slorenzo chelini       llvm::cl::desc("Test rewrite of subtensor(tensor.pad) into "
1011ad9b266Slorenzo chelini                      "tensor.pad(subtensor)"),
10224199f53SMatthias Springer       llvm::cl::init(false)};
10333d2a780SThomas Raoux   Option<bool> testSplitReduction{
10433d2a780SThomas Raoux       *this, "test-split-reduction",
10533d2a780SThomas Raoux       llvm::cl::desc("Test split reduction transformation"),
10633d2a780SThomas Raoux       llvm::cl::init(false)};
1078faf35c0SMatthias Springer   ListOption<int64_t> peeledLoops{
1088faf35c0SMatthias Springer       *this, "peeled-loops",
109*62a4e6abSFangrui Song       llvm::cl::desc("Loops to be peeled when test-tile-pattern")};
1108faf35c0SMatthias Springer   ListOption<int64_t> tileSizes{
1118faf35c0SMatthias Springer       *this, "tile-sizes",
112*62a4e6abSFangrui Song       llvm::cl::desc("Linalg tile sizes for test-tile-pattern")};
113a4a654d3SMatthias Springer   Option<bool> skipPartial{
114a4a654d3SMatthias Springer       *this, "skip-partial",
115a4a654d3SMatthias Springer       llvm::cl::desc("Skip loops inside partial iterations during peeling"),
116a4a654d3SMatthias Springer       llvm::cl::init(false)};
1172190f8a8SMatthias Springer   Option<std::string> loopType{
1182190f8a8SMatthias Springer       *this, "loop-type",
1192190f8a8SMatthias Springer       llvm::cl::desc("Specify the type of loops to generate: for, parallel or "
1202190f8a8SMatthias Springer                      "tiled_loop"),
1212190f8a8SMatthias Springer       llvm::cl::init("for")};
12265bdeddbSOkwan Kwon   Option<bool> testBubbleUpExtractSliceOpPattern{
12365bdeddbSOkwan Kwon       *this, "test-bubble-up-extract-slice-op-pattern",
12465bdeddbSOkwan Kwon       llvm::cl::desc("Test rewrite of linalgOp + extract_slice into "
12565bdeddbSOkwan Kwon                      "extract_slice + linalgOp"),
12665bdeddbSOkwan Kwon       llvm::cl::init(false)};
1273fef2d26SRiver Riddle };
128be0a7e9fSMehdi Amini } // namespace
1293fef2d26SRiver Riddle 
applyPatterns(func::FuncOp funcOp)13058ceae95SRiver Riddle static void applyPatterns(func::FuncOp funcOp) {
1313fef2d26SRiver Riddle   MLIRContext *ctx = funcOp.getContext();
1323fef2d26SRiver Riddle   RewritePatternSet patterns(ctx);
1333fef2d26SRiver Riddle 
1343fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
1353fef2d26SRiver Riddle   // Linalg tiling patterns.
1363fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
1374a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1384a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
1394a661602SNicolas Vasilache       LinalgTilingOptions().setTileSizes({2000, 3000, 4000}),
140195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "MEM"),
141195730a6SRiver Riddle                                  StringAttr::get(ctx, "L3")));
1424a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1434a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
1444a661602SNicolas Vasilache       LinalgTilingOptions().setTileSizes({200, 300, 400}),
145195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "L3"),
146195730a6SRiver Riddle                                  StringAttr::get(ctx, "L2")));
1474a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1484a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
1494a661602SNicolas Vasilache       LinalgTilingOptions().setTileSizes({20, 30, 40}),
150195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "L2"),
151195730a6SRiver Riddle                                  StringAttr::get(ctx, "L1")));
1524a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1534a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
1544a661602SNicolas Vasilache       LinalgTilingOptions().setTileSizes({2, 3, 4}),
155195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "L1"),
156195730a6SRiver Riddle                                  StringAttr::get(ctx, "REG")));
1573fef2d26SRiver Riddle 
1584a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1594a661602SNicolas Vasilache       MatvecOp::getOperationName(), ctx,
1603fef2d26SRiver Riddle       LinalgTilingOptions().setTileSizes({5, 6}).setLoopType(
1613fef2d26SRiver Riddle           LinalgTilingLoopType::ParallelLoops),
162195730a6SRiver Riddle       LinalgTransformationFilter(ArrayRef<StringAttr>{},
163195730a6SRiver Riddle                                  StringAttr::get(ctx, "L1")));
1643fef2d26SRiver Riddle 
1654a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1664a661602SNicolas Vasilache       DotOp::getOperationName(), ctx, LinalgTilingOptions().setTileSizes(8000),
1673fef2d26SRiver Riddle       LinalgTransformationFilter(
168195730a6SRiver Riddle           ArrayRef<StringAttr>{StringAttr::get(ctx, "MEM"),
169195730a6SRiver Riddle                                StringAttr::get(ctx, "L3"),
170195730a6SRiver Riddle                                StringAttr::get(ctx, "L2")},
171195730a6SRiver Riddle           StringAttr::get(ctx, "REG")));
1723fef2d26SRiver Riddle 
1733fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
1743fef2d26SRiver Riddle   // Linalg tiling and permutation patterns.
1753fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
1764a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1774a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
1783fef2d26SRiver Riddle       LinalgTilingOptions()
1793fef2d26SRiver Riddle           .setTileSizes({2000, 3000, 4000})
1803fef2d26SRiver Riddle           .setInterchange({1, 2, 0}),
181195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "__with_perm__"),
182195730a6SRiver Riddle                                  StringAttr::get(ctx, "L2__with_perm__")));
1834a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1844a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
1853fef2d26SRiver Riddle       LinalgTilingOptions()
1863fef2d26SRiver Riddle           .setTileSizes({200, 300, 400})
1873fef2d26SRiver Riddle           .setInterchange({1, 0, 2}),
188195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "L2__with_perm__"),
189195730a6SRiver Riddle                                  StringAttr::get(ctx, "L1__with_perm__")));
1904a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1914a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
1924a661602SNicolas Vasilache       LinalgTilingOptions().setTileSizes({20, 30, 40}),
193195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "L1__with_perm__"),
194195730a6SRiver Riddle                                  StringAttr::get(ctx, "REG__with_perm__")));
1953fef2d26SRiver Riddle 
1964a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
1974a661602SNicolas Vasilache       MatvecOp::getOperationName(), ctx,
1984a661602SNicolas Vasilache       LinalgTilingOptions().setTileSizes({5, 6}).setInterchange({1, 0}),
199195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "__with_perm__"),
200195730a6SRiver Riddle                                  StringAttr::get(ctx, "L1__with_perm__")));
2013fef2d26SRiver Riddle 
2024a661602SNicolas Vasilache   patterns.add<LinalgTilingPattern>(
2034a661602SNicolas Vasilache       MatmulOp::getOperationName(), ctx,
2043fef2d26SRiver Riddle       LinalgTilingOptions()
2053fef2d26SRiver Riddle           .setTileSizes({16, 8, 4})
2063fef2d26SRiver Riddle           .setInterchange({1, 2, 0})
2073fef2d26SRiver Riddle           .setLoopType(LinalgTilingLoopType::ParallelLoops),
2083fef2d26SRiver Riddle       LinalgTransformationFilter(
209195730a6SRiver Riddle           StringAttr::get(ctx, "par__with_perm__"),
210195730a6SRiver Riddle           StringAttr::get(ctx, "after_par__with_perm__")));
2113fef2d26SRiver Riddle 
2123fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2133fef2d26SRiver Riddle   // Linalg to loops patterns.
2143fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2153fef2d26SRiver Riddle   patterns.add<LinalgLoweringPattern<DotOp>>(
2163fef2d26SRiver Riddle       ctx,
2173fef2d26SRiver Riddle       /*loweringType=*/LinalgLoweringType::Loops,
218195730a6SRiver Riddle       LinalgTransformationFilter(StringAttr::get(ctx, "REG")));
2193fef2d26SRiver Riddle 
2203fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2213fef2d26SRiver Riddle   // Linalg distribution patterns.
2223fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2233fef2d26SRiver Riddle   LinalgLoopDistributionOptions distributionOptions;
2243fef2d26SRiver Riddle 
2253fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2263fef2d26SRiver Riddle   // Linalg to vector contraction patterns.
2273fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2283fef2d26SRiver Riddle   patterns.add<LinalgVectorizationPattern>(
229195730a6SRiver Riddle       ctx, LinalgTransformationFilter(StringAttr::get(ctx, "VECTORIZE"))
230ebc81537SAlexander Belyaev                .addOpFilter<MatmulOp, FillOp, GenericOp>());
231ebc81537SAlexander Belyaev   patterns.add<CopyVectorizationPattern>(ctx);
2323fef2d26SRiver Riddle 
2333fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2343fef2d26SRiver Riddle   // Linalg generic interchange pattern.
2353fef2d26SRiver Riddle   //===--------------------------------------------------------------------===//
2363fef2d26SRiver Riddle   patterns.add<GenericOpInterchangePattern>(
2373fef2d26SRiver Riddle       ctx,
2383fef2d26SRiver Riddle       /*interchangeVector=*/ArrayRef<unsigned>{1, 2, 0},
239195730a6SRiver Riddle       LinalgTransformationFilter(ArrayRef<StringAttr>{},
240195730a6SRiver Riddle                                  StringAttr::get(ctx, "PERMUTED")));
2413fef2d26SRiver Riddle 
2423fef2d26SRiver Riddle   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
2433fef2d26SRiver Riddle 
2443fef2d26SRiver Riddle   // Drop the marker.
2453fef2d26SRiver Riddle   funcOp.walk([](LinalgOp op) {
2463fef2d26SRiver Riddle     op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
2473fef2d26SRiver Riddle   });
2483fef2d26SRiver Riddle }
2493fef2d26SRiver Riddle 
2503fef2d26SRiver Riddle template <typename IdOp, typename NProcsOp>
2513fef2d26SRiver Riddle static SmallVector<ProcInfo, 2>
getGpuProcIds(OpBuilder & b,Location loc,ArrayRef<Range> parallelLoopRanges)2523fef2d26SRiver Riddle getGpuProcIds(OpBuilder &b, Location loc, ArrayRef<Range> parallelLoopRanges) {
2533fef2d26SRiver Riddle   size_t count = std::min<size_t>(3, parallelLoopRanges.size());
2543fef2d26SRiver Riddle   SmallVector<ProcInfo, 2> procInfo(count);
2553fef2d26SRiver Riddle   Type indexType = b.getIndexType();
2563fef2d26SRiver Riddle   for (unsigned i = 0; i < count; ++i) {
257aae51255SMogball     gpu::Dimension dim = *gpu::symbolizeDimension(i);
258aae51255SMogball     procInfo[count - 1 - i] = {b.create<IdOp>(loc, indexType, dim),
259aae51255SMogball                                b.create<NProcsOp>(loc, indexType, dim)};
2603fef2d26SRiver Riddle   }
2613fef2d26SRiver Riddle   return procInfo;
2623fef2d26SRiver Riddle }
2633fef2d26SRiver Riddle 
fillTileAndDistributePatterns(MLIRContext * context,RewritePatternSet & patterns)2643fef2d26SRiver Riddle static void fillTileAndDistributePatterns(MLIRContext *context,
2653fef2d26SRiver Riddle                                           RewritePatternSet &patterns) {
2663fef2d26SRiver Riddle   {
2673fef2d26SRiver Riddle     LinalgLoopDistributionOptions cyclicNprocsEqNiters;
2683fef2d26SRiver Riddle     cyclicNprocsEqNiters.distributionMethod.resize(
2693fef2d26SRiver Riddle         2, DistributionMethod::CyclicNumProcsEqNumIters);
2703fef2d26SRiver Riddle     cyclicNprocsEqNiters.procInfo =
2713fef2d26SRiver Riddle         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
2724a661602SNicolas Vasilache     patterns.add<LinalgTilingPattern>(
2734a661602SNicolas Vasilache         MatmulOp::getOperationName(), context,
2743fef2d26SRiver Riddle         LinalgTilingOptions()
2753fef2d26SRiver Riddle             .setTileSizes({8, 8, 4})
2763fef2d26SRiver Riddle             .setLoopType(LinalgTilingLoopType::ParallelLoops)
2773fef2d26SRiver Riddle             .setDistributionOptions(cyclicNprocsEqNiters),
2783fef2d26SRiver Riddle         LinalgTransformationFilter(
279195730a6SRiver Riddle             StringAttr::get(context, "distribute1"),
280195730a6SRiver Riddle             StringAttr::get(context, "after_distribute1")));
2813fef2d26SRiver Riddle   }
2823fef2d26SRiver Riddle 
2833fef2d26SRiver Riddle   {
2843fef2d26SRiver Riddle     LinalgLoopDistributionOptions cyclicNprocsGeNiters;
2853fef2d26SRiver Riddle     cyclicNprocsGeNiters.distributionMethod.resize(
2863fef2d26SRiver Riddle         2, DistributionMethod::CyclicNumProcsGeNumIters);
2873fef2d26SRiver Riddle     cyclicNprocsGeNiters.procInfo =
2883fef2d26SRiver Riddle         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
2894a661602SNicolas Vasilache     patterns.add<LinalgTilingPattern>(
2904a661602SNicolas Vasilache         MatmulOp::getOperationName(), context,
2913fef2d26SRiver Riddle         LinalgTilingOptions()
2923fef2d26SRiver Riddle             .setTileSizes({8, 8, 4})
2933fef2d26SRiver Riddle             .setLoopType(LinalgTilingLoopType::ParallelLoops)
2943fef2d26SRiver Riddle             .setDistributionOptions(cyclicNprocsGeNiters),
2953fef2d26SRiver Riddle         LinalgTransformationFilter(
296195730a6SRiver Riddle             StringAttr::get(context, "distribute2"),
297195730a6SRiver Riddle             StringAttr::get(context, "after_distribute2")));
2983fef2d26SRiver Riddle   }
2993fef2d26SRiver Riddle 
3003fef2d26SRiver Riddle   {
3013fef2d26SRiver Riddle     LinalgLoopDistributionOptions cyclicNprocsDefault;
3023fef2d26SRiver Riddle     cyclicNprocsDefault.distributionMethod.resize(2,
3033fef2d26SRiver Riddle                                                   DistributionMethod::Cyclic);
3043fef2d26SRiver Riddle     cyclicNprocsDefault.procInfo =
3053fef2d26SRiver Riddle         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
3064a661602SNicolas Vasilache     patterns.add<LinalgTilingPattern>(
3074a661602SNicolas Vasilache         MatmulOp::getOperationName(), context,
3083fef2d26SRiver Riddle         LinalgTilingOptions()
3093fef2d26SRiver Riddle             .setTileSizes({8, 8, 4})
3103fef2d26SRiver Riddle             .setLoopType(LinalgTilingLoopType::ParallelLoops)
3113fef2d26SRiver Riddle             .setDistributionOptions(cyclicNprocsDefault),
3123fef2d26SRiver Riddle         LinalgTransformationFilter(
313195730a6SRiver Riddle             StringAttr::get(context, "distribute3"),
314195730a6SRiver Riddle             StringAttr::get(context, "after_distribute3")));
3153fef2d26SRiver Riddle   }
3163fef2d26SRiver Riddle 
3173fef2d26SRiver Riddle   {
3183fef2d26SRiver Riddle     LinalgLoopDistributionOptions cyclicNprocsMixed1;
3193fef2d26SRiver Riddle     cyclicNprocsMixed1.distributionMethod = {
3203fef2d26SRiver Riddle         DistributionMethod::CyclicNumProcsEqNumIters,
3213fef2d26SRiver Riddle         DistributionMethod::CyclicNumProcsGeNumIters};
3223fef2d26SRiver Riddle     cyclicNprocsMixed1.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
3234a661602SNicolas Vasilache     patterns.add<LinalgTilingPattern>(
3244a661602SNicolas Vasilache         MatmulOp::getOperationName(), context,
3253fef2d26SRiver Riddle         LinalgTilingOptions()
3263fef2d26SRiver Riddle             .setTileSizes({8, 8, 4})
3273fef2d26SRiver Riddle             .setLoopType(LinalgTilingLoopType::ParallelLoops)
3283fef2d26SRiver Riddle             .setDistributionOptions(cyclicNprocsMixed1),
3293fef2d26SRiver Riddle         LinalgTransformationFilter(
330195730a6SRiver Riddle             StringAttr::get(context, "distribute4"),
331195730a6SRiver Riddle             StringAttr::get(context, "after_distribute4")));
3323fef2d26SRiver Riddle   }
3333fef2d26SRiver Riddle 
3343fef2d26SRiver Riddle   {
3353fef2d26SRiver Riddle     LinalgLoopDistributionOptions cyclicNprocsMixed2;
3363fef2d26SRiver Riddle     cyclicNprocsMixed2.distributionMethod = {
3373fef2d26SRiver Riddle         DistributionMethod::CyclicNumProcsGeNumIters,
3383fef2d26SRiver Riddle         DistributionMethod::Cyclic};
3393fef2d26SRiver Riddle     cyclicNprocsMixed2.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
3404a661602SNicolas Vasilache     patterns.add<LinalgTilingPattern>(
3414a661602SNicolas Vasilache         MatmulOp::getOperationName(), context,
3423fef2d26SRiver Riddle         LinalgTilingOptions()
3433fef2d26SRiver Riddle             .setTileSizes({8, 8, 4})
3443fef2d26SRiver Riddle             .setLoopType(LinalgTilingLoopType::ParallelLoops)
3453fef2d26SRiver Riddle             .setDistributionOptions(cyclicNprocsMixed2),
3463fef2d26SRiver Riddle         LinalgTransformationFilter(
347195730a6SRiver Riddle             StringAttr::get(context, "distribute5"),
348195730a6SRiver Riddle             StringAttr::get(context, "after_distribute5")));
3493fef2d26SRiver Riddle   }
3503fef2d26SRiver Riddle 
3513fef2d26SRiver Riddle   {
3523fef2d26SRiver Riddle     LinalgLoopDistributionOptions cyclicNprocsMixed3;
3533fef2d26SRiver Riddle     cyclicNprocsMixed3.distributionMethod = {
3543fef2d26SRiver Riddle         DistributionMethod::Cyclic,
3553fef2d26SRiver Riddle         DistributionMethod::CyclicNumProcsEqNumIters};
3563fef2d26SRiver Riddle     cyclicNprocsMixed3.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
3573fef2d26SRiver Riddle 
3584a661602SNicolas Vasilache     patterns.add<LinalgTilingPattern>(
3594a661602SNicolas Vasilache         MatmulOp::getOperationName(), context,
3603fef2d26SRiver Riddle         LinalgTilingOptions()
3613fef2d26SRiver Riddle             .setTileSizes({8, 8, 4})
3623fef2d26SRiver Riddle             .setLoopType(LinalgTilingLoopType::ParallelLoops)
3633fef2d26SRiver Riddle             .setDistributionOptions(cyclicNprocsMixed3),
3643fef2d26SRiver Riddle         LinalgTransformationFilter(
365195730a6SRiver Riddle             StringAttr::get(context, "distribute6"),
366195730a6SRiver Riddle             StringAttr::get(context, "after_distribute6")));
3673fef2d26SRiver Riddle   }
3683fef2d26SRiver Riddle 
3693fef2d26SRiver Riddle   {
3703fef2d26SRiver Riddle     LinalgLoopDistributionOptions cyclicNprocsEqNiters;
3713fef2d26SRiver Riddle     cyclicNprocsEqNiters.distributionMethod.resize(2,
3723fef2d26SRiver Riddle                                                    DistributionMethod::Cyclic);
3733fef2d26SRiver Riddle     cyclicNprocsEqNiters.procInfo =
3743fef2d26SRiver Riddle         getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
3754a661602SNicolas Vasilache     patterns.add<LinalgTilingPattern>(
3764a661602SNicolas Vasilache         MatmulOp::getOperationName(), context,
3773fef2d26SRiver Riddle         LinalgTilingOptions()
3783fef2d26SRiver Riddle             .setTileSizes({8, 8, 4})
3793fef2d26SRiver Riddle             .setLoopType(LinalgTilingLoopType::Loops)
3803fef2d26SRiver Riddle             .setDistributionOptions(cyclicNprocsEqNiters),
3813fef2d26SRiver Riddle         LinalgTransformationFilter(
382195730a6SRiver Riddle             StringAttr::get(context, "tensors_distribute1"),
383195730a6SRiver Riddle             StringAttr::get(context, "tensors_after_distribute1")));
3843fef2d26SRiver Riddle   }
3853fef2d26SRiver Riddle }
3863fef2d26SRiver Riddle 
fillTileFuseAndDistributePatterns(MLIRContext * context,RewritePatternSet & patterns)387748bf4bbSHanhan Wang static void fillTileFuseAndDistributePatterns(MLIRContext *context,
388748bf4bbSHanhan Wang                                               RewritePatternSet &patterns) {
389748bf4bbSHanhan Wang   LinalgLoopDistributionOptions cyclicNprocsEqNiters;
390748bf4bbSHanhan Wang   cyclicNprocsEqNiters.distributionMethod.resize(2, DistributionMethod::Cyclic);
391748bf4bbSHanhan Wang   cyclicNprocsEqNiters.procInfo = getGpuProcIds<gpu::BlockIdOp, gpu::GridDimOp>;
392748bf4bbSHanhan Wang   patterns.add<LinalgTileAndFuseTensorOpsPattern>(
393748bf4bbSHanhan Wang       MatmulOp::getOperationName(), context,
394748bf4bbSHanhan Wang       LinalgTilingAndFusionOptions()
395748bf4bbSHanhan Wang           .setTileSizes({8, 8, 4})
396748bf4bbSHanhan Wang           .setDistributionOptions(cyclicNprocsEqNiters),
397748bf4bbSHanhan Wang       LinalgTransformationFilter(
398748bf4bbSHanhan Wang           StringAttr::get(context, "tensors_fuse_distribute1"),
399748bf4bbSHanhan Wang           StringAttr::get(context, "tensors_after_fuse_distribute1")));
400748bf4bbSHanhan Wang }
401748bf4bbSHanhan Wang 
applyVectorTransferForwardingPatterns(func::FuncOp funcOp)40258ceae95SRiver Riddle static void applyVectorTransferForwardingPatterns(func::FuncOp funcOp) {
4033fef2d26SRiver Riddle   RewritePatternSet forwardPattern(funcOp.getContext());
4043fef2d26SRiver Riddle   forwardPattern.add<LinalgCopyVTRForwardingPattern>(funcOp.getContext());
4053fef2d26SRiver Riddle   forwardPattern.add<LinalgCopyVTWForwardingPattern>(funcOp.getContext());
4063fef2d26SRiver Riddle   (void)applyPatternsAndFoldGreedily(funcOp, std::move(forwardPattern));
4073fef2d26SRiver Riddle }
4083fef2d26SRiver Riddle 
applyLinalgToVectorPatterns(func::FuncOp funcOp)40958ceae95SRiver Riddle static void applyLinalgToVectorPatterns(func::FuncOp funcOp) {
4103fef2d26SRiver Riddle   RewritePatternSet patterns(funcOp.getContext());
411ebc81537SAlexander Belyaev   auto *ctx = funcOp.getContext();
4123fef2d26SRiver Riddle   patterns.add<LinalgVectorizationPattern>(
413ebc81537SAlexander Belyaev       ctx, LinalgTransformationFilter()
414ebc81537SAlexander Belyaev                .addOpFilter<ContractionOpInterface, FillOp, GenericOp>());
415ebc81537SAlexander Belyaev   patterns.add<CopyVectorizationPattern>(ctx);
416fd0c6f53SAlexander Belyaev   populatePadOpVectorizationPatterns(patterns);
4176bb7d247SNicolas Vasilache   populateConvolutionVectorizationPatterns(patterns);
4183fef2d26SRiver Riddle   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
4193fef2d26SRiver Riddle }
4203fef2d26SRiver Riddle 
applyPadTensorToGenericPatterns(func::FuncOp funcOp)42158ceae95SRiver Riddle static void applyPadTensorToGenericPatterns(func::FuncOp funcOp) {
4220804a88eSNicolas Agostini   RewritePatternSet patterns(funcOp.getContext());
423fd0c6f53SAlexander Belyaev   patterns.add<PadOpTransformationPattern>(funcOp.getContext());
4240804a88eSNicolas Agostini   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
4250804a88eSNicolas Agostini }
4260804a88eSNicolas Agostini 
applyGeneralizePadTensorPatterns(func::FuncOp funcOp)42758ceae95SRiver Riddle static void applyGeneralizePadTensorPatterns(func::FuncOp funcOp) {
42835df2f6fSYi Zhang   RewritePatternSet patterns(funcOp.getContext());
429fd0c6f53SAlexander Belyaev   patterns.add<GeneralizePadOpPattern>(funcOp.getContext());
43035df2f6fSYi Zhang   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
43135df2f6fSYi Zhang }
43235df2f6fSYi Zhang 
applyExtractSliceOfPadTensorSwapPattern(func::FuncOp funcOp)43358ceae95SRiver Riddle static void applyExtractSliceOfPadTensorSwapPattern(func::FuncOp funcOp) {
43424199f53SMatthias Springer   RewritePatternSet patterns(funcOp.getContext());
435060208b4SMatthias Springer   patterns.add<ExtractSliceOfPadTensorSwapPattern>(funcOp.getContext());
43624199f53SMatthias Springer   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
43724199f53SMatthias Springer }
43824199f53SMatthias Springer 
applyTilePattern(func::FuncOp funcOp,const std::string & loopType,ArrayRef<int64_t> tileSizes,ArrayRef<int64_t> peeledLoops,bool scalarizeDynamicDims)43958ceae95SRiver Riddle static void applyTilePattern(func::FuncOp funcOp, const std::string &loopType,
440d20d0e14STobias Gysi                              ArrayRef<int64_t> tileSizes,
4412190f8a8SMatthias Springer                              ArrayRef<int64_t> peeledLoops,
442fb1def9cSMatthias Springer                              bool scalarizeDynamicDims) {
4433fef2d26SRiver Riddle   MLIRContext *context = funcOp.getContext();
4443fef2d26SRiver Riddle   RewritePatternSet tilingPattern(context);
4452190f8a8SMatthias Springer   LinalgTilingLoopType type =
4462190f8a8SMatthias Springer       llvm::StringSwitch<LinalgTilingLoopType>(loopType)
4472190f8a8SMatthias Springer           .Case("for", LinalgTilingLoopType::Loops)
4482190f8a8SMatthias Springer           .Case("affine", LinalgTilingLoopType::AffineLoops)
4491a829d2dSAlexander Belyaev           .Case("parallel", LinalgTilingLoopType::ParallelLoops);
4502190f8a8SMatthias Springer   auto linalgTilingOptions = linalg::LinalgTilingOptions()
4512190f8a8SMatthias Springer                                  .setPeeledLoops(peeledLoops)
4522190f8a8SMatthias Springer                                  .setLoopType(type);
453fb1def9cSMatthias Springer   if (scalarizeDynamicDims) {
454fb1def9cSMatthias Springer     linalgTilingOptions.scalarizeDynamicDims();
455fb1def9cSMatthias Springer     assert(tileSizes.empty() &&
456fb1def9cSMatthias Springer            "tileSizes and scalarizeDynamicDims is mutually exclusive");
457fb1def9cSMatthias Springer   } else {
458fb1def9cSMatthias Springer     linalgTilingOptions.setTileSizes(tileSizes);
459fb1def9cSMatthias Springer   }
4604a661602SNicolas Vasilache   linalg::LinalgTransformationFilter f(StringAttr::get(context, "tile"));
4614a661602SNicolas Vasilache   TilingPatterns<linalg::MatmulOp, linalg::GenericOp>::insert(
4624a661602SNicolas Vasilache       tilingPattern, linalgTilingOptions, f);
4633fef2d26SRiver Riddle   (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
4643fef2d26SRiver Riddle }
4653fef2d26SRiver Riddle 
applySplitReduction(func::FuncOp funcOp)46658ceae95SRiver Riddle static void applySplitReduction(func::FuncOp funcOp) {
46733d2a780SThomas Raoux   RewritePatternSet patterns(funcOp.getContext());
46833d2a780SThomas Raoux   linalg::populateSplitReductionPattern(
46933d2a780SThomas Raoux       patterns,
47033d2a780SThomas Raoux       [](LinalgOp op) {
47133d2a780SThomas Raoux         unsigned insertDimIndex = op.getNumLoops() - 1;
47233d2a780SThomas Raoux         return std::make_pair(4, insertDimIndex);
47333d2a780SThomas Raoux       },
47433d2a780SThomas Raoux       LinalgTransformationFilter(
47533d2a780SThomas Raoux           ArrayRef<StringAttr>{},
47633d2a780SThomas Raoux           StringAttr::get(funcOp.getContext(), "SPLIT")));
47733d2a780SThomas Raoux   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
47833d2a780SThomas Raoux }
47933d2a780SThomas Raoux 
applyBubbleUpExtractSliceOpPattern(func::FuncOp funcOp)48058ceae95SRiver Riddle static void applyBubbleUpExtractSliceOpPattern(func::FuncOp funcOp) {
48165bdeddbSOkwan Kwon   RewritePatternSet patterns(funcOp.getContext());
48265bdeddbSOkwan Kwon   populateBubbleUpExtractSliceOpPatterns(patterns);
48365bdeddbSOkwan Kwon   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
48465bdeddbSOkwan Kwon }
48565bdeddbSOkwan Kwon 
4863fef2d26SRiver Riddle /// Apply transformations specified as patterns.
runOnOperation()48741574554SRiver Riddle void TestLinalgTransforms::runOnOperation() {
4883fef2d26SRiver Riddle   auto lambda = [&](void *) {
48941574554SRiver Riddle     getOperation().walk([](LinalgOp op) {
4903fef2d26SRiver Riddle       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
4913fef2d26SRiver Riddle     });
4923fef2d26SRiver Riddle   };
4933fef2d26SRiver Riddle   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
4943fef2d26SRiver Riddle 
4953fef2d26SRiver Riddle   if (testTileAndDistributionOptions) {
4963fef2d26SRiver Riddle     RewritePatternSet patterns(&getContext());
4973fef2d26SRiver Riddle     fillTileAndDistributePatterns(&getContext(), patterns);
49841574554SRiver Riddle     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
4993fef2d26SRiver Riddle     return;
5003fef2d26SRiver Riddle   }
501748bf4bbSHanhan Wang   if (testTileFuseAndDistributionOptions) {
502748bf4bbSHanhan Wang     RewritePatternSet patterns(&getContext());
503748bf4bbSHanhan Wang     fillTileFuseAndDistributePatterns(&getContext(), patterns);
504748bf4bbSHanhan Wang     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
505748bf4bbSHanhan Wang     return;
506748bf4bbSHanhan Wang   }
5073fef2d26SRiver Riddle   if (testPatterns)
50841574554SRiver Riddle     return applyPatterns(getOperation());
5093fef2d26SRiver Riddle   if (testVectorTransferForwardingPatterns)
51041574554SRiver Riddle     return applyVectorTransferForwardingPatterns(getOperation());
5113fef2d26SRiver Riddle   if (testGenericToVectorPattern)
51241574554SRiver Riddle     return applyLinalgToVectorPatterns(getOperation());
5130804a88eSNicolas Agostini   if (testTransformPadTensor)
51441574554SRiver Riddle     return applyPadTensorToGenericPatterns(getOperation());
51535df2f6fSYi Zhang   if (testGeneralizePadTensor)
51641574554SRiver Riddle     return applyGeneralizePadTensorPatterns(getOperation());
51724199f53SMatthias Springer   if (testSwapSubTensorPadTensor)
51841574554SRiver Riddle     return applyExtractSliceOfPadTensorSwapPattern(getOperation());
5198faf35c0SMatthias Springer   if (testTilePattern)
52041574554SRiver Riddle     return applyTilePattern(getOperation(), loopType, tileSizes, peeledLoops,
521a8f69be6STobias Gysi                             /*scalarizeDynamicDims=*/false);
522fb1def9cSMatthias Springer   if (testTileScalarizeDynamicDims)
52341574554SRiver Riddle     return applyTilePattern(getOperation(), loopType, tileSizes,
524fb1def9cSMatthias Springer                             /*peeledLoops=*/{}, /*scalarizeDynamicDims=*/true);
52533d2a780SThomas Raoux   if (testSplitReduction)
52633d2a780SThomas Raoux     return applySplitReduction(getOperation());
52765bdeddbSOkwan Kwon   if (testBubbleUpExtractSliceOpPattern)
52865bdeddbSOkwan Kwon     return applyBubbleUpExtractSliceOpPattern(getOperation());
5293fef2d26SRiver Riddle }
5303fef2d26SRiver Riddle 
5313fef2d26SRiver Riddle namespace mlir {
5323fef2d26SRiver Riddle namespace test {
registerTestLinalgTransforms()5333fef2d26SRiver Riddle void registerTestLinalgTransforms() {
534b5e22e6dSMehdi Amini   PassRegistration<TestLinalgTransforms>();
5353fef2d26SRiver Riddle }
5363fef2d26SRiver Riddle } // namespace test
5373fef2d26SRiver Riddle } // namespace mlir
538