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 
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 
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 
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 
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 
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 
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)
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 
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 
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 
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