1 //===- SafeStackColoring.cpp - SafeStack frame coloring -------------------===//
2 //
3 // The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9
10 #include "SafeStackColoring.h"
11 #include "llvm/ADT/BitVector.h"
12 #include "llvm/ADT/DenseMap.h"
13 #include "llvm/ADT/DepthFirstIterator.h"
14 #include "llvm/ADT/SmallVector.h"
15 #include "llvm/Config/llvm-config.h"
16 #include "llvm/IR/BasicBlock.h"
17 #include "llvm/IR/CFG.h"
18 #include "llvm/IR/Instruction.h"
19 #include "llvm/IR/Instructions.h"
20 #include "llvm/IR/IntrinsicInst.h"
21 #include "llvm/IR/Intrinsics.h"
22 #include "llvm/IR/User.h"
23 #include "llvm/Support/Casting.h"
24 #include "llvm/Support/CommandLine.h"
25 #include "llvm/Support/Compiler.h"
26 #include "llvm/Support/Debug.h"
27 #include "llvm/Support/raw_ostream.h"
28 #include <cassert>
29 #include <tuple>
30 #include <utility>
31
32 using namespace llvm;
33 using namespace llvm::safestack;
34
35 #define DEBUG_TYPE "safestackcoloring"
36
37 // Disabled by default due to PR32143.
38 static cl::opt<bool> ClColoring("safe-stack-coloring",
39 cl::desc("enable safe stack coloring"),
40 cl::Hidden, cl::init(false));
41
getLiveRange(AllocaInst * AI)42 const StackColoring::LiveRange &StackColoring::getLiveRange(AllocaInst *AI) {
43 const auto IT = AllocaNumbering.find(AI);
44 assert(IT != AllocaNumbering.end());
45 return LiveRanges[IT->second];
46 }
47
readMarker(Instruction * I,bool * IsStart)48 bool StackColoring::readMarker(Instruction *I, bool *IsStart) {
49 if (!I->isLifetimeStartOrEnd())
50 return false;
51
52 auto *II = cast<IntrinsicInst>(I);
53 *IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
54 return true;
55 }
56
removeAllMarkers()57 void StackColoring::removeAllMarkers() {
58 for (auto *I : Markers) {
59 auto *Op = dyn_cast<Instruction>(I->getOperand(1));
60 I->eraseFromParent();
61 // Remove the operand bitcast, too, if it has no more uses left.
62 if (Op && Op->use_empty())
63 Op->eraseFromParent();
64 }
65 }
66
collectMarkers()67 void StackColoring::collectMarkers() {
68 InterestingAllocas.resize(NumAllocas);
69 DenseMap<BasicBlock *, SmallDenseMap<Instruction *, Marker>> BBMarkerSet;
70
71 // Compute the set of start/end markers per basic block.
72 for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
73 AllocaInst *AI = Allocas[AllocaNo];
74 SmallVector<Instruction *, 8> WorkList;
75 WorkList.push_back(AI);
76 while (!WorkList.empty()) {
77 Instruction *I = WorkList.pop_back_val();
78 for (User *U : I->users()) {
79 if (auto *BI = dyn_cast<BitCastInst>(U)) {
80 WorkList.push_back(BI);
81 continue;
82 }
83 auto *UI = dyn_cast<Instruction>(U);
84 if (!UI)
85 continue;
86 bool IsStart;
87 if (!readMarker(UI, &IsStart))
88 continue;
89 if (IsStart)
90 InterestingAllocas.set(AllocaNo);
91 BBMarkerSet[UI->getParent()][UI] = {AllocaNo, IsStart};
92 Markers.push_back(UI);
93 }
94 }
95 }
96
97 // Compute instruction numbering. Only the following instructions are
98 // considered:
99 // * Basic block entries
100 // * Lifetime markers
101 // For each basic block, compute
102 // * the list of markers in the instruction order
103 // * the sets of allocas whose lifetime starts or ends in this BB
104 LLVM_DEBUG(dbgs() << "Instructions:\n");
105 unsigned InstNo = 0;
106 for (BasicBlock *BB : depth_first(&F)) {
107 LLVM_DEBUG(dbgs() << " " << InstNo << ": BB " << BB->getName() << "\n");
108 unsigned BBStart = InstNo++;
109
110 BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
111 BlockInfo.Begin.resize(NumAllocas);
112 BlockInfo.End.resize(NumAllocas);
113 BlockInfo.LiveIn.resize(NumAllocas);
114 BlockInfo.LiveOut.resize(NumAllocas);
115
116 auto &BlockMarkerSet = BBMarkerSet[BB];
117 if (BlockMarkerSet.empty()) {
118 unsigned BBEnd = InstNo;
119 BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
120 continue;
121 }
122
123 auto ProcessMarker = [&](Instruction *I, const Marker &M) {
124 LLVM_DEBUG(dbgs() << " " << InstNo << ": "
125 << (M.IsStart ? "start " : "end ") << M.AllocaNo
126 << ", " << *I << "\n");
127
128 BBMarkers[BB].push_back({InstNo, M});
129
130 InstructionNumbering[I] = InstNo++;
131
132 if (M.IsStart) {
133 if (BlockInfo.End.test(M.AllocaNo))
134 BlockInfo.End.reset(M.AllocaNo);
135 BlockInfo.Begin.set(M.AllocaNo);
136 } else {
137 if (BlockInfo.Begin.test(M.AllocaNo))
138 BlockInfo.Begin.reset(M.AllocaNo);
139 BlockInfo.End.set(M.AllocaNo);
140 }
141 };
142
143 if (BlockMarkerSet.size() == 1) {
144 ProcessMarker(BlockMarkerSet.begin()->getFirst(),
145 BlockMarkerSet.begin()->getSecond());
146 } else {
147 // Scan the BB to determine the marker order.
148 for (Instruction &I : *BB) {
149 auto It = BlockMarkerSet.find(&I);
150 if (It == BlockMarkerSet.end())
151 continue;
152 ProcessMarker(&I, It->getSecond());
153 }
154 }
155
156 unsigned BBEnd = InstNo;
157 BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
158 }
159 NumInst = InstNo;
160 }
161
calculateLocalLiveness()162 void StackColoring::calculateLocalLiveness() {
163 bool changed = true;
164 while (changed) {
165 changed = false;
166
167 for (BasicBlock *BB : depth_first(&F)) {
168 BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
169
170 // Compute LiveIn by unioning together the LiveOut sets of all preds.
171 BitVector LocalLiveIn;
172 for (auto *PredBB : predecessors(BB)) {
173 LivenessMap::const_iterator I = BlockLiveness.find(PredBB);
174 // If a predecessor is unreachable, ignore it.
175 if (I == BlockLiveness.end())
176 continue;
177 LocalLiveIn |= I->second.LiveOut;
178 }
179
180 // Compute LiveOut by subtracting out lifetimes that end in this
181 // block, then adding in lifetimes that begin in this block. If
182 // we have both BEGIN and END markers in the same basic block
183 // then we know that the BEGIN marker comes after the END,
184 // because we already handle the case where the BEGIN comes
185 // before the END when collecting the markers (and building the
186 // BEGIN/END vectors).
187 BitVector LocalLiveOut = LocalLiveIn;
188 LocalLiveOut.reset(BlockInfo.End);
189 LocalLiveOut |= BlockInfo.Begin;
190
191 // Update block LiveIn set, noting whether it has changed.
192 if (LocalLiveIn.test(BlockInfo.LiveIn)) {
193 changed = true;
194 BlockInfo.LiveIn |= LocalLiveIn;
195 }
196
197 // Update block LiveOut set, noting whether it has changed.
198 if (LocalLiveOut.test(BlockInfo.LiveOut)) {
199 changed = true;
200 BlockInfo.LiveOut |= LocalLiveOut;
201 }
202 }
203 } // while changed.
204 }
205
calculateLiveIntervals()206 void StackColoring::calculateLiveIntervals() {
207 for (auto IT : BlockLiveness) {
208 BasicBlock *BB = IT.getFirst();
209 BlockLifetimeInfo &BlockInfo = IT.getSecond();
210 unsigned BBStart, BBEnd;
211 std::tie(BBStart, BBEnd) = BlockInstRange[BB];
212
213 BitVector Started, Ended;
214 Started.resize(NumAllocas);
215 Ended.resize(NumAllocas);
216 SmallVector<unsigned, 8> Start;
217 Start.resize(NumAllocas);
218
219 // LiveIn ranges start at the first instruction.
220 for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
221 if (BlockInfo.LiveIn.test(AllocaNo)) {
222 Started.set(AllocaNo);
223 Start[AllocaNo] = BBStart;
224 }
225 }
226
227 for (auto &It : BBMarkers[BB]) {
228 unsigned InstNo = It.first;
229 bool IsStart = It.second.IsStart;
230 unsigned AllocaNo = It.second.AllocaNo;
231
232 if (IsStart) {
233 assert(!Started.test(AllocaNo) || Start[AllocaNo] == BBStart);
234 if (!Started.test(AllocaNo)) {
235 Started.set(AllocaNo);
236 Ended.reset(AllocaNo);
237 Start[AllocaNo] = InstNo;
238 }
239 } else {
240 assert(!Ended.test(AllocaNo));
241 if (Started.test(AllocaNo)) {
242 LiveRanges[AllocaNo].AddRange(Start[AllocaNo], InstNo);
243 Started.reset(AllocaNo);
244 }
245 Ended.set(AllocaNo);
246 }
247 }
248
249 for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
250 if (Started.test(AllocaNo))
251 LiveRanges[AllocaNo].AddRange(Start[AllocaNo], BBEnd);
252 }
253 }
254
255 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
dumpAllocas()256 LLVM_DUMP_METHOD void StackColoring::dumpAllocas() {
257 dbgs() << "Allocas:\n";
258 for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
259 dbgs() << " " << AllocaNo << ": " << *Allocas[AllocaNo] << "\n";
260 }
261
dumpBlockLiveness()262 LLVM_DUMP_METHOD void StackColoring::dumpBlockLiveness() {
263 dbgs() << "Block liveness:\n";
264 for (auto IT : BlockLiveness) {
265 BasicBlock *BB = IT.getFirst();
266 BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
267 auto BlockRange = BlockInstRange[BB];
268 dbgs() << " BB [" << BlockRange.first << ", " << BlockRange.second
269 << "): begin " << BlockInfo.Begin << ", end " << BlockInfo.End
270 << ", livein " << BlockInfo.LiveIn << ", liveout "
271 << BlockInfo.LiveOut << "\n";
272 }
273 }
274
dumpLiveRanges()275 LLVM_DUMP_METHOD void StackColoring::dumpLiveRanges() {
276 dbgs() << "Alloca liveness:\n";
277 for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
278 LiveRange &Range = LiveRanges[AllocaNo];
279 dbgs() << " " << AllocaNo << ": " << Range << "\n";
280 }
281 }
282 #endif
283
run()284 void StackColoring::run() {
285 LLVM_DEBUG(dumpAllocas());
286
287 for (unsigned I = 0; I < NumAllocas; ++I)
288 AllocaNumbering[Allocas[I]] = I;
289 LiveRanges.resize(NumAllocas);
290
291 collectMarkers();
292
293 if (!ClColoring) {
294 for (auto &R : LiveRanges) {
295 R.SetMaximum(1);
296 R.AddRange(0, 1);
297 }
298 return;
299 }
300
301 for (auto &R : LiveRanges)
302 R.SetMaximum(NumInst);
303 for (unsigned I = 0; I < NumAllocas; ++I)
304 if (!InterestingAllocas.test(I))
305 LiveRanges[I] = getFullLiveRange();
306
307 calculateLocalLiveness();
308 LLVM_DEBUG(dumpBlockLiveness());
309 calculateLiveIntervals();
310 LLVM_DEBUG(dumpLiveRanges());
311 }
312