Add support for probestack on ARM targets

[ryan-summers]

Summary

There is a desire to add support for probestack functionality on ARM (32-bit) targets for embedded applications, as the stack probe functionality can be used, along with a memory protection unit, to protect bare-metal applications from stack overflow issues (or to detect these issues so they can be acted upon).

See https://github.com/rust-embedded/wg/issues/460 for more information as to why this is useful

This diff provides the ability to use function annotations for the probestack and probe-stack-size to define the various probestack stubs and page protection size. This patch does not enable any of this functionality in clang, and is only targeting the LLVM library currently (as the final usecase for this patch is intended to be in rustc).

Implementation is largely based off the X86 implementation of stack probes. As such, I have tagged some of the reviewers from the original review (https://reviews.llvm.org/D68720) for this diff.

[ryan-summers]

llvm/lib/Target/ARM/ARMFrameLowering.cpp

591

The no-stack-arg-probe FnAttribute is no longer present for Windows, which isn't desirable as the removal would imply a break in backwards compatibility.

765

It looks like this unintentionally is removing the check for WIndows stack probe requirements. I'll need to reinclude this condition

967

Similar to the above, this changes the conditional for Windows requiring stack probes, which it shouldn't.

2301

Same as earlier - we need to preserve stack probe emission for Windows

llvm/lib/Target/ARM/ARMFrameLowering.cpp

static int sizeOfSPAdjustment(const MachineInstr &MI) {
   int count = 0;
   // ARM and Thumb2 push/pop insts have explicit "sp, sp" operands (+
   // pred) so the list starts at 4.
   for (int i = MI.getNumOperands() - 1; i >= 4; --i)
     count += RegSize;
   return count;
 }
 
-static bool WindowsRequiresStackProbe(const MachineFunction &MF,
-                                      size_t StackSizeInBytes) {
-  const MachineFrameInfo &MFI = MF.getFrameInfo();
-  const Function &F = MF.getFunction();
-  unsigned StackProbeSize = (MFI.getStackProtectorIndex() > 0) ? 4080 : 4096;
-
-  StackProbeSize =
-      F.getFnAttributeAsParsedInteger("stack-probe-size", StackProbeSize);
-  return (StackSizeInBytes >= StackProbeSize) &&
-         !F.hasFnAttribute("no-stack-arg-probe");
-}
-
 namespace {
 
 struct StackAdjustingInsts {
   struct InstInfo {
     MachineBasicBlock::iterator I;
     unsigned SPAdjust;
     bool BeforeFPSet;
   };
void ARMFrameLowering::emitPrologue(MachineFunction &MF,
   bool isARM = !AFI->isThumbFunction();
   Align Alignment = STI.getFrameLowering()->getStackAlign();
   unsigned ArgRegsSaveSize = AFI->getArgRegsSaveSize();
   unsigned NumBytes = MFI.getStackSize();
   const std::vector<CalleeSavedInfo> &CSI = MFI.getCalleeSavedInfo();
   int FPCXTSaveSize = 0;
   bool NeedsWinCFI = needsWinCFI(MF);
 
+  const bool EmitStackProbeCall =
+      STI.getTargetLowering()->hasStackProbeSymbol(MF);
+  unsigned StackProbeSize = STI.getTargetLowering()->getStackProbeSize(MF);
+
   // Debug location must be unknown since the first debug location is used
   // to determine the end of the prologue.
   DebugLoc dl;
 
   Register FramePtr = RegInfo->getFrameRegister(MF);
 
   // Determine the sizes of each callee-save spill areas and record which frame
   // belongs to which callee-save spill areas.
   unsigned GPRCS1Size = 0, GPRCS2Size = 0, DPRCSSize = 0;
   int FramePtrSpillFI = 0;
   int D8SpillFI = 0;
 
   // All calls are tail calls in GHC calling conv, and functions have no
   // prologue/epilogue.
   if (MF.getFunction().getCallingConv() == CallingConv::GHC)
     return;
 
   StackAdjustingInsts DefCFAOffsetCandidates;
   bool HasFP = hasFP(MF);
 
   if (!AFI->hasStackFrame() &&
-      (!STI.isTargetWindows() || !WindowsRequiresStackProbe(MF, NumBytes))) {
+      !(EmitStackProbeCall && NumBytes >= StackProbeSize)) {
     if (NumBytes != 0) {
       emitSPUpdate(isARM, MBB, MBBI, dl, TII, -NumBytes,
                    MachineInstr::FrameSetup);
       DefCFAOffsetCandidates.addInst(std::prev(MBBI), NumBytes, true);
     }
     if (!NeedsWinCFI)
       DefCFAOffsetCandidates.emitDefCFAOffsets(MBB, dl, TII, HasFP);
     if (NeedsWinCFI && MBBI != MBB.begin()) {
if (GPRCS2Size > 0 && STI.splitFramePointerPush(MF)) {
     GPRCS2Push = LastPush = MBBI++;
     DefCFAOffsetCandidates.addInst(LastPush, GPRCS2Size);
   }
 
   bool NeedsWinCFIStackAlloc = NeedsWinCFI;
   if (STI.splitFramePointerPush(MF) && HasFP)
     NeedsWinCFIStackAlloc = false;
 
-  if (STI.isTargetWindows() && WindowsRequiresStackProbe(MF, NumBytes)) {
+  if (NumBytes >= StackProbeSize && EmitStackProbeCall) {
     uint32_t NumWords = NumBytes >> 2;
 
     if (NumWords < 65536) {
       BuildMI(MBB, MBBI, dl, TII.get(ARM::t2MOVi16), ARM::R4)
           .addImm(NumWords)
           .setMIFlags(MachineInstr::FrameSetup)
           .add(predOps(ARMCC::AL));
     } else {
       // Split into two instructions here, instead of using t2MOVi32imm,
       // to allow inserting accurate SEH instructions (including accurate
       // instruction size for each of them).
       BuildMI(MBB, MBBI, dl, TII.get(ARM::t2MOVi16), ARM::R4)
           .addImm(NumWords & 0xffff)
           .setMIFlags(MachineInstr::FrameSetup)
           .add(predOps(ARMCC::AL));
       BuildMI(MBB, MBBI, dl, TII.get(ARM::t2MOVTi16), ARM::R4)
           .addReg(ARM::R4)
           .addImm(NumWords >> 16)
           .setMIFlags(MachineInstr::FrameSetup)
           .add(predOps(ARMCC::AL));
     }
 
+    StringRef Symbol = STI.getTargetLowering()->getStackProbeSymbolName(MF);
+
     switch (TM.getCodeModel()) {
     case CodeModel::Tiny:
       llvm_unreachable("Tiny code model not available on ARM.");
     case CodeModel::Small:
     case CodeModel::Medium:
     case CodeModel::Kernel:
       BuildMI(MBB, MBBI, dl, TII.get(ARM::tBL))
           .add(predOps(ARMCC::AL))
-          .addExternalSymbol("__chkstk")
+          .addExternalSymbol(MF.createExternalSymbolName(Symbol))
           .addReg(ARM::R4, RegState::Implicit)
           .setMIFlags(MachineInstr::FrameSetup);
       break;
     case CodeModel::Large:
       BuildMI(MBB, MBBI, dl, TII.get(ARM::t2MOVi32imm), ARM::R12)
-        .addExternalSymbol("__chkstk")
+          .addExternalSymbol(MF.createExternalSymbolName(Symbol))
           .setMIFlags(MachineInstr::FrameSetup);
 
       BuildMI(MBB, MBBI, dl, TII.get(ARM::tBLXr))
           .add(predOps(ARMCC::AL))
           .addReg(ARM::R12, RegState::Kill)
           .addReg(ARM::R4, RegState::Implicit)
           .setMIFlags(MachineInstr::FrameSetup);
       break;
     }
const ARMBaseInstrInfo &TII =
       *static_cast<const ARMBaseInstrInfo *>(MF.getSubtarget().getInstrInfo());
   ARMFunctionInfo *AFI = MF.getInfo<ARMFunctionInfo>();
   MachineFrameInfo &MFI = MF.getFrameInfo();
   MachineRegisterInfo &MRI = MF.getRegInfo();
   const TargetRegisterInfo *TRI = MF.getSubtarget().getRegisterInfo();
   (void)TRI;  // Silence unused warning in non-assert builds.
   Register FramePtr = RegInfo->getFrameRegister(MF);
 
+  const bool EmitStackProbeCall =
+      STI.getTargetLowering()->hasStackProbeSymbol(MF);
+  unsigned StackProbeSize = STI.getTargetLowering()->getStackProbeSize(MF);
+
   // Spill R4 if Thumb2 function requires stack realignment - it will be used as
   // scratch register. Also spill R4 if Thumb2 function has varsized objects,
   // since it's not always possible to restore sp from fp in a single
   // instruction.
   // FIXME: It will be better just to find spare register here.
   if (AFI->isThumb2Function() &&
       (MFI.hasVarSizedObjects() || RegInfo->hasStackRealignment(MF)))
     SavedRegs.set(ARM::R4);
 
   // If a stack probe will be emitted, spill R4 and LR, since they are
   // clobbered by the stack probe call.
   // This estimate should be a safe, conservative estimate. The actual
   // stack probe is enabled based on the size of the local objects;
   // this estimate also includes the varargs store size.
-  if (STI.isTargetWindows() &&
-      WindowsRequiresStackProbe(MF, MFI.estimateStackSize(MF))) {
+  if (MFI.estimateStackSize(MF) >= StackProbeSize && EmitStackProbeCall) {
     SavedRegs.set(ARM::R4);
     SavedRegs.set(ARM::LR);
   }
 
   if (AFI->isThumb1OnlyFunction()) {
     // Spill LR if Thumb1 function uses variable length argument lists.
     if (AFI->getArgRegsSaveSize() > 0)
       SavedRegs.set(ARM::LR);

llvm/lib/Target/ARM/ARMISelLowering.h

public:
     bool convertSetCCLogicToBitwiseLogic(EVT VT) const override {
       return VT.isScalarInteger();
     }
 
     bool supportSwiftError() const override {
       return true;
     }
 
+    bool hasStackProbeSymbol(const MachineFunction &MF) const override;
+    StringRef getStackProbeSymbolName(const MachineFunction &MF) const override;
+    unsigned getStackProbeSize(const MachineFunction &MF) const;
+
     bool hasStandaloneRem(EVT VT) const override {
       return HasStandaloneRem;
     }
 
     ShiftLegalizationStrategy
     preferredShiftLegalizationStrategy(SelectionDAG &DAG, SDNode *N,
                                        unsigned ExpansionFactor) const override;
 

llvm/lib/Target/ARM/ARMISelLowering.cpp

if (!ConstRotation)
       return nullptr; // Invalid rotation for arm_mve_vcaddq
 
     return B.CreateIntrinsic(Intrinsic::arm_mve_vcaddq, Ty,
                              {ConstHalving, ConstRotation, InputA, InputB});
   }
 
   return nullptr;
 }
+
+/// Returns true if stack probing through a function call is requested.
+bool ARMTargetLowering::hasStackProbeSymbol(const MachineFunction &MF) const {
+  return !getStackProbeSymbolName(MF).empty();
+}
+
+/// Returns the name of the symbol used to emit stack probes or the empty
+/// string if not applicable.
+StringRef
+ARMTargetLowering::getStackProbeSymbolName(const MachineFunction &MF) const {
+  // If the function specifically requests stack probes, emit them.
+  if (MF.getFunction().hasFnAttribute("probe-stack"))
+    return MF.getFunction().getFnAttribute("probe-stack").getValueAsString();
+
+  if (Subtarget->isTargetWindows() &&
+      !MF.getFunction().hasFnAttribute("no-stack-arg-probe"))
+    return "__chkstk";
+
+  // Generally, if we aren't on Windows, the platform ABI does not include
+  // support for stack probes, so don't emit them.
+  return "";
+}
+
+unsigned ARMTargetLowering::getStackProbeSize(const MachineFunction &MF) const {
+  // The default stack probe size is 4096 if the function has no stackprobesize
+  // attribute.
+  const MachineFrameInfo &MFI = MF.getFrameInfo();
+
+  unsigned StackProbeSize = 4096;
+
+  if (Subtarget->isTargetWindows())
+    StackProbeSize = (MFI.getStackProtectorIndex() > 0) ? 4080 : 4096;
+
+  const Function &Fn = MF.getFunction();
+  if (Fn.hasFnAttribute("stack-probe-size"))
+    Fn.getFnAttribute("stack-probe-size")
+        .getValueAsString()
+        .getAsInteger(0, StackProbeSize);
+  return StackProbeSize;
+}

llvm/test/CodeGen/ARM/stack-probe-size.ll

+; This tests that the probe-stack attribute properly calls the probe-stack
+; function.
+;
+; RUN: llc < %s -mtriple=arm-none  | FileCheck %s
+
+define void @_test1() "probe-stack"="__probestack" "stack-probe-size"="0" {
+    %array = alloca [1024 x i8], align 16
+    ret void
+
+; CHECK-LABEL:      test1:
+; CHECK:            movw r4, #256
+; CHECK-NEXT:       bl __probestack
+; CHECK-NEXT:       sub.w sp, sp, r4
+}
+
+define void @_test2() "probe-stack"="__probestack" "stack-probe-size"="1024" {
+    %array = alloca [1024 x i8], align 16
+    ret void
+
+; CHECK-LABEL:      test2:
+; CHECK:            movw r4, #256
+; CHECK-NEXT:       bl __probestack
+; CHECK-NEXT:       sub.w sp, sp, r4
+}
+
+define void @_test3() "probe-stack"="__probestack" "stack-probe-size"="8192" {
+    %array = alloca [8192 x i8], align 16
+    ret void
+
+; CHECK-LABEL:      test3:
+; CHECK:            movw r4, #2048
+; CHECK-NEXT:       bl __probestack
+; CHECK-NEXT:       sub.w sp, sp, r4
+}
+
+define void @_test4() "probe-stack"="__probestack" "stack-probe-size"="8192" {
+    %array = alloca [4096 x i8], align 16
+    ret void
+
+; CHECK-LABEL:      test4:
+; CHECK-NOT:        bl __probestack
+}

llvm/test/CodeGen/ARM/stack-probe.ll

+; This tests that the probe-stack attribute properly calls the probe-stack
+; function.
+;
+; RUN: llc < %s -mtriple=arm-none  | FileCheck %s
+
+define void @_test1() "probe-stack"="__probestack" {
+    %array = alloca [4096 x i8], align 16
+    ret void
+
+; CHECK-LABEL:      test1:
+; CHECK:            movw r4, #1024
+; CHECK-NEXT:       bl __probestack
+; CHECK-NEXT:       sub.w sp, sp, r4
+}