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 "SafeStackColoring.h"
10#include "llvm/ADT/BitVector.h"
11#include "llvm/ADT/DenseMap.h"
12#include "llvm/ADT/DepthFirstIterator.h"
13#include "llvm/ADT/SmallVector.h"
14#include "llvm/Config/llvm-config.h"
15#include "llvm/IR/BasicBlock.h"
16#include "llvm/IR/CFG.h"
17#include "llvm/IR/Instruction.h"
18#include "llvm/IR/Instructions.h"
19#include "llvm/IR/IntrinsicInst.h"
20#include "llvm/IR/Intrinsics.h"
21#include "llvm/IR/User.h"
22#include "llvm/Support/Casting.h"
23#include "llvm/Support/CommandLine.h"
24#include "llvm/Support/Compiler.h"
25#include "llvm/Support/Debug.h"
26#include "llvm/Support/raw_ostream.h"
27#include <cassert>
28#include <tuple>
29#include <utility>
30
31using namespace llvm;
32using namespace llvm::safestack;
33
34#define DEBUG_TYPE "safestackcoloring"
35
36// Disabled by default due to PR32143.
37static cl::opt<bool> ClColoring("safe-stack-coloring",
38                                cl::desc("enable safe stack coloring"),
39                                cl::Hidden, cl::init(false));
40
41const StackColoring::LiveRange &StackColoring::getLiveRange(AllocaInst *AI) {
42  const auto IT = AllocaNumbering.find(AI);
43  assert(IT != AllocaNumbering.end());
44  return LiveRanges[IT->second];
45}
46
47bool StackColoring::readMarker(Instruction *I, bool *IsStart) {
48  if (!I->isLifetimeStartOrEnd())
49    return false;
50
51  auto *II = cast<IntrinsicInst>(I);
52  *IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
53  return true;
54}
55
56void StackColoring::removeAllMarkers() {
57  for (auto *I : Markers) {
58    auto *Op = dyn_cast<Instruction>(I->getOperand(1));
59    I->eraseFromParent();
60    // Remove the operand bitcast, too, if it has no more uses left.
61    if (Op && Op->use_empty())
62      Op->eraseFromParent();
63  }
64}
65
66void StackColoring::collectMarkers() {
67  InterestingAllocas.resize(NumAllocas);
68  DenseMap<BasicBlock *, SmallDenseMap<Instruction *, Marker>> BBMarkerSet;
69
70  // Compute the set of start/end markers per basic block.
71  for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
72    AllocaInst *AI = Allocas[AllocaNo];
73    SmallVector<Instruction *, 8> WorkList;
74    WorkList.push_back(AI);
75    while (!WorkList.empty()) {
76      Instruction *I = WorkList.pop_back_val();
77      for (User *U : I->users()) {
78        if (auto *BI = dyn_cast<BitCastInst>(U)) {
79          WorkList.push_back(BI);
80          continue;
81        }
82        auto *UI = dyn_cast<Instruction>(U);
83        if (!UI)
84          continue;
85        bool IsStart;
86        if (!readMarker(UI, &IsStart))
87          continue;
88        if (IsStart)
89          InterestingAllocas.set(AllocaNo);
90        BBMarkerSet[UI->getParent()][UI] = {AllocaNo, IsStart};
91        Markers.push_back(UI);
92      }
93    }
94  }
95
96  // Compute instruction numbering. Only the following instructions are
97  // considered:
98  // * Basic block entries
99  // * Lifetime markers
100  // For each basic block, compute
101  // * the list of markers in the instruction order
102  // * the sets of allocas whose lifetime starts or ends in this BB
103  LLVM_DEBUG(dbgs() << "Instructions:\n");
104  unsigned InstNo = 0;
105  for (BasicBlock *BB : depth_first(&F)) {
106    LLVM_DEBUG(dbgs() << "  " << InstNo << ": BB " << BB->getName() << "\n");
107    unsigned BBStart = InstNo++;
108
109    BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
110    BlockInfo.Begin.resize(NumAllocas);
111    BlockInfo.End.resize(NumAllocas);
112    BlockInfo.LiveIn.resize(NumAllocas);
113    BlockInfo.LiveOut.resize(NumAllocas);
114
115    auto &BlockMarkerSet = BBMarkerSet[BB];
116    if (BlockMarkerSet.empty()) {
117      unsigned BBEnd = InstNo;
118      BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
119      continue;
120    }
121
122    auto ProcessMarker = [&](Instruction *I, const Marker &M) {
123      LLVM_DEBUG(dbgs() << "  " << InstNo << ":  "
124                        << (M.IsStart ? "start " : "end   ") << M.AllocaNo
125                        << ", " << *I << "\n");
126
127      BBMarkers[BB].push_back({InstNo, M});
128
129      InstructionNumbering[I] = InstNo++;
130
131      if (M.IsStart) {
132        if (BlockInfo.End.test(M.AllocaNo))
133          BlockInfo.End.reset(M.AllocaNo);
134        BlockInfo.Begin.set(M.AllocaNo);
135      } else {
136        if (BlockInfo.Begin.test(M.AllocaNo))
137          BlockInfo.Begin.reset(M.AllocaNo);
138        BlockInfo.End.set(M.AllocaNo);
139      }
140    };
141
142    if (BlockMarkerSet.size() == 1) {
143      ProcessMarker(BlockMarkerSet.begin()->getFirst(),
144                    BlockMarkerSet.begin()->getSecond());
145    } else {
146      // Scan the BB to determine the marker order.
147      for (Instruction &I : *BB) {
148        auto It = BlockMarkerSet.find(&I);
149        if (It == BlockMarkerSet.end())
150          continue;
151        ProcessMarker(&I, It->getSecond());
152      }
153    }
154
155    unsigned BBEnd = InstNo;
156    BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
157  }
158  NumInst = InstNo;
159}
160
161void StackColoring::calculateLocalLiveness() {
162  bool changed = true;
163  while (changed) {
164    changed = false;
165
166    for (BasicBlock *BB : depth_first(&F)) {
167      BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
168
169      // Compute LiveIn by unioning together the LiveOut sets of all preds.
170      BitVector LocalLiveIn;
171      for (auto *PredBB : predecessors(BB)) {
172        LivenessMap::const_iterator I = BlockLiveness.find(PredBB);
173        // If a predecessor is unreachable, ignore it.
174        if (I == BlockLiveness.end())
175          continue;
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
205void 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)
255LLVM_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
261LLVM_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
274LLVM_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
283void 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