lib/CodeGen/SafeStackColoring.cpp

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