• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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