1 //===- StackLifetime.cpp - Alloca Lifetime Analysis -----------------------===//
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/ADT/STLExtras.h"
12 #include "llvm/ADT/SmallVector.h"
13 #include "llvm/ADT/StringExtras.h"
14 #include "llvm/Analysis/ValueTracking.h"
15 #include "llvm/Config/llvm-config.h"
16 #include "llvm/IR/AssemblyAnnotationWriter.h"
17 #include "llvm/IR/BasicBlock.h"
18 #include "llvm/IR/CFG.h"
19 #include "llvm/IR/InstIterator.h"
20 #include "llvm/IR/Instructions.h"
21 #include "llvm/IR/IntrinsicInst.h"
22 #include "llvm/IR/Value.h"
23 #include "llvm/Support/Casting.h"
24 #include "llvm/Support/Compiler.h"
25 #include "llvm/Support/Debug.h"
26 #include "llvm/Support/FormattedStream.h"
27 #include <algorithm>
28 #include <tuple>
29 
30 using namespace llvm;
31 
32 #define DEBUG_TYPE "stack-lifetime"
33 
34 const StackLifetime::LiveRange &
35 StackLifetime::getLiveRange(const AllocaInst *AI) const {
36   const auto IT = AllocaNumbering.find(AI);
37   assert(IT != AllocaNumbering.end());
38   return LiveRanges[IT->second];
39 }
40 
41 bool StackLifetime::isReachable(const Instruction *I) const {
42   return BlockInstRange.find(I->getParent()) != BlockInstRange.end();
43 }
44 
45 bool StackLifetime::isAliveAfter(const AllocaInst *AI,
46                                  const Instruction *I) const {
47   const BasicBlock *BB = I->getParent();
48   auto ItBB = BlockInstRange.find(BB);
49   assert(ItBB != BlockInstRange.end() && "Unreachable is not expected");
50 
51   // Search the block for the first instruction following 'I'.
52   auto It = std::upper_bound(Instructions.begin() + ItBB->getSecond().first + 1,
53                              Instructions.begin() + ItBB->getSecond().second, I,
54                              [](const Instruction *L, const Instruction *R) {
55                                return L->comesBefore(R);
56                              });
57   --It;
58   unsigned InstNum = It - Instructions.begin();
59   return getLiveRange(AI).test(InstNum);
60 }
61 
62 // Returns unique alloca annotated by lifetime marker only if
63 // markers has the same size and points to the alloca start.
64 static const AllocaInst *findMatchingAlloca(const IntrinsicInst &II,
65                                             const DataLayout &DL) {
66   const AllocaInst *AI = findAllocaForValue(II.getArgOperand(1), true);
67   if (!AI)
68     return nullptr;
69 
70   auto AllocaSizeInBits = AI->getAllocationSizeInBits(DL);
71   if (!AllocaSizeInBits)
72     return nullptr;
73   int64_t AllocaSize = *AllocaSizeInBits / 8;
74 
75   auto *Size = dyn_cast<ConstantInt>(II.getArgOperand(0));
76   if (!Size)
77     return nullptr;
78   int64_t LifetimeSize = Size->getSExtValue();
79 
80   if (LifetimeSize != -1 && LifetimeSize != AllocaSize)
81     return nullptr;
82 
83   return AI;
84 }
85 
86 void StackLifetime::collectMarkers() {
87   InterestingAllocas.resize(NumAllocas);
88   DenseMap<const BasicBlock *, SmallDenseMap<const IntrinsicInst *, Marker>>
89       BBMarkerSet;
90 
91   const DataLayout &DL = F.getParent()->getDataLayout();
92 
93   // Compute the set of start/end markers per basic block.
94   for (const BasicBlock *BB : depth_first(&F)) {
95     for (const Instruction &I : *BB) {
96       const IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
97       if (!II || !II->isLifetimeStartOrEnd())
98         continue;
99       const AllocaInst *AI = findMatchingAlloca(*II, DL);
100       if (!AI) {
101         HasUnknownLifetimeStartOrEnd = true;
102         continue;
103       }
104       auto It = AllocaNumbering.find(AI);
105       if (It == AllocaNumbering.end())
106         continue;
107       auto AllocaNo = It->second;
108       bool IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
109       if (IsStart)
110         InterestingAllocas.set(AllocaNo);
111       BBMarkerSet[BB][II] = {AllocaNo, IsStart};
112     }
113   }
114 
115   // Compute instruction numbering. Only the following instructions are
116   // considered:
117   // * Basic block entries
118   // * Lifetime markers
119   // For each basic block, compute
120   // * the list of markers in the instruction order
121   // * the sets of allocas whose lifetime starts or ends in this BB
122   LLVM_DEBUG(dbgs() << "Instructions:\n");
123   for (const BasicBlock *BB : depth_first(&F)) {
124     LLVM_DEBUG(dbgs() << "  " << Instructions.size() << ": BB " << BB->getName()
125                       << "\n");
126     auto BBStart = Instructions.size();
127     Instructions.push_back(nullptr);
128 
129     BlockLifetimeInfo &BlockInfo =
130         BlockLiveness.try_emplace(BB, NumAllocas).first->getSecond();
131 
132     auto &BlockMarkerSet = BBMarkerSet[BB];
133     if (BlockMarkerSet.empty()) {
134       BlockInstRange[BB] = std::make_pair(BBStart, Instructions.size());
135       continue;
136     }
137 
138     auto ProcessMarker = [&](const IntrinsicInst *I, const Marker &M) {
139       LLVM_DEBUG(dbgs() << "  " << Instructions.size() << ":  "
140                         << (M.IsStart ? "start " : "end   ") << M.AllocaNo
141                         << ", " << *I << "\n");
142 
143       BBMarkers[BB].push_back({Instructions.size(), M});
144       Instructions.push_back(I);
145 
146       if (M.IsStart) {
147         BlockInfo.End.reset(M.AllocaNo);
148         BlockInfo.Begin.set(M.AllocaNo);
149       } else {
150         BlockInfo.Begin.reset(M.AllocaNo);
151         BlockInfo.End.set(M.AllocaNo);
152       }
153     };
154 
155     if (BlockMarkerSet.size() == 1) {
156       ProcessMarker(BlockMarkerSet.begin()->getFirst(),
157                     BlockMarkerSet.begin()->getSecond());
158     } else {
159       // Scan the BB to determine the marker order.
160       for (const Instruction &I : *BB) {
161         const IntrinsicInst *II = dyn_cast<IntrinsicInst>(&I);
162         if (!II)
163           continue;
164         auto It = BlockMarkerSet.find(II);
165         if (It == BlockMarkerSet.end())
166           continue;
167         ProcessMarker(II, It->getSecond());
168       }
169     }
170 
171     BlockInstRange[BB] = std::make_pair(BBStart, Instructions.size());
172   }
173 }
174 
175 void StackLifetime::calculateLocalLiveness() {
176   bool Changed = true;
177   while (Changed) {
178     Changed = false;
179 
180     for (const BasicBlock *BB : depth_first(&F)) {
181       BlockLifetimeInfo &BlockInfo = BlockLiveness.find(BB)->getSecond();
182 
183       // Compute LiveIn by unioning together the LiveOut sets of all preds.
184       BitVector LocalLiveIn;
185       for (const auto *PredBB : predecessors(BB)) {
186         LivenessMap::const_iterator I = BlockLiveness.find(PredBB);
187         // If a predecessor is unreachable, ignore it.
188         if (I == BlockLiveness.end())
189           continue;
190         switch (Type) {
191         case LivenessType::May:
192           LocalLiveIn |= I->second.LiveOut;
193           break;
194         case LivenessType::Must:
195           if (LocalLiveIn.empty())
196             LocalLiveIn = I->second.LiveOut;
197           else
198             LocalLiveIn &= I->second.LiveOut;
199           break;
200         }
201       }
202 
203       // Compute LiveOut by subtracting out lifetimes that end in this
204       // block, then adding in lifetimes that begin in this block.  If
205       // we have both BEGIN and END markers in the same basic block
206       // then we know that the BEGIN marker comes after the END,
207       // because we already handle the case where the BEGIN comes
208       // before the END when collecting the markers (and building the
209       // BEGIN/END vectors).
210       BitVector LocalLiveOut = LocalLiveIn;
211       LocalLiveOut.reset(BlockInfo.End);
212       LocalLiveOut |= BlockInfo.Begin;
213 
214       // Update block LiveIn set, noting whether it has changed.
215       if (LocalLiveIn.test(BlockInfo.LiveIn)) {
216         BlockInfo.LiveIn |= LocalLiveIn;
217       }
218 
219       // Update block LiveOut set, noting whether it has changed.
220       if (LocalLiveOut.test(BlockInfo.LiveOut)) {
221         Changed = true;
222         BlockInfo.LiveOut |= LocalLiveOut;
223       }
224     }
225   } // while changed.
226 }
227 
228 void StackLifetime::calculateLiveIntervals() {
229   for (auto IT : BlockLiveness) {
230     const BasicBlock *BB = IT.getFirst();
231     BlockLifetimeInfo &BlockInfo = IT.getSecond();
232     unsigned BBStart, BBEnd;
233     std::tie(BBStart, BBEnd) = BlockInstRange[BB];
234 
235     BitVector Started, Ended;
236     Started.resize(NumAllocas);
237     Ended.resize(NumAllocas);
238     SmallVector<unsigned, 8> Start;
239     Start.resize(NumAllocas);
240 
241     // LiveIn ranges start at the first instruction.
242     for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
243       if (BlockInfo.LiveIn.test(AllocaNo)) {
244         Started.set(AllocaNo);
245         Start[AllocaNo] = BBStart;
246       }
247     }
248 
249     for (auto &It : BBMarkers[BB]) {
250       unsigned InstNo = It.first;
251       bool IsStart = It.second.IsStart;
252       unsigned AllocaNo = It.second.AllocaNo;
253 
254       if (IsStart) {
255         if (!Started.test(AllocaNo)) {
256           Started.set(AllocaNo);
257           Ended.reset(AllocaNo);
258           Start[AllocaNo] = InstNo;
259         }
260       } else {
261         if (Started.test(AllocaNo)) {
262           LiveRanges[AllocaNo].addRange(Start[AllocaNo], InstNo);
263           Started.reset(AllocaNo);
264         }
265         Ended.set(AllocaNo);
266       }
267     }
268 
269     for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
270       if (Started.test(AllocaNo))
271         LiveRanges[AllocaNo].addRange(Start[AllocaNo], BBEnd);
272   }
273 }
274 
275 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
276 LLVM_DUMP_METHOD void StackLifetime::dumpAllocas() const {
277   dbgs() << "Allocas:\n";
278   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
279     dbgs() << "  " << AllocaNo << ": " << *Allocas[AllocaNo] << "\n";
280 }
281 
282 LLVM_DUMP_METHOD void StackLifetime::dumpBlockLiveness() const {
283   dbgs() << "Block liveness:\n";
284   for (auto IT : BlockLiveness) {
285     const BasicBlock *BB = IT.getFirst();
286     const BlockLifetimeInfo &BlockInfo = BlockLiveness.find(BB)->getSecond();
287     auto BlockRange = BlockInstRange.find(BB)->getSecond();
288     dbgs() << "  BB (" << BB->getName() << ") [" << BlockRange.first << ", " << BlockRange.second
289            << "): begin " << BlockInfo.Begin << ", end " << BlockInfo.End
290            << ", livein " << BlockInfo.LiveIn << ", liveout "
291            << BlockInfo.LiveOut << "\n";
292   }
293 }
294 
295 LLVM_DUMP_METHOD void StackLifetime::dumpLiveRanges() const {
296   dbgs() << "Alloca liveness:\n";
297   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
298     dbgs() << "  " << AllocaNo << ": " << LiveRanges[AllocaNo] << "\n";
299 }
300 #endif
301 
302 StackLifetime::StackLifetime(const Function &F,
303                              ArrayRef<const AllocaInst *> Allocas,
304                              LivenessType Type)
305     : F(F), Type(Type), Allocas(Allocas), NumAllocas(Allocas.size()) {
306   LLVM_DEBUG(dumpAllocas());
307 
308   for (unsigned I = 0; I < NumAllocas; ++I)
309     AllocaNumbering[Allocas[I]] = I;
310 
311   collectMarkers();
312 }
313 
314 void StackLifetime::run() {
315   if (HasUnknownLifetimeStartOrEnd) {
316     // There is marker which we can't assign to a specific alloca, so we
317     // fallback to the most conservative results for the type.
318     switch (Type) {
319     case LivenessType::May:
320       LiveRanges.resize(NumAllocas, getFullLiveRange());
321       break;
322     case LivenessType::Must:
323       LiveRanges.resize(NumAllocas, LiveRange(Instructions.size()));
324       break;
325     }
326     return;
327   }
328 
329   LiveRanges.resize(NumAllocas, LiveRange(Instructions.size()));
330   for (unsigned I = 0; I < NumAllocas; ++I)
331     if (!InterestingAllocas.test(I))
332       LiveRanges[I] = getFullLiveRange();
333 
334   calculateLocalLiveness();
335   LLVM_DEBUG(dumpBlockLiveness());
336   calculateLiveIntervals();
337   LLVM_DEBUG(dumpLiveRanges());
338 }
339 
340 class StackLifetime::LifetimeAnnotationWriter
341     : public AssemblyAnnotationWriter {
342   const StackLifetime &SL;
343 
344   void printInstrAlive(unsigned InstrNo, formatted_raw_ostream &OS) {
345     SmallVector<StringRef, 16> Names;
346     for (const auto &KV : SL.AllocaNumbering) {
347       if (SL.LiveRanges[KV.getSecond()].test(InstrNo))
348         Names.push_back(KV.getFirst()->getName());
349     }
350     llvm::sort(Names);
351     OS << "  ; Alive: <" << llvm::join(Names, " ") << ">\n";
352   }
353 
354   void emitBasicBlockStartAnnot(const BasicBlock *BB,
355                                 formatted_raw_ostream &OS) override {
356     auto ItBB = SL.BlockInstRange.find(BB);
357     if (ItBB == SL.BlockInstRange.end())
358       return; // Unreachable.
359     printInstrAlive(ItBB->getSecond().first, OS);
360   }
361 
362   void printInfoComment(const Value &V, formatted_raw_ostream &OS) override {
363     const Instruction *Instr = dyn_cast<Instruction>(&V);
364     if (!Instr || !SL.isReachable(Instr))
365       return;
366 
367     SmallVector<StringRef, 16> Names;
368     for (const auto &KV : SL.AllocaNumbering) {
369       if (SL.isAliveAfter(KV.getFirst(), Instr))
370         Names.push_back(KV.getFirst()->getName());
371     }
372     llvm::sort(Names);
373     OS << "\n  ; Alive: <" << llvm::join(Names, " ") << ">\n";
374   }
375 
376 public:
377   LifetimeAnnotationWriter(const StackLifetime &SL) : SL(SL) {}
378 };
379 
380 void StackLifetime::print(raw_ostream &OS) {
381   LifetimeAnnotationWriter AAW(*this);
382   F.print(OS, &AAW);
383 }
384 
385 PreservedAnalyses StackLifetimePrinterPass::run(Function &F,
386                                                 FunctionAnalysisManager &AM) {
387   SmallVector<const AllocaInst *, 8> Allocas;
388   for (auto &I : instructions(F))
389     if (const AllocaInst *AI = dyn_cast<AllocaInst>(&I))
390       Allocas.push_back(AI);
391   StackLifetime SL(F, Allocas, Type);
392   SL.run();
393   SL.print(OS);
394   return PreservedAnalyses::all();
395 }
396 
397 void StackLifetimePrinterPass::printPipeline(
398     raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
399   static_cast<PassInfoMixin<StackLifetimePrinterPass> *>(this)->printPipeline(
400       OS, MapClassName2PassName);
401   OS << "<";
402   switch (Type) {
403   case StackLifetime::LivenessType::May:
404     OS << "may";
405     break;
406   case StackLifetime::LivenessType::Must:
407     OS << "must";
408     break;
409   }
410   OS << ">";
411 }
412