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 // Disabled by default due to PR32143.
  24 static cl::opt<bool> ClColoring("safe-stack-coloring",
  25                                 cl::desc("enable safe stack coloring"),
  26                                 cl::Hidden, cl::init(false));
  27
  28 const StackColoring::LiveRange &StackColoring::getLiveRange(AllocaInst *AI) {
  29   const auto IT = AllocaNumbering.find(AI);
  30   assert(IT != AllocaNumbering.end());
  31   return LiveRanges[IT->second];
  32 }
  33
  34 bool StackColoring::readMarker(Instruction *I, bool *IsStart) {
  35   auto *II = dyn_cast<IntrinsicInst>(I);
  36   if (!II || (II->getIntrinsicID() != Intrinsic::lifetime_start &&
  37               II->getIntrinsicID() != Intrinsic::lifetime_end))
  38     return false;
  39
  40   *IsStart = II->getIntrinsicID() == Intrinsic::lifetime_start;
  41   return true;
  42 }
  43
  44 void StackColoring::removeAllMarkers() {
  45   for (auto *I : Markers) {
  46     auto *Op = dyn_cast<Instruction>(I->getOperand(1));
  47     I->eraseFromParent();
  48     // Remove the operand bitcast, too, if it has no more uses left.
  49     if (Op && Op->use_empty())
  50       Op->eraseFromParent();
  51   }
  52 }
  53
  54 void StackColoring::collectMarkers() {
  55   InterestingAllocas.resize(NumAllocas);
  56   DenseMap<BasicBlock *, SmallDenseMap<Instruction *, Marker>> BBMarkerSet;
  57
  58   // Compute the set of start/end markers per basic block.
  59   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
  60     AllocaInst *AI = Allocas[AllocaNo];
  61     SmallVector<Instruction *, 8> WorkList;
  62     WorkList.push_back(AI);
  63     while (!WorkList.empty()) {
  64       Instruction *I = WorkList.pop_back_val();
  65       for (User *U : I->users()) {
  66         if (auto *BI = dyn_cast<BitCastInst>(U)) {
  67           WorkList.push_back(BI);
  68           continue;
  69         }
  70         auto *UI = dyn_cast<Instruction>(U);
  71         if (!UI)
  72           continue;
  73         bool IsStart;
  74         if (!readMarker(UI, &IsStart))
  75           continue;
  76         if (IsStart)
  77           InterestingAllocas.set(AllocaNo);
  78         BBMarkerSet[UI->getParent()][UI] = {AllocaNo, IsStart};
  79         Markers.push_back(UI);
  80       }
  81     }
  82   }
  83
  84   // Compute instruction numbering. Only the following instructions are
  85   // considered:
  86   // * Basic block entries
  87   // * Lifetime markers
  88   // For each basic block, compute
  89   // * the list of markers in the instruction order
  90   // * the sets of allocas whose lifetime starts or ends in this BB
  91   DEBUG(dbgs() << "Instructions:\n");
  92   unsigned InstNo = 0;
  93   for (BasicBlock *BB : depth_first(&F)) {
  94     DEBUG(dbgs() << "  " << InstNo << ": BB " << BB->getName() << "\n");
  95     unsigned BBStart = InstNo++;
  96
  97     BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
  98     BlockInfo.Begin.resize(NumAllocas);
  99     BlockInfo.End.resize(NumAllocas);
 100     BlockInfo.LiveIn.resize(NumAllocas);
 101     BlockInfo.LiveOut.resize(NumAllocas);
 102
 103     auto &BlockMarkerSet = BBMarkerSet[BB];
 104     if (BlockMarkerSet.empty()) {
 105       unsigned BBEnd = InstNo;
 106       BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
 107       continue;
 108     }
 109
 110     auto ProcessMarker = [&](Instruction *I, const Marker &M) {
 111       DEBUG(dbgs() << "  " << InstNo << ":  "
 112                    << (M.IsStart ? "start " : "end   ") << M.AllocaNo << ", "
 113                    << *I << "\n");
 114
 115       BBMarkers[BB].push_back({InstNo, M});
 116
 117       InstructionNumbering[I] = InstNo++;
 118
 119       if (M.IsStart) {
 120         if (BlockInfo.End.test(M.AllocaNo))
 121           BlockInfo.End.reset(M.AllocaNo);
 122         BlockInfo.Begin.set(M.AllocaNo);
 123       } else {
 124         if (BlockInfo.Begin.test(M.AllocaNo))
 125           BlockInfo.Begin.reset(M.AllocaNo);
 126         BlockInfo.End.set(M.AllocaNo);
 127       }
 128     };
 129
 130     if (BlockMarkerSet.size() == 1) {
 131       ProcessMarker(BlockMarkerSet.begin()->getFirst(),
 132                     BlockMarkerSet.begin()->getSecond());
 133     } else {
 134       // Scan the BB to determine the marker order.
 135       for (Instruction &I : *BB) {
 136         auto It = BlockMarkerSet.find(&I);
 137         if (It == BlockMarkerSet.end())
 138           continue;
 139         ProcessMarker(&I, It->getSecond());
 140       }
 141     }
 142
 143     unsigned BBEnd = InstNo;
 144     BlockInstRange[BB] = std::make_pair(BBStart, BBEnd);
 145   }
 146   NumInst = InstNo;
 147 }
 148
 149 void StackColoring::calculateLocalLiveness() {
 150   bool changed = true;
 151   while (changed) {
 152     changed = false;
 153
 154     for (BasicBlock *BB : depth_first(&F)) {
 155       BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
 156
 157       // Compute LiveIn by unioning together the LiveOut sets of all preds.
 158       BitVector LocalLiveIn;
 159       for (auto *PredBB : predecessors(BB)) {
 160         LivenessMap::const_iterator I = BlockLiveness.find(PredBB);
 161         assert(I != BlockLiveness.end() && "Predecessor not found");
 162         LocalLiveIn |= I->second.LiveOut;
 163       }
 164
 165       // Compute LiveOut by subtracting out lifetimes that end in this
 166       // block, then adding in lifetimes that begin in this block.  If
 167       // we have both BEGIN and END markers in the same basic block
 168       // then we know that the BEGIN marker comes after the END,
 169       // because we already handle the case where the BEGIN comes
 170       // before the END when collecting the markers (and building the
 171       // BEGIN/END vectors).
 172       BitVector LocalLiveOut = LocalLiveIn;
 173       LocalLiveOut.reset(BlockInfo.End);
 174       LocalLiveOut |= BlockInfo.Begin;
 175
 176       // Update block LiveIn set, noting whether it has changed.
 177       if (LocalLiveIn.test(BlockInfo.LiveIn)) {
 178         changed = true;
 179         BlockInfo.LiveIn |= LocalLiveIn;
 180       }
 181
 182       // Update block LiveOut set, noting whether it has changed.
 183       if (LocalLiveOut.test(BlockInfo.LiveOut)) {
 184         changed = true;
 185         BlockInfo.LiveOut |= LocalLiveOut;
 186       }
 187     }
 188   } // while changed.
 189 }
 190
 191 void StackColoring::calculateLiveIntervals() {
 192   for (auto IT : BlockLiveness) {
 193     BasicBlock *BB = IT.getFirst();
 194     BlockLifetimeInfo &BlockInfo = IT.getSecond();
 195     unsigned BBStart, BBEnd;
 196     std::tie(BBStart, BBEnd) = BlockInstRange[BB];
 197
 198     BitVector Started, Ended;
 199     Started.resize(NumAllocas);
 200     Ended.resize(NumAllocas);
 201     SmallVector<unsigned, 8> Start;
 202     Start.resize(NumAllocas);
 203
 204     // LiveIn ranges start at the first instruction.
 205     for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
 206       if (BlockInfo.LiveIn.test(AllocaNo)) {
 207         Started.set(AllocaNo);
 208         Start[AllocaNo] = BBStart;
 209       }
 210     }
 211
 212     for (auto &It : BBMarkers[BB]) {
 213       unsigned InstNo = It.first;
 214       bool IsStart = It.second.IsStart;
 215       unsigned AllocaNo = It.second.AllocaNo;
 216
 217       if (IsStart) {
 218         assert(!Started.test(AllocaNo) || Start[AllocaNo] == BBStart);
 219         if (!Started.test(AllocaNo)) {
 220           Started.set(AllocaNo);
 221           Ended.reset(AllocaNo);
 222           Start[AllocaNo] = InstNo;
 223         }
 224       } else {
 225         assert(!Ended.test(AllocaNo));
 226         if (Started.test(AllocaNo)) {
 227           LiveRanges[AllocaNo].AddRange(Start[AllocaNo], InstNo);
 228           Started.reset(AllocaNo);
 229         }
 230         Ended.set(AllocaNo);
 231       }
 232     }
 233
 234     for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
 235       if (Started.test(AllocaNo))
 236         LiveRanges[AllocaNo].AddRange(Start[AllocaNo], BBEnd);
 237   }
 238 }
 239
 240 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 241 LLVM_DUMP_METHOD void StackColoring::dumpAllocas() {
 242   dbgs() << "Allocas:\n";
 243   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
 244     dbgs() << "  " << AllocaNo << ": " << *Allocas[AllocaNo] << "\n";
 245 }
 246
 247 LLVM_DUMP_METHOD void StackColoring::dumpBlockLiveness() {
 248   dbgs() << "Block liveness:\n";
 249   for (auto IT : BlockLiveness) {
 250     BasicBlock *BB = IT.getFirst();
 251     BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
 252     auto BlockRange = BlockInstRange[BB];
 253     dbgs() << "  BB [" << BlockRange.first << ", " << BlockRange.second
 254            << "): begin " << BlockInfo.Begin << ", end " << BlockInfo.End
 255            << ", livein " << BlockInfo.LiveIn << ", liveout "
 256            << BlockInfo.LiveOut << "\n";
 257   }
 258 }
 259
 260 LLVM_DUMP_METHOD void StackColoring::dumpLiveRanges() {
 261   dbgs() << "Alloca liveness:\n";
 262   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
 263     LiveRange &Range = LiveRanges[AllocaNo];
 264     dbgs() << "  " << AllocaNo << ": " << Range << "\n";
 265   }
 266 }
 267 #endif
 268
 269 void StackColoring::run() {
 270   DEBUG(dumpAllocas());
 271
 272   for (unsigned I = 0; I < NumAllocas; ++I)
 273     AllocaNumbering[Allocas[I]] = I;
 274   LiveRanges.resize(NumAllocas);
 275
 276   collectMarkers();
 277
 278   if (!ClColoring) {
 279     for (auto &R : LiveRanges) {
 280       R.SetMaximum(1);
 281       R.AddRange(0, 1);
 282     }
 283     return;
 284   }
 285
 286   for (auto &R : LiveRanges)
 287     R.SetMaximum(NumInst);
 288   for (unsigned I = 0; I < NumAllocas; ++I)
 289     if (!InterestingAllocas.test(I))
 290       LiveRanges[I] = getFullLiveRange();
 291
 292   calculateLocalLiveness();
 293   DEBUG(dumpBlockLiveness());
 294   calculateLiveIntervals();
 295   DEBUG(dumpLiveRanges());
 296 }