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
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
|
//===- LoopInvariantCodeMotion.cpp - Code to perform loop fusion-----------===//
//
// Copyright 2019 The MLIR Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// =============================================================================
//
// This file implements loop invariant code motion.
//
//===----------------------------------------------------------------------===//
#include <iomanip>
#include <sstream>
#include "mlir/AffineOps/AffineOps.h"
#include "mlir/Analysis/AffineAnalysis.h"
#include "mlir/Analysis/AffineStructures.h"
#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/Utils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
#include "mlir/Pass/Pass.h"
#include "mlir/StandardOps/Ops.h"
#include "mlir/Transforms/LoopUtils.h"
#include "mlir/Transforms/Passes.h"
#include "mlir/Transforms/Utils.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
#define DEBUG_TYPE "licm"
using llvm::SetVector;
using namespace mlir;
namespace {
/// Loop invariant code motion (LICM) pass.
/// TODO(asabne) : The pass is missing zero-trip tests.
/// TODO(asabne) : Check for the presence of side effects before hoisting.
struct LoopInvariantCodeMotion : public FunctionPass<LoopInvariantCodeMotion> {
void runOnFunction() override;
void runOnAffineForOp(AffineForOp forOp);
std::vector<AffineForOp> forOps;
};
} // end anonymous namespace
FunctionPassBase *mlir::createLoopInvariantCodeMotionPass() {
return new LoopInvariantCodeMotion();
}
void LoopInvariantCodeMotion::runOnAffineForOp(AffineForOp forOp) {
auto *loopBody = forOp.getBody();
// This is the place where hoisted instructions would reside.
FuncBuilder b(forOp.getOperation());
// This vector is used to place loop invariant operations.
SmallVector<Operation *, 8> opsToMove;
SetVector<Operation *> loopDefinedOps;
// Generate forward slice which contains ops that fall under the transitive
// definition closure following the loop induction variable.
getForwardSlice(forOp, &loopDefinedOps);
LLVM_DEBUG(for (auto i
: loopDefinedOps) {
(i->print(llvm::dbgs() << "\nLoop-dependent op\n"));
});
for (auto &op : *loopBody) {
// If the operation is loop invariant, insert it into opsToMove.
if (!op.isa<AffineForOp>() && !op.isa<AffineTerminatorOp>() &&
loopDefinedOps.count(&op) != 1) {
LLVM_DEBUG(op.print(llvm::dbgs() << "\nLICM'ing op\n"));
opsToMove.push_back(&op);
}
}
// For all instructions that we found to be invariant, place them sequentially
// right before the for loop.
for (auto *op : opsToMove) {
op->moveBefore(forOp);
}
LLVM_DEBUG(forOp.getOperation()->print(llvm::dbgs() << "\nModified loop\n"));
// If the for loop body has a single operation (the terminator), erase it.
if (forOp.getBody()->getOperations().size() == 1) {
assert(forOp.getBody()->getOperations().front().isa<AffineTerminatorOp>());
forOp.erase();
}
}
void LoopInvariantCodeMotion::runOnFunction() {
forOps.clear();
// Gather all loops in a function, and order them in innermost-loop-first
// order. This way, we first LICM from the inner loop, and place the ops in
// the outer loop, which in turn can be further LICM'ed. This saves iterating
// on the inner loop operations while LICMing through the outer loop.
getFunction().walk<AffineForOp>(
[&](AffineForOp forOp) { forOps.push_back(forOp); });
// We gather loops first, and then go over them later because we don't want to
// mess the iterators up.
for (auto op : forOps) {
LLVM_DEBUG(op.getOperation()->print(llvm::dbgs() << "\nOriginal loop\n"));
runOnAffineForOp(op);
}
}
static PassRegistration<LoopInvariantCodeMotion>
pass("loop-invariant-code-motion",
"Hoist loop invariant instructions outside of the loop");
|
