contrib/llvm/lib/CodeGen/SafeStackColoring.cpp

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