1//===- TestLinalgFusionTransforms.cpp - Test Linalg fusion 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 fusion patterns.
10//
11//===----------------------------------------------------------------------===//
12
13#include "mlir/Dialect/Affine/IR/AffineOps.h"
14#include "mlir/Dialect/Func/IR/FuncOps.h"
15#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
16#include "mlir/Dialect/SCF/Transforms/Patterns.h"
17#include "mlir/Pass/Pass.h"
18#include "mlir/Pass/PassManager.h"
19#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
20#include "mlir/Transforms/Passes.h"
21
22using namespace mlir;
23using namespace mlir::linalg;
24
25static LogicalResult fuseLinalgOpsGreedily(func::FuncOp f) {
26 OpBuilder b(f);
27
28 // Save original Linalg ops, we only want to make a pass over those.
29 SmallVector<LinalgOp, 8> linalgOps;
30 f.walk([&](LinalgOp op) {
31 // TODO: support multi-results.
32 if (op->getNumResults() <= 1)
33 linalgOps.push_back(op);
34 });
35
36 // Tile and Fuse for tensors inputs (TODO: all tensor operands).
37 bool changed = false;
38 for (LinalgOp linalgOp : llvm::reverse(linalgOps)) {
39 for (OpOperand &opOperand : linalgOp->getOpOperands()) {
40 if (isa<MemRefType>(opOperand.get().getType()))
41 continue;
42 if (isa<RankedTensorType>(opOperand.get().getType())) {
43 // Tile and Fuse tensor input.
44 if (opOperand.getOperandNumber() >= linalgOp.getNumDpsInputs())
45 continue;
46 auto info = fuseProducerOfTensor(b, opOperand);
47 if (failed(info))
48 continue;
49 auto *originalOp = info->originalProducer.getOperation();
50 auto *originalOpInLinalgOpsVector = llvm::find(linalgOps, originalOp);
51 *originalOpInLinalgOpsVector = info->fusedProducer;
52 // Don't mark for erasure in the tensor case, let DCE handle this.
53 changed = true;
54 }
55 }
56 }
57
58 return changed ? success() : failure();
59}
60
61namespace {
62struct TestLinalgGreedyFusion
63 : public PassWrapper<TestLinalgGreedyFusion, OperationPass<func::FuncOp>> {
64 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestLinalgGreedyFusion)
65
66 void getDependentDialects(DialectRegistry &registry) const override {
67 registry.insert<affine::AffineDialect, linalg::LinalgDialect,
68 memref::MemRefDialect, scf::SCFDialect>();
69 }
70 StringRef getArgument() const final { return "test-linalg-greedy-fusion"; }
71 StringRef getDescription() const final {
72 return "Test Linalg fusion by applying a greedy test transformation.";
73 }
74 void runOnOperation() override {
75 MLIRContext *context = &getContext();
76 RewritePatternSet patterns =
77 linalg::getLinalgTilingCanonicalizationPatterns(ctx: context);
78 patterns.add<ExtractSliceOfPadTensorSwapPattern>(arg&: context);
79 scf::populateSCFForLoopCanonicalizationPatterns(patterns);
80 FrozenRewritePatternSet frozenPatterns(std::move(patterns));
81 OpPassManager pm(func::FuncOp::getOperationName());
82 pm.addPass(createLoopInvariantCodeMotionPass());
83 pm.addPass(createCanonicalizerPass());
84 pm.addPass(createCSEPass());
85 do {
86 (void)applyPatternsGreedily(getOperation(), frozenPatterns);
87 if (failed(runPipeline(pipeline&: pm, op: getOperation())))
88 this->signalPassFailure();
89 } while (succeeded(fuseLinalgOpsGreedily(getOperation())));
90 }
91};
92} // namespace
93
94namespace mlir {
95namespace test {
96void registerTestLinalgGreedyFusion() {
97 PassRegistration<TestLinalgGreedyFusion>();
98}
99
100} // namespace test
101} // namespace mlir
102

source code of mlir/test/lib/Dialect/Linalg/TestLinalgFusionTransforms.cpp