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