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 ®istry) 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