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/Config/llvm-config.h"
  16 #include "llvm/IR/BasicBlock.h"
  17 #include "llvm/IR/CFG.h"
  18 #include "llvm/IR/Instruction.h"
  19 #include "llvm/IR/Instructions.h"
  20 #include "llvm/IR/IntrinsicInst.h"
  21 #include "llvm/IR/Intrinsics.h"
  22 #include "llvm/IR/User.h"
  23 #include "llvm/Support/Casting.h"
  24 #include "llvm/Support/CommandLine.h"
  25 #include "llvm/Support/Compiler.h"
  26 #include "llvm/Support/Debug.h"
  27 #include "llvm/Support/raw_ostream.h"
  28 #include <cassert>
  29 #include <tuple>
  30 #include <utility>
  31
  32 using namespace llvm;
  33 using namespace llvm::safestack;
  34
  35 #define DEBUG_TYPE "safestackcoloring"
  36
  37 // Disabled by default due to PR32143.
  38 static cl::opt<bool> ClColoring("safe-stack-coloring",
  39                                 cl::desc("enable safe stack coloring"),
  40                                 cl::Hidden, cl::init(false));
  41
  42 const StackColoring::LiveRange &StackColoring::getLiveRange(AllocaInst *AI) {
  43   const auto IT = AllocaNumbering.find(AI);
  44   assert(IT != AllocaNumbering.end());
  45   return LiveRanges[IT->second];
  46 }
  47
  48 bool StackColoring::readMarker(Instruction *I, bool *IsStart) {
  49   if (!I->isLifetimeStartOrEnd())
  50     return false;
  51
  52   auto *II = cast<IntrinsicInst>(I);
  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   LLVM_DEBUG(dbgs() << "Instructions:\n");
 105   unsigned InstNo = 0;
 106   for (BasicBlock *BB : depth_first(&F)) {
 107     LLVM_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       LLVM_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         // If a predecessor is unreachable, ignore it.
 175         if (I == BlockLiveness.end())
 176           continue;
 177         LocalLiveIn |= I->second.LiveOut;
 178       }
 179
 180       // Compute LiveOut by subtracting out lifetimes that end in this
 181       // block, then adding in lifetimes that begin in this block.  If
 182       // we have both BEGIN and END markers in the same basic block
 183       // then we know that the BEGIN marker comes after the END,
 184       // because we already handle the case where the BEGIN comes
 185       // before the END when collecting the markers (and building the
 186       // BEGIN/END vectors).
 187       BitVector LocalLiveOut = LocalLiveIn;
 188       LocalLiveOut.reset(BlockInfo.End);
 189       LocalLiveOut |= BlockInfo.Begin;
 190
 191       // Update block LiveIn set, noting whether it has changed.
 192       if (LocalLiveIn.test(BlockInfo.LiveIn)) {
 193         changed = true;
 194         BlockInfo.LiveIn |= LocalLiveIn;
 195       }
 196
 197       // Update block LiveOut set, noting whether it has changed.
 198       if (LocalLiveOut.test(BlockInfo.LiveOut)) {
 199         changed = true;
 200         BlockInfo.LiveOut |= LocalLiveOut;
 201       }
 202     }
 203   } // while changed.
 204 }
 205
 206 void StackColoring::calculateLiveIntervals() {
 207   for (auto IT : BlockLiveness) {
 208     BasicBlock *BB = IT.getFirst();
 209     BlockLifetimeInfo &BlockInfo = IT.getSecond();
 210     unsigned BBStart, BBEnd;
 211     std::tie(BBStart, BBEnd) = BlockInstRange[BB];
 212
 213     BitVector Started, Ended;
 214     Started.resize(NumAllocas);
 215     Ended.resize(NumAllocas);
 216     SmallVector<unsigned, 8> Start;
 217     Start.resize(NumAllocas);
 218
 219     // LiveIn ranges start at the first instruction.
 220     for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
 221       if (BlockInfo.LiveIn.test(AllocaNo)) {
 222         Started.set(AllocaNo);
 223         Start[AllocaNo] = BBStart;
 224       }
 225     }
 226
 227     for (auto &It : BBMarkers[BB]) {
 228       unsigned InstNo = It.first;
 229       bool IsStart = It.second.IsStart;
 230       unsigned AllocaNo = It.second.AllocaNo;
 231
 232       if (IsStart) {
 233         assert(!Started.test(AllocaNo) || Start[AllocaNo] == BBStart);
 234         if (!Started.test(AllocaNo)) {
 235           Started.set(AllocaNo);
 236           Ended.reset(AllocaNo);
 237           Start[AllocaNo] = InstNo;
 238         }
 239       } else {
 240         assert(!Ended.test(AllocaNo));
 241         if (Started.test(AllocaNo)) {
 242           LiveRanges[AllocaNo].AddRange(Start[AllocaNo], InstNo);
 243           Started.reset(AllocaNo);
 244         }
 245         Ended.set(AllocaNo);
 246       }
 247     }
 248
 249     for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
 250       if (Started.test(AllocaNo))
 251         LiveRanges[AllocaNo].AddRange(Start[AllocaNo], BBEnd);
 252   }
 253 }
 254
 255 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 256 LLVM_DUMP_METHOD void StackColoring::dumpAllocas() {
 257   dbgs() << "Allocas:\n";
 258   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo)
 259     dbgs() << "  " << AllocaNo << ": " << *Allocas[AllocaNo] << "\n";
 260 }
 261
 262 LLVM_DUMP_METHOD void StackColoring::dumpBlockLiveness() {
 263   dbgs() << "Block liveness:\n";
 264   for (auto IT : BlockLiveness) {
 265     BasicBlock *BB = IT.getFirst();
 266     BlockLifetimeInfo &BlockInfo = BlockLiveness[BB];
 267     auto BlockRange = BlockInstRange[BB];
 268     dbgs() << "  BB [" << BlockRange.first << ", " << BlockRange.second
 269            << "): begin " << BlockInfo.Begin << ", end " << BlockInfo.End
 270            << ", livein " << BlockInfo.LiveIn << ", liveout "
 271            << BlockInfo.LiveOut << "\n";
 272   }
 273 }
 274
 275 LLVM_DUMP_METHOD void StackColoring::dumpLiveRanges() {
 276   dbgs() << "Alloca liveness:\n";
 277   for (unsigned AllocaNo = 0; AllocaNo < NumAllocas; ++AllocaNo) {
 278     LiveRange &Range = LiveRanges[AllocaNo];
 279     dbgs() << "  " << AllocaNo << ": " << Range << "\n";
 280   }
 281 }
 282 #endif
 283
 284 void StackColoring::run() {
 285   LLVM_DEBUG(dumpAllocas());
 286
 287   for (unsigned I = 0; I < NumAllocas; ++I)
 288     AllocaNumbering[Allocas[I]] = I;
 289   LiveRanges.resize(NumAllocas);
 290
 291   collectMarkers();
 292
 293   if (!ClColoring) {
 294     for (auto &R : LiveRanges) {
 295       R.SetMaximum(1);
 296       R.AddRange(0, 1);
 297     }
 298     return;
 299   }
 300
 301   for (auto &R : LiveRanges)
 302     R.SetMaximum(NumInst);
 303   for (unsigned I = 0; I < NumAllocas; ++I)
 304     if (!InterestingAllocas.test(I))
 305       LiveRanges[I] = getFullLiveRange();
 306
 307   calculateLocalLiveness();
 308   LLVM_DEBUG(dumpBlockLiveness());
 309   calculateLiveIntervals();
 310   LLVM_DEBUG(dumpLiveRanges());
 311 }