reference, declarationdefinition
definition → references, declarations, derived classes, virtual overrides
reference to multiple definitions → definitions
unreferenced
    1
    2
    3
    4
    5
    6
    7
    8
    9
   10
   11
   12
   13
   14
   15
   16
   17
   18
   19
   20
   21
   22
   23
   24
   25
   26
   27
   28
   29
   30
   31
   32
   33
   34
   35
   36
   37
   38
   39
   40
   41
   42
   43
   44
   45
   46
   47
   48
   49
   50
   51
   52
   53
   54
   55
   56
   57
   58
   59
   60
   61
   62
   63
   64
   65
   66
   67
   68
   69
   70
   71
   72
   73
   74
   75
   76
   77
   78
   79
   80
   81
   82
   83
   84
   85
   86
   87
   88
   89
   90
   91
   92
   93
   94
   95
   96
   97
   98
   99
  100
  101
  102
  103
  104
  105
  106
  107
  108
  109
  110
  111
  112
  113
  114
  115
  116
  117
  118
  119
  120
  121
  122
  123
  124
  125
  126
  127
  128
  129
  130
  131
  132
  133
  134
  135
  136
  137
  138
  139
  140
  141
  142
  143
  144
  145
  146
  147
  148
  149
  150
  151
  152
  153
  154
  155
  156
  157
  158
  159
  160
  161
  162
  163
  164
  165
  166
  167
  168
  169
  170
  171
  172
  173
  174
  175
  176
  177
  178
  179
  180
  181
  182
  183
  184
  185
  186
  187
  188
  189
  190
  191
  192
  193
  194
  195
  196
  197
  198
  199
  200
  201
  202
  203
  204
  205
  206
  207
  208
  209
//===-- AArch64StackTaggingPreRA.cpp --- Stack Tagging for AArch64 -----===//
//
//                     The LLVM Compiler Infrastructure
//
// This file is distributed under the University of Illinois Open Source
// License. See LICENSE.TXT for details.
//
//===----------------------------------------------------------------------===//


#include "AArch64.h"
#include "AArch64MachineFunctionInfo.h"
#include "AArch64InstrInfo.h"
#include "llvm/ADT/DepthFirstIterator.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/CodeGen/MachineBranchProbabilityInfo.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineFunction.h"
#include "llvm/CodeGen/MachineFunctionPass.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineLoopInfo.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
#include "llvm/CodeGen/MachineTraceMetrics.h"
#include "llvm/CodeGen/Passes.h"
#include "llvm/CodeGen/TargetInstrInfo.h"
#include "llvm/CodeGen/TargetRegisterInfo.h"
#include "llvm/CodeGen/TargetSubtargetInfo.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

#define DEBUG_TYPE "aarch64-stack-tagging-pre-ra"

enum UncheckedLdStMode { UncheckedNever, UncheckedSafe, UncheckedAlways };

cl::opt<UncheckedLdStMode> ClUncheckedLdSt(
    "stack-tagging-unchecked-ld-st", cl::Hidden,
    cl::init(UncheckedSafe),
    cl::desc(
        "Unconditionally apply unchecked-ld-st optimization (even for large "
        "stack frames, or in the presence of variable sized allocas)."),
    cl::values(
        clEnumValN(UncheckedNever, "never", "never apply unchecked-ld-st"),
        clEnumValN(
            UncheckedSafe, "safe",
            "apply unchecked-ld-st when the target is definitely within range"),
        clEnumValN(UncheckedAlways, "always", "always apply unchecked-ld-st")));

namespace {

class AArch64StackTaggingPreRA : public MachineFunctionPass {
  MachineFunction *MF;
  AArch64FunctionInfo *AFI;
  MachineFrameInfo *MFI;
  MachineRegisterInfo *MRI;
  const AArch64RegisterInfo *TRI;
  const AArch64InstrInfo *TII;

  SmallVector<MachineInstr*, 16> ReTags;

public:
  static char ID;
  AArch64StackTaggingPreRA() : MachineFunctionPass(ID) {
    initializeAArch64StackTaggingPreRAPass(*PassRegistry::getPassRegistry());
  }

  bool mayUseUncheckedLoadStore();
  void uncheckUsesOf(unsigned TaggedReg, int FI);
  void uncheckLoadsAndStores();

  bool runOnMachineFunction(MachineFunction &Func) override;
  StringRef getPassName() const override {
    return "AArch64 Stack Tagging PreRA";
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    MachineFunctionPass::getAnalysisUsage(AU);
  }
};
} // end anonymous namespace

char AArch64StackTaggingPreRA::ID = 0;

INITIALIZE_PASS_BEGIN(AArch64StackTaggingPreRA, "aarch64-stack-tagging-pre-ra",
                      "AArch64 Stack Tagging PreRA Pass", false, false)
