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