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