lib/CodeGen/SafeStackColoring.h

   1 //===- SafeStackColoring.h - SafeStack frame coloring ----------*- C++ -*--===//
   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 #ifndef LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H
  10 #define LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H
  11
  12 #include "llvm/ADT/ArrayRef.h"
  13 #include "llvm/ADT/BitVector.h"
  14 #include "llvm/ADT/DenseMap.h"
  15 #include "llvm/ADT/SmallVector.h"
  16 #include "llvm/IR/Instructions.h"
  17 #include "llvm/Support/raw_ostream.h"
  18 #include <cassert>
  19 #include <utility>
  20
  21 namespace llvm {
  22
  23 class BasicBlock;
  24 class Function;
  25 class Instruction;
  26
  27 namespace safestack {
  28
  29 /// Compute live ranges of allocas.
  30 /// Live ranges are represented as sets of "interesting" instructions, which are
  31 /// defined as instructions that may start or end an alloca's lifetime. These
  32 /// are:
  33 /// * lifetime.start and lifetime.end intrinsics
  34 /// * first instruction of any basic block
  35 /// Interesting instructions are numbered in the depth-first walk of the CFG,
  36 /// and in the program order inside each basic block.
  37 class StackColoring {
  38   /// A class representing liveness information for a single basic block.
  39   /// Each bit in the BitVector represents the liveness property
  40   /// for a different stack slot.
  41   struct BlockLifetimeInfo {
  42     /// Which slots BEGINs in each basic block.
  43     BitVector Begin;
  44
  45     /// Which slots ENDs in each basic block.
  46     BitVector End;
  47
  48     /// Which slots are marked as LIVE_IN, coming into each basic block.
  49     BitVector LiveIn;
  50
  51     /// Which slots are marked as LIVE_OUT, coming out of each basic block.
  52     BitVector LiveOut;
  53   };
  54
  55 public:
  56   /// This class represents a set of interesting instructions where an alloca is
  57   /// live.
  58   struct LiveRange {
  59     BitVector bv;
  60
  61     void SetMaximum(int size) { bv.resize(size); }
  62     void AddRange(unsigned start, unsigned end) { bv.set(start, end); }
  63
  64     bool Overlaps(const LiveRange &Other) const {
  65       return bv.anyCommon(Other.bv);
  66     }
  67
  68     void Join(const LiveRange &Other) { bv |= Other.bv; }
  69   };
  70
  71 private:
  72   Function &F;
  73
  74   /// Maps active slots (per bit) for each basic block.
  75   using LivenessMap = DenseMap<BasicBlock *, BlockLifetimeInfo>;
  76   LivenessMap BlockLiveness;
  77
  78   /// Number of interesting instructions.
  79   int NumInst = -1;
  80
  81   /// Numeric ids for interesting instructions.
  82   DenseMap<Instruction *, unsigned> InstructionNumbering;
  83
  84   /// A range [Start, End) of instruction ids for each basic block.
  85   /// Instructions inside each BB have monotonic and consecutive ids.
  86   DenseMap<const BasicBlock *, std::pair<unsigned, unsigned>> BlockInstRange;
  87
  88   ArrayRef<AllocaInst *> Allocas;
  89   unsigned NumAllocas;
  90   DenseMap<AllocaInst *, unsigned> AllocaNumbering;
  91
  92   /// LiveRange for allocas.
  93   SmallVector<LiveRange, 8> LiveRanges;
  94
  95   /// The set of allocas that have at least one lifetime.start. All other
  96   /// allocas get LiveRange that corresponds to the entire function.
  97   BitVector InterestingAllocas;
  98   SmallVector<Instruction *, 8> Markers;
  99
 100   struct Marker {
 101     unsigned AllocaNo;
 102     bool IsStart;
 103   };
 104
 105   /// List of {InstNo, {AllocaNo, IsStart}} for each BB, ordered by InstNo.
 106   DenseMap<BasicBlock *, SmallVector<std::pair<unsigned, Marker>, 4>> BBMarkers;
 107
 108   void dumpAllocas();
 109   void dumpBlockLiveness();
 110   void dumpLiveRanges();
 111
 112   bool readMarker(Instruction *I, bool *IsStart);
 113   void collectMarkers();
 114   void calculateLocalLiveness();
 115   void calculateLiveIntervals();
 116
 117 public:
 118   StackColoring(Function &F, ArrayRef<AllocaInst *> Allocas)
 119       : F(F), Allocas(Allocas), NumAllocas(Allocas.size()) {}
 120
 121   void run();
 122   void removeAllMarkers();
 123
 124   /// Returns a set of "interesting" instructions where the given alloca is
 125   /// live. Not all instructions in a function are interesting: we pick a set
 126   /// that is large enough for LiveRange::Overlaps to be correct.
 127   const LiveRange &getLiveRange(AllocaInst *AI);
 128
 129   /// Returns a live range that represents an alloca that is live throughout the
 130   /// entire function.
 131   LiveRange getFullLiveRange() {
 132     assert(NumInst >= 0);
 133     LiveRange R;
 134     R.SetMaximum(NumInst);
 135     R.AddRange(0, NumInst);
 136     return R;
 137   }
 138 };
 139
 140 static inline raw_ostream &operator<<(raw_ostream &OS, const BitVector &V) {
 141   OS << "{";
 142   int idx = V.find_first();
 143   bool first = true;
 144   while (idx >= 0) {
 145     if (!first) {
 146       OS << ", ";
 147     }
 148     first = false;
 149     OS << idx;
 150     idx = V.find_next(idx);
 151   }
 152   OS << "}";
 153   return OS;
 154 }
 155
 156 static inline raw_ostream &operator<<(raw_ostream &OS,
 157                                       const StackColoring::LiveRange &R) {
 158   return OS << R.bv;
 159 }
 160
 161 } // end namespace safestack
 162
 163 } // end namespace llvm
 164
 165 #endif // LLVM_LIB_CODEGEN_SAFESTACKCOLORING_H