contrib/llvm/lib/CodeGen/SafeStackColoring.cpp

   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 }