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