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