1 //===- TosaToSCF.cpp - Lowering Tosa to SCF Dialect -----------------------===//
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 // These rewriters lower from the Tosa to the SCF dialect.
10 //
11 //===----------------------------------------------------------------------===//
12
13 #include "mlir/Conversion/TosaToSCF/TosaToSCF.h"
14 #include "mlir/Dialect/SCF/IR/SCF.h"
15 #include "mlir/Dialect/Tensor/IR/Tensor.h"
16 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
17 #include "mlir/IR/BlockAndValueMapping.h"
18 #include "mlir/IR/PatternMatch.h"
19 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
20
21 using namespace mlir;
22 using namespace tosa;
23
inlineIfCase(Region & srcRegion,Region & dstRegion,OperandRange operands,PatternRewriter & rewriter)24 static void inlineIfCase(Region &srcRegion, Region &dstRegion,
25 OperandRange operands, PatternRewriter &rewriter) {
26 rewriter.cloneRegionBefore(srcRegion, &dstRegion.front());
27 rewriter.eraseBlock(&dstRegion.back());
28
29 Block *headBlock = &dstRegion.front();
30 for (auto it : llvm::zip(headBlock->getArguments(), operands))
31 std::get<0>(it).replaceAllUsesWith(std::get<1>(it));
32
33 auto yield = cast<YieldOp>(headBlock->getTerminator());
34 rewriter.setInsertionPoint(yield);
35 rewriter.create<scf::YieldOp>(yield.getLoc(), yield.getInputs());
36 rewriter.eraseOp(yield);
37
38 headBlock->eraseArguments(
39 llvm::to_vector<4>(llvm::seq<unsigned>(0, headBlock->getNumArguments())));
40 }
41
inlineWhileCase(Region & srcRegion,Region & dstRegion,PatternRewriter & rewriter,bool isCond)42 static void inlineWhileCase(Region &srcRegion, Region &dstRegion,
43 PatternRewriter &rewriter, bool isCond) {
44 rewriter.cloneRegionBefore(srcRegion, &dstRegion.back());
45 rewriter.eraseBlock(&dstRegion.back());
46
47 Block *headBlock = &dstRegion.front();
48
49 auto yield = cast<YieldOp>(headBlock->getTerminator());
50 rewriter.setInsertionPoint(yield);
51 if (isCond) {
52 auto condition =
53 rewriter.create<tensor::ExtractOp>(yield.getLoc(), yield.getOperand(0));
54 rewriter.create<scf::ConditionOp>(yield.getLoc(), condition,
55 headBlock->getArguments());
56 } else {
57 rewriter.setInsertionPoint(yield);
58 rewriter.create<scf::YieldOp>(yield.getLoc(), yield.getInputs());
59 }
60 rewriter.eraseOp(yield);
61 }
62
63 namespace {
64
65 class IfOpConverter : public OpRewritePattern<tosa::IfOp> {
66 public:
67 using OpRewritePattern<tosa::IfOp>::OpRewritePattern;
68
matchAndRewrite(tosa::IfOp op,PatternRewriter & rewriter) const69 LogicalResult matchAndRewrite(tosa::IfOp op,
70 PatternRewriter &rewriter) const final {
71 auto condition =
72 rewriter.create<tensor::ExtractOp>(op.getLoc(), op.getCond());
73 auto newIf = rewriter.create<scf::IfOp>(op.getLoc(), op.getResultTypes(),
74 condition, true);
75
76 inlineIfCase(op.getThenBranch(), newIf.getThenRegion(), op.getInputs(),
77 rewriter);
78 inlineIfCase(op.getElseBranch(), newIf.getElseRegion(), op.getInputs(),
79 rewriter);
80
81 rewriter.replaceOp(op, newIf.getResults());
82 return success();
83 }
84 };
85
86 class WhileOpConverter : public OpRewritePattern<tosa::WhileOp> {
87 public:
88 using OpRewritePattern<tosa::WhileOp>::OpRewritePattern;
89
matchAndRewrite(tosa::WhileOp op,PatternRewriter & rewriter) const90 LogicalResult matchAndRewrite(tosa::WhileOp op,
91 PatternRewriter &rewriter) const final {
92 auto newWhile = rewriter.create<scf::WhileOp>(
93 op.getLoc(), op.getResultTypes(), op.getInputs());
94 rewriter.createBlock(&newWhile.getBefore());
95 rewriter.createBlock(&newWhile.getAfter());
96
97 inlineWhileCase(op.getCond(), newWhile.getBefore(), rewriter, true);
98 inlineWhileCase(op.getBody(), newWhile.getAfter(), rewriter, false);
99
100 rewriter.replaceOp(op, newWhile.getResults());
101
102 return success();
103 }
104 };
105
106 } // namespace
107
populateTosaToSCFConversionPatterns(RewritePatternSet * patterns)108 void mlir::tosa::populateTosaToSCFConversionPatterns(
109 RewritePatternSet *patterns) {
110 patterns->add<IfOpConverter>(patterns->getContext());
111 patterns->add<WhileOpConverter>(patterns->getContext());
112 }
113