1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
|
//===- LoopRangeFolding.cpp - Code to perform loop range folding-----------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements loop range folding.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/SCF/Transforms/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/IR/IRMapping.h"
namespace mlir {
#define GEN_PASS_DEF_SCFFORLOOPRANGEFOLDING
#include "mlir/Dialect/SCF/Transforms/Passes.h.inc"
} // namespace mlir
using namespace mlir;
using namespace mlir::scf;
namespace {
struct ForLoopRangeFolding
: public impl::SCFForLoopRangeFoldingBase<ForLoopRangeFolding> {
void runOnOperation() override;
};
} // namespace
void ForLoopRangeFolding::runOnOperation() {
getOperation()->walk([&](ForOp op) {
Value indVar = op.getInductionVar();
auto canBeFolded = [&](Value value) {
return op.isDefinedOutsideOfLoop(value) || value == indVar;
};
// Fold until a fixed point is reached
while (true) {
// If the induction variable is used more than once, we can't fold its
// arith ops into the loop range
if (!indVar.hasOneUse())
break;
Operation *user = *indVar.getUsers().begin();
if (!isa<arith::AddIOp, arith::MulIOp>(user))
break;
if (!llvm::all_of(user->getOperands(), canBeFolded))
break;
OpBuilder b(op);
IRMapping lbMap;
lbMap.map(indVar, op.getLowerBound());
IRMapping ubMap;
ubMap.map(indVar, op.getUpperBound());
IRMapping stepMap;
stepMap.map(indVar, op.getStep());
if (isa<arith::AddIOp>(user)) {
Operation *lbFold = b.clone(*user, lbMap);
Operation *ubFold = b.clone(*user, ubMap);
op.setLowerBound(lbFold->getResult(0));
op.setUpperBound(ubFold->getResult(0));
} else if (isa<arith::MulIOp>(user)) {
Operation *lbFold = b.clone(*user, lbMap);
Operation *ubFold = b.clone(*user, ubMap);
Operation *stepFold = b.clone(*user, stepMap);
op.setLowerBound(lbFold->getResult(0));
op.setUpperBound(ubFold->getResult(0));
op.setStep(stepFold->getResult(0));
}
ValueRange wrapIndvar(indVar);
user->replaceAllUsesWith(wrapIndvar);
user->erase();
}
});
}
std::unique_ptr<Pass> mlir::createForLoopRangeFoldingPass() {
return std::make_unique<ForLoopRangeFolding>();
}
|
