1 //===- SafeStackColoring.cpp - SafeStack frame coloring -------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "llvm/Analysis/StackLifetime.h"
10 #include "llvm/ADT/DepthFirstIterator.h"
11 #include "llvm/Config/llvm-config.h"
12 #include "llvm/IR/BasicBlock.h"
13 #include "llvm/IR/CFG.h"
14 #include "llvm/IR/Instructions.h"
15 #include "llvm/IR/IntrinsicInst.h"
16 #include "llvm/IR/Intrinsics.h"
17 #include "llvm/IR/User.h"
18 #include "llvm/Support/Casting.h"
19 #include "llvm/Support/CommandLine.h"
20 #include "llvm/Support/Compiler.h"
21 #include "llvm/Support/Debug.h"
22 #include <tuple>
23 
24 using namespace llvm;
25 using namespace llvm::safestack;
26 
27 #define DEBUG_TYPE "safestackcoloring"
28 
29 const StackColoring::LiveRange &
30 StackColoring::getLiveRange(const AllocaInst *AI) const {
31   const auto IT = AllocaNumbering.find(AI);
32   assert(IT != AllocaNumbering.end());
33   return LiveRanges[IT->second];
34 }
35 
36 static bool readMarker(const Instruction *I, bool *IsStart) {
37   if (!I->isLifetimeStartOrEnd())
38     return false;
39 
40   auto *II = cast<IntrinsicInst>(I);
41   *IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
42   return true;
43 }
44 
45 std::vector<const IntrinsicInst *> StackColoring::getMarkers() const {
46   std::vector<const IntrinsicInst *> Markers;
47   for (auto &M : InstructionNumbering)
48     if (M.getFirst()->isLifetimeStartOrEnd())
49       Markers.push_back(M.getFirst());
50   return Markers;
51 }
52 
53 void StackColoring::collectMarkers() {
54   InterestingAllocas.resize(NumAllocas);
55   DenseMap<const BasicBlock *, SmallDenseMap<const IntrinsicInst *, Marker>>
56       BBMarkerSet;
57 
58   // Compute the set of start/end markers per basic block.
59   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
60     const AllocaInst *AI = Allocas[AllocaNo];
61     SmallVector<const Instruction *, 8> WorkList;
62     WorkList.push_back(AI);
63     while (!WorkList.empty()) {
64       const Instruction *I = WorkList.pop_back_val();
65       for (const User *U : I->users()) {
66         if (auto *BI = dyn_cast<BitCastInst>(U)) {
67           WorkList.push_back(BI);
68           continue;
69         }
70         auto *UI = dyn_cast<IntrinsicInst>(U);
71         if (!UI)
72           continue;
73         bool IsStart;
74         if (!readMarker(UI, &IsStart))
75           continue;
76         if (IsStart)
77           InterestingAllocas.set(AllocaNo);
78         BBMarkerSet[UI->getParent()][UI] = {AllocaNo, IsStart};
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   LLVM_DEBUG(dbgs() << "Instructions:\n");
91   unsigned InstNo = 0;
92   for (const BasicBlock *BB : depth_first(&F)) {
93     LLVM_DEBUG(dbgs() << "  " << InstNo << ": BB " << BB->getName() << "\n");
94     unsigned BBStart = InstNo++;
95 
96     BlockLifetimeInfo &BlockInfo =
97         BlockLiveness.try_emplace(BB, NumAllocas).first->getSecond();
98 
99     auto &BlockMarkerSet = BBMarkerSet[BB];
100     if (BlockMarkerSet.empty()) {
101       unsigned BBEnd = InstNo;
102       BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
103       continue;
104     }
105 
106     auto ProcessMarker = [&](const IntrinsicInst *I, const Marker &M) {
107       LLVM_DEBUG(dbgs() << "  " << InstNo << ":  "
108                         << (M.IsStart ? "start " : "end   ") << M.AllocaNo
109                         << ", " << *I << "\n");
110 
111       BBMarkers[BB].push_back({InstNo, M});
112 
113       InstructionNumbering[I] = InstNo++;
114 
115       if (M.IsStart) {
116         BlockInfo.End.reset(M.AllocaNo);
117         BlockInfo.Begin.set(M.AllocaNo);
118       } else {
119         BlockInfo.Begin.reset(M.AllocaNo);
120         BlockInfo.End.set(M.AllocaNo);
121       }
122     };
123 
124     if (BlockMarkerSet.size() == 1) {
125       ProcessMarker(BlockMarkerSet.begin()->getFirst(),
126                     BlockMarkerSet.begin()->getSecond());
127     } else {
128       // Scan the BB to determine the marker order.
129       for (const Instruction &I : *BB) {
130         const IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
131         if (!II)
132           continue;
133         auto It = BlockMarkerSet.find(II);
134         if (It == BlockMarkerSet.end())
135           continue;
136         ProcessMarker(II, It->getSecond());
137       }
138     }
139 
140     unsigned BBEnd = InstNo;
141     BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
142   }
143   NumInst = InstNo;
144 }
145 
146 void StackColoring::calculateLocalLiveness() {
147   bool Changed = true;
148   while (Changed) {
149     Changed = false;
150 
151     for (const BasicBlock *BB : depth_first(&F)) {
152       BlockLifetimeInfo &BlockInfo = BlockLiveness.find(BB)->getSecond();
153 
154       // Compute LiveIn by unioning together the LiveOut sets of all preds.
155       BitVector LocalLiveIn;
156       for (auto *PredBB : predecessors(BB)) {
157         LivenessMap::const_iterator I = BlockLiveness.find(PredBB);
158         // If a predecessor is unreachable, ignore it.
159         if (I == BlockLiveness.end())
160           continue;
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 
190 void StackColoring::calculateLiveIntervals() {
191   for (auto IT : BlockLiveness) {
192     const 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) || Start[AllocaNo] == BBStart);
218         if (!Started.test(AllocaNo)) {
219           Started.set(AllocaNo);
220           Ended.reset(AllocaNo);
221           Start[AllocaNo] = InstNo;
222         }
223       } else {
224         assert(!Ended.test(AllocaNo));
225         if (Started.test(AllocaNo)) {
226           LiveRanges[AllocaNo].addRange(Start[AllocaNo], InstNo);
227           Started.reset(AllocaNo);
228         }
229         Ended.set(AllocaNo);
230       }
231     }
232 
233     for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
234       if (Started.test(AllocaNo))
235         LiveRanges[AllocaNo].addRange(Start[AllocaNo], BBEnd);
236   }
237 }
238 
239 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
240 LLVM_DUMP_METHOD void StackColoring::dumpAllocas() const {
241   dbgs() << "Allocas:\n";
242   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
243     dbgs() << "  " << AllocaNo << ": " << *Allocas[AllocaNo] << "\n";
244 }
245 
246 LLVM_DUMP_METHOD void StackColoring::dumpBlockLiveness() const {
247   dbgs() << "Block liveness:\n";
248   for (auto IT : BlockLiveness) {
249     const BasicBlock *BB = IT.getFirst();
250     const BlockLifetimeInfo &BlockInfo = BlockLiveness.find(BB)->getSecond();
251     auto BlockRange = BlockInstRange.find(BB)->getSecond();
252     dbgs() << "  BB [" << BlockRange.first << ", " << BlockRange.second
253            << "): begin " << BlockInfo.Begin << ", end " << BlockInfo.End
254            << ", livein " << BlockInfo.LiveIn << ", liveout "
255            << BlockInfo.LiveOut << "\n";
256   }
257 }
258 
259 LLVM_DUMP_METHOD void StackColoring::dumpLiveRanges() const {
260   dbgs() << "Alloca liveness:\n";
261   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
262     dbgs() << "  " << AllocaNo << ": " << LiveRanges[AllocaNo] << "\n";
263 }
264 #endif
265 
266 StackColoring::StackColoring(const Function &F,
267                              ArrayRef<const AllocaInst *> Allocas)
268     : F(F), Allocas(Allocas), NumAllocas(Allocas.size()) {
269   LLVM_DEBUG(dumpAllocas());
270 
271   for (unsigned I = 0; I < NumAllocas; ++I)
272     AllocaNumbering[Allocas[I]] = I;
273 
274   collectMarkers();
275 }
276 
277 void StackColoring::run() {
278   LiveRanges.resize(NumAllocas, LiveRange(NumInst));
279   for (unsigned I = 0; I < NumAllocas; ++I)
280     if (!InterestingAllocas.test(I))
281       LiveRanges[I] = getFullLiveRange();
282 
283   calculateLocalLiveness();
284   LLVM_DEBUG(dumpBlockLiveness());
285   calculateLiveIntervals();
286   LLVM_DEBUG(dumpLiveRanges());
287 }
288