INITIALIZE_PASS_END(AArch64StackTaggingPreRA, "aarch64-stack-tagging-pre-ra",
                    "AArch64 Stack Tagging PreRA Pass", false, false)

FunctionPass *llvm::createAArch64StackTaggingPreRAPass() {
  return new AArch64StackTaggingPreRA();
}

static bool isUncheckedLoadOrStoreOpcode(unsigned Opcode) {
  switch (Opcode) {
  case AArch64::LDRWui:
  case AArch64::LDRSHWui:
  case AArch64::LDRXui:
  case AArch64::LDRBui:
  case AArch64::LDRBBui:
  case AArch64::LDRHui:
  case AArch64::LDRSui:
  case AArch64::LDRDui:
  case AArch64::LDRQui:
  case AArch64::STRWui:
  case AArch64::STRXui:
  case AArch64::STRBui:
  case AArch64::STRBBui:
  case AArch64::STRHui:
  case AArch64::STRSui:
  case AArch64::STRDui:
  case AArch64::STRQui:
    return true;
  default:
    return false;
  }
}

bool AArch64StackTaggingPreRA::mayUseUncheckedLoadStore() {
  if (ClUncheckedLdSt == UncheckedNever)
    return false;
  else if (ClUncheckedLdSt == UncheckedAlways)
    return true;

  // This estimate can be improved if we had harder guarantees about stack frame
  // layout. With LocalStackAllocation we can estimate SP offset to any
  // preallocated slot. AArch64FrameLowering::orderFrameObjects could put tagged
  // objects ahead of non-tagged ones, but that's not always desirable.
  //
  // Underestimating SP offset here may require the use of LDG to materialize
  // the tagged address of the stack slot, along with a scratch register
  // allocation (post-regalloc!).
  //
  // For now we do the safe thing here and require that the entire stack frame
  // is within range of the shortest of the unchecked instructions.
  unsigned FrameSize = 0;
  for (unsigned i = 0, e = MFI->getObjectIndexEnd(); i != e; ++i)
    FrameSize += MFI->getObjectSize(i);
  bool EntireFrameReachableFromSP = FrameSize < 0xf00;
  return !MFI->hasVarSizedObjects() && EntireFrameReachableFromSP;
}

void AArch64StackTaggingPreRA::uncheckUsesOf(unsigned TaggedReg, int FI) {
  for (auto UI = MRI->use_instr_begin(TaggedReg), E = MRI->use_instr_end();
       UI != E;) {
    MachineInstr *UseI = &*(UI++);
    if (isUncheckedLoadOrStoreOpcode(UseI->getOpcode())) {
      // FI operand is always the one before the immediate offset.
      unsigned OpIdx = TII->getLoadStoreImmIdx(UseI->getOpcode()) - 1;
      if (UseI->getOperand(OpIdx).isReg() &&
          UseI->getOperand(OpIdx).getReg() == TaggedReg) {
        UseI->getOperand(OpIdx).ChangeToFrameIndex(FI);
        UseI->getOperand(OpIdx).setTargetFlags(AArch64II::MO_TAGGED);
      }
    } else if (UseI->isCopy() &&
               Register::isVirtualRegister(UseI->getOperand(0).getReg())) {
      uncheckUsesOf(UseI->getOperand(0).getReg(), FI);
    }
  }
}

void AArch64StackTaggingPreRA::uncheckLoadsAndStores() {
  for (auto *I : ReTags) {
    unsigned TaggedReg = I->getOperand(0).getReg();
    int FI = I->getOperand(1).getIndex();
    uncheckUsesOf(TaggedReg, FI);
  }
}

bool AArch64StackTaggingPreRA::runOnMachineFunction(MachineFunction &Func) {
  MF = &Func;
  MRI = &MF->getRegInfo();
  AFI = MF->getInfo<AArch64FunctionInfo>();
  TII = static_cast<const AArch64InstrInfo *>(MF->getSubtarget().getInstrInfo());
  TRI = static_cast<const AArch64RegisterInfo *>(
      MF->getSubtarget().getRegisterInfo());
  MFI = &MF->getFrameInfo();
  ReTags.clear();

  assert(MRI->isSSA());

  LLVM_DEBUG(dbgs() << "********** AArch64 Stack Tagging PreRA **********\n"
                    << "********** Function: " << MF->getName() << '\n');

  SmallSetVector<int, 8> TaggedSlots;
  for (auto &BB : *MF) {
    for (auto &I : BB) {
      if (I.getOpcode() == AArch64::TAGPstack) {
        ReTags.push_back(&I);
        int FI = I.getOperand(1).getIndex();
        TaggedSlots.insert(FI);
        // There should be no offsets in TAGP yet.
        assert(I.getOperand(2).getImm() == 0);
      }
    }
  }

  if (ReTags.empty())
    return false;

  if (mayUseUncheckedLoadStore())
    uncheckLoadsAndStores();

  return true;
}