lib/Target/AArch64/AArch64StackTaggingPreRA.cpp

   1 //===-- AArch64StackTaggingPreRA.cpp --- Stack Tagging for AArch64 -----===//
   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
  11 #include "AArch64.h"
  12 #include "AArch64MachineFunctionInfo.h"
  13 #include "AArch64InstrInfo.h"
  14 #include "llvm/ADT/DepthFirstIterator.h"
  15 #include "llvm/ADT/SetVector.h"
  16 #include "llvm/ADT/MapVector.h"
  17 #include "llvm/ADT/Statistic.h"
  18 #include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
  19 #include "llvm/CodeGen/MachineFrameInfo.h"
  20 #include "llvm/CodeGen/MachineFunction.h"
  21 #include "llvm/CodeGen/MachineFunctionPass.h"
  22 #include "llvm/CodeGen/MachineInstrBuilder.h"
  23 #include "llvm/CodeGen/MachineLoopInfo.h"
  24 #include "llvm/CodeGen/MachineRegisterInfo.h"
  25 #include "llvm/CodeGen/MachineTraceMetrics.h"
  26 #include "llvm/CodeGen/Passes.h"
  27 #include "llvm/CodeGen/TargetInstrInfo.h"
  28 #include "llvm/CodeGen/TargetRegisterInfo.h"
  29 #include "llvm/CodeGen/TargetSubtargetInfo.h"
  30 #include "llvm/Support/CommandLine.h"
  31 #include "llvm/Support/Debug.h"
  32 #include "llvm/Support/raw_ostream.h"
  33
  34 using namespace llvm;
  35
  36 #define DEBUG_TYPE "aarch64-stack-tagging-pre-ra"
  37
  38 enum UncheckedLdStMode { UncheckedNever, UncheckedSafe, UncheckedAlways };
  39
  40 cl::opt<UncheckedLdStMode> ClUncheckedLdSt(
  41     "stack-tagging-unchecked-ld-st", cl::Hidden,
  42     cl::init(UncheckedSafe),
  43     cl::desc(
  44         "Unconditionally apply unchecked-ld-st optimization (even for large "
  45         "stack frames, or in the presence of variable sized allocas)."),
  46     cl::values(
  47         clEnumValN(UncheckedNever, "never", "never apply unchecked-ld-st"),
  48         clEnumValN(
  49             UncheckedSafe, "safe",
  50             "apply unchecked-ld-st when the target is definitely within range"),
  51         clEnumValN(UncheckedAlways, "always", "always apply unchecked-ld-st")));
  52
  53 namespace {
  54
  55 class AArch64StackTaggingPreRA : public MachineFunctionPass {
  56   MachineFunction *MF;
  57   AArch64FunctionInfo *AFI;
  58   MachineFrameInfo *MFI;
  59   MachineRegisterInfo *MRI;
  60   const AArch64RegisterInfo *TRI;
  61   const AArch64InstrInfo *TII;
  62
  63   SmallVector<MachineInstr*, 16> ReTags;
  64
  65 public:
  66   static char ID;
  67   AArch64StackTaggingPreRA() : MachineFunctionPass(ID) {
  68     initializeAArch64StackTaggingPreRAPass(*PassRegistry::getPassRegistry());
  69   }
  70
  71   bool mayUseUncheckedLoadStore();
  72   void uncheckUsesOf(unsigned TaggedReg, int FI);
  73   void uncheckLoadsAndStores();
  74
  75   bool runOnMachineFunction(MachineFunction &Func) override;
  76   StringRef getPassName() const override {
  77     return "AArch64 Stack Tagging PreRA";
  78   }
  79
  80   void getAnalysisUsage(AnalysisUsage &AU) const override {
  81     AU.setPreservesCFG();
  82     MachineFunctionPass::getAnalysisUsage(AU);
  83   }
  84 };
  85 } // end anonymous namespace
  86
  87 char AArch64StackTaggingPreRA::ID = 0;
  88
  89 INITIALIZE_PASS_BEGIN(AArch64StackTaggingPreRA, "aarch64-stack-tagging-pre-ra",
  90                       "AArch64 Stack Tagging PreRA Pass", false, false)
  91 INITIALIZE_PASS_END(AArch64StackTaggingPreRA, "aarch64-stack-tagging-pre-ra",
  92                     "AArch64 Stack Tagging PreRA Pass", false, false)
  93
  94 FunctionPass *llvm::createAArch64StackTaggingPreRAPass() {
  95   return new AArch64StackTaggingPreRA();
  96 }
  97
  98 static bool isUncheckedLoadOrStoreOpcode(unsigned Opcode) {
  99   switch (Opcode) {
 100   case AArch64::LDRWui:
 101   case AArch64::LDRSHWui:
 102   case AArch64::LDRXui:
 103   case AArch64::LDRBui:
 104   case AArch64::LDRBBui:
 105   case AArch64::LDRHui:
 106   case AArch64::LDRSui:
 107   case AArch64::LDRDui:
 108   case AArch64::LDRQui:
 109   case AArch64::STRWui:
 110   case AArch64::STRXui:
 111   case AArch64::STRBui:
 112   case AArch64::STRBBui:
 113   case AArch64::STRHui:
 114   case AArch64::STRSui:
 115   case AArch64::STRDui:
 116   case AArch64::STRQui:
 117     return true;
 118   default:
 119     return false;
 120   }
 121 }
 122
 123 bool AArch64StackTaggingPreRA::mayUseUncheckedLoadStore() {
 124   if (ClUncheckedLdSt == UncheckedNever)
 125     return false;
 126   else if (ClUncheckedLdSt == UncheckedAlways)
 127     return true;
 128
 129   // This estimate can be improved if we had harder guarantees about stack frame
 130   // layout. With LocalStackAllocation we can estimate SP offset to any
 131   // preallocated slot. AArch64FrameLowering::orderFrameObjects could put tagged
 132   // objects ahead of non-tagged ones, but that's not always desirable.
 133   //
 134   // Underestimating SP offset here may require the use of LDG to materialize
 135   // the tagged address of the stack slot, along with a scratch register
 136   // allocation (post-regalloc!).
 137   //
 138   // For now we do the safe thing here and require that the entire stack frame
 139   // is within range of the shortest of the unchecked instructions.
 140   unsigned FrameSize = 0;
 141   for (unsigned i = 0, e = MFI->getObjectIndexEnd(); i != e; ++i)
 142     FrameSize += MFI->getObjectSize(i);
 143   bool EntireFrameReachableFromSP = FrameSize < 0xf00;
 144   return !MFI->hasVarSizedObjects() && EntireFrameReachableFromSP;
 145 }
 146
 147 void AArch64StackTaggingPreRA::uncheckUsesOf(unsigned TaggedReg, int FI) {
 148   for (auto UI = MRI->use_instr_begin(TaggedReg), E = MRI->use_instr_end();
 149        UI != E;) {
 150     MachineInstr *UseI = &*(UI++);
 151     if (isUncheckedLoadOrStoreOpcode(UseI->getOpcode())) {
 152       // FI operand is always the one before the immediate offset.
 153       unsigned OpIdx = TII->getLoadStoreImmIdx(UseI->getOpcode()) - 1;
 154       if (UseI->getOperand(OpIdx).isReg() &&
 155           UseI->getOperand(OpIdx).getReg() == TaggedReg) {
 156         UseI->getOperand(OpIdx).ChangeToFrameIndex(FI);
 157         UseI->getOperand(OpIdx).setTargetFlags(AArch64II::MO_TAGGED);
 158       }
 159     } else if (UseI->isCopy() &&
 160                Register::isVirtualRegister(UseI->getOperand(0).getReg())) {
 161       uncheckUsesOf(UseI->getOperand(0).getReg(), FI);
 162     }
 163   }
 164 }
 165
 166 void AArch64StackTaggingPreRA::uncheckLoadsAndStores() {
 167   for (auto *I : ReTags) {
 168     unsigned TaggedReg = I->getOperand(0).getReg();
 169     int FI = I->getOperand(1).getIndex();
 170     uncheckUsesOf(TaggedReg, FI);
 171   }
 172 }
 173
 174 bool AArch64StackTaggingPreRA::runOnMachineFunction(MachineFunction &Func) {
 175   MF = &Func;
 176   MRI = &MF->getRegInfo();
 177   AFI = MF->getInfo<AArch64FunctionInfo>();
 178   TII = static_cast<const AArch64InstrInfo *>(MF->getSubtarget().getInstrInfo());
 179   TRI = static_cast<const AArch64RegisterInfo *>(
 180       MF->getSubtarget().getRegisterInfo());
 181   MFI = &MF->getFrameInfo();
 182   ReTags.clear();
 183
 184   assert(MRI->isSSA());
 185
 186   LLVM_DEBUG(dbgs() << "********** AArch64 Stack Tagging PreRA **********\n"
 187                     << "********** Function: " << MF->getName() << '\n');
 188
 189   SmallSetVector<int, 8> TaggedSlots;
 190   for (auto &BB : *MF) {
 191     for (auto &I : BB) {
 192       if (I.getOpcode() == AArch64::TAGPstack) {
 193         ReTags.push_back(&I);
 194         int FI = I.getOperand(1).getIndex();
 195         TaggedSlots.insert(FI);
 196         // There should be no offsets in TAGP yet.
 197         assert(I.getOperand(2).getImm() == 0);
 198       }
 199     }
 200   }
 201
 202   if (ReTags.empty())
 203     return false;
 204
 205   if (mayUseUncheckedLoadStore())
 206     uncheckLoadsAndStores();
 207
 208   return true;
 209 }