1 //===- SafeStackColoring.h - SafeStack frame coloring ----------*- C++ -*--===//
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 #ifndef LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H
11 #define LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H
12 
13 #include "llvm/ADT/ArrayRef.h"
14 #include "llvm/ADT/BitVector.h"
15 #include "llvm/ADT/DenseMap.h"
16 #include "llvm/ADT/SmallVector.h"
17 #include "llvm/IR/Instructions.h"
18 #include "llvm/Support/raw_ostream.h"
19 #include <cassert>
20 #include <utility>
21 
22 namespace llvm {
23 
24 class BasicBlock;
25 class Function;
26 class Instruction;
27 
28 namespace safestack {
29 
30 /// Compute live ranges of allocas.
31 /// Live ranges are represented as sets of "interesting" instructions, which are
32 /// defined as instructions that may start or end an alloca's lifetime. These
33 /// are:
34 /// * lifetime.start and lifetime.end intrinsics
35 /// * first instruction of any basic block
36 /// Interesting instructions are numbered in the depth-first walk of the CFG,
37 /// and in the program order inside each basic block.
38 class StackColoring {
39   /// A class representing liveness information for a single basic block.
40   /// Each bit in the BitVector represents the liveness property
41   /// for a different stack slot.
42   struct BlockLifetimeInfo {
43     /// Which slots BEGINs in each basic block.
44     BitVector Begin;
45 
46     /// Which slots ENDs in each basic block.
47     BitVector End;
48 
49     /// Which slots are marked as LIVE_IN, coming into each basic block.
50     BitVector LiveIn;
51 
52     /// Which slots are marked as LIVE_OUT, coming out of each basic block.
53     BitVector LiveOut;
54   };
55 
56 public:
57   /// This class represents a set of interesting instructions where an alloca is
58   /// live.
59   struct LiveRange {
60     BitVector bv;
61 
SetMaximumLiveRange62     void SetMaximum(int size) { bv.resize(size); }
AddRangeLiveRange63     void AddRange(unsigned start, unsigned end) { bv.set(start, end); }
64 
OverlapsLiveRange65     bool Overlaps(const LiveRange &Other) const {
66       return bv.anyCommon(Other.bv);
67     }
68 
JoinLiveRange69     void Join(const LiveRange &Other) { bv |= Other.bv; }
70   };
71 
72 private:
73   Function &F;
74 
75   /// Maps active slots (per bit) for each basic block.
76   using LivenessMap = DenseMap<BasicBlock *, BlockLifetimeInfo>;
77   LivenessMap BlockLiveness;
78 
79   /// Number of interesting instructions.
80   int NumInst = -1;
81 
82   /// Numeric ids for interesting instructions.
83   DenseMap<Instruction *, unsigned> InstructionNumbering;
84 
85   /// A range [Start, End) of instruction ids for each basic block.
86   /// Instructions inside each BB have monotonic and consecutive ids.
87   DenseMap<const BasicBlock *, std::pair<unsigned, unsigned>> BlockInstRange;
88 
89   ArrayRef<AllocaInst *> Allocas;
90   unsigned NumAllocas;
91   DenseMap<AllocaInst *, unsigned> AllocaNumbering;
92 
93   /// LiveRange for allocas.
94   SmallVector<LiveRange, 8> LiveRanges;
95 
96   /// The set of allocas that have at least one lifetime.start. All other
97   /// allocas get LiveRange that corresponds to the entire function.
98   BitVector InterestingAllocas;
99   SmallVector<Instruction *, 8> Markers;
100 
101   struct Marker {
102     unsigned AllocaNo;
103     bool IsStart;
104   };
105 
106   /// List of {InstNo, {AllocaNo, IsStart}} for each BB, ordered by InstNo.
107   DenseMap<BasicBlock *, SmallVector<std::pair<unsigned, Marker>, 4>> BBMarkers;
108 
109   void dumpAllocas();
110   void dumpBlockLiveness();
111   void dumpLiveRanges();
112 
113   bool readMarker(Instruction *I, bool *IsStart);
114   void collectMarkers();
115   void calculateLocalLiveness();
116   void calculateLiveIntervals();
117 
118 public:
StackColoring(Function & F,ArrayRef<AllocaInst * > Allocas)119   StackColoring(Function &F, ArrayRef<AllocaInst *> Allocas)
120       : F(F), Allocas(Allocas), NumAllocas(Allocas.size()) {}
121 
122   void run();
123   void removeAllMarkers();
124 
125   /// Returns a set of "interesting" instructions where the given alloca is
126   /// live. Not all instructions in a function are interesting: we pick a set
127   /// that is large enough for LiveRange::Overlaps to be correct.
128   const LiveRange &getLiveRange(AllocaInst *AI);
129 
130   /// Returns a live range that represents an alloca that is live throughout the
131   /// entire function.
getFullLiveRange()132   LiveRange getFullLiveRange() {
133     assert(NumInst >= 0);
134     LiveRange R;
135     R.SetMaximum(NumInst);
136     R.AddRange(0, NumInst);
137     return R;
138   }
139 };
140 
141 static inline raw_ostream &operator<<(raw_ostream &OS, const BitVector &V) {
142   OS << "{";
143   int idx = V.find_first();
144   bool first = true;
145   while (idx >= 0) {
146     if (!first) {
147       OS << ", ";
148     }
149     first = false;
150     OS << idx;
151     idx = V.find_next(idx);
152   }
153   OS << "}";
154   return OS;
155 }
156 
157 static inline raw_ostream &operator<<(raw_ostream &OS,
158                                       const StackColoring::LiveRange &R) {
159   return OS << R.bv;
160 }
161 
162 } // end namespace safestack
163 
164 } // end namespace llvm
165 
166 #endif // LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H
167