[RISCV] Add initial stack clash protection

rzinsly · rzinsly · commit 132e8598b2d0 · 2024-11-26T17:00:42.000-03:00
Enable `-fstack-clash-protection` for RISCV and stack probe for function
prologues.
We probe the stack by creating an unrolled loop that allocates and probe
the stack in ProbeSize chunks, this is not ideal if the loop has many
iterations.
diff --git a/clang/lib/Driver/ToolChains/Clang.cpp b/clang/lib/Driver/ToolChains/Clang.cpp
@@ -3774,7 +3774,8 @@ static void RenderSCPOptions(const ToolChain &TC, const ArgList &Args,
     return;
 
   if (!EffectiveTriple.isX86() && !EffectiveTriple.isSystemZ() &&
-      !EffectiveTriple.isPPC64() && !EffectiveTriple.isAArch64())
+      !EffectiveTriple.isPPC64() && !EffectiveTriple.isAArch64() &&
+      !EffectiveTriple.isRISCV())
     return;
 
   Args.addOptInFlag(CmdArgs, options::OPT_fstack_clash_protection,
diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp b/llvm/lib/Target/RISCV/RISCVFrameLowering.cpp
@@ -610,26 +610,74 @@ static MCCFIInstruction createDefCFAOffset(const TargetRegisterInfo &TRI,
                                         Comment.str());
 }
 
+// Allocate stack space and probe it if necessary.
 void RISCVFrameLowering::allocateStack(MachineBasicBlock &MBB,
                                        MachineBasicBlock::iterator MBBI,
-                                       MachineFunction &MF, StackOffset Offset,
-                                       uint64_t RealStackSize,
-                                       bool EmitCFI) const {
+                                       MachineFunction &MF, uint64_t Offset,
+                                       uint64_t RealStackSize, bool EmitCFI,
+                                       bool NeedProbe,
+                                       uint64_t ProbeSize) const {
   DebugLoc DL;
   const RISCVRegisterInfo *RI = STI.getRegisterInfo();
   const RISCVInstrInfo *TII = STI.getInstrInfo();
 
-  RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg, Offset, MachineInstr::FrameSetup,
-                getStackAlign());
+  // Simply allocate the stack if it's not big enough to require a probe.
+  if (!NeedProbe || Offset <= ProbeSize) {
+    RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg, StackOffset::getFixed(-Offset),
+                  MachineInstr::FrameSetup, getStackAlign());
 
-  if (EmitCFI) {
-    // Emit ".cfi_def_cfa_offset RealStackSize"
-    unsigned CFIIndex = MF.addFrameInst(
-		MCCFIInstruction::cfiDefCfaOffset(nullptr, RealStackSize));
-    BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
-        .addCFIIndex(CFIIndex)
-        .setMIFlag(MachineInstr::FrameSetup);
+    if (EmitCFI) {
+      // Emit ".cfi_def_cfa_offset RealStackSize"
+      unsigned CFIIndex = MF.addFrameInst(
+          MCCFIInstruction::cfiDefCfaOffset(nullptr, RealStackSize));
+      BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex)
+          .setMIFlag(MachineInstr::FrameSetup);
+    }
+
+    return;
   }
+
+  // Do an unrolled probe loop.
+  uint64_t CurrentOffset = 0;
+  bool IsRV64 = STI.is64Bit();
+  while (CurrentOffset + ProbeSize <= Offset) {
+    RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg,
+                  StackOffset::getFixed(-ProbeSize), MachineInstr::FrameSetup,
+                  getStackAlign());
+    // s[d|w] zero, 0(sp)
+    BuildMI(MBB, MBBI, DL, TII->get(IsRV64 ? RISCV::SD : RISCV::SW))
+        .addReg(RISCV::X0)
+        .addReg(SPReg)
+        .addImm(0)
+        .setMIFlags(MachineInstr::FrameSetup);
+
+    CurrentOffset += ProbeSize;
+    if (EmitCFI) {
+      // Emit ".cfi_def_cfa_offset CurrentOffset"
+      unsigned CFIIndex = MF.addFrameInst(
+          MCCFIInstruction::cfiDefCfaOffset(nullptr, CurrentOffset));
+      BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex)
+          .setMIFlag(MachineInstr::FrameSetup);
+    }
+  }
+
+  uint64_t Residual = Offset - CurrentOffset;
+  if (Residual) {
+    RI->adjustReg(MBB, MBBI, DL, SPReg, SPReg, StackOffset::getFixed(-Residual),
+                  MachineInstr::FrameSetup, getStackAlign());
+    if (EmitCFI) {
+      // Emit ".cfi_def_cfa_offset Offset"
+      unsigned CFIIndex =
+          MF.addFrameInst(MCCFIInstruction::cfiDefCfaOffset(nullptr, Offset));
+      BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
+          .addCFIIndex(CFIIndex)
+          .setMIFlag(MachineInstr::FrameSetup);
+    }
+  }
+
+  return;
 }
 
 void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
@@ -746,11 +794,14 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
                                           getPushOrLibCallsSavedInfo(MF, CSI));
   }
 
-  if (StackSize != 0) {
-    // Allocate space on the stack if necessary.
-    allocateStack(MBB, MBBI, MF, StackOffset::getFixed(-StackSize),
-		  RealStackSize, /*EmitCFI=*/ true);
-  }
+  // Allocate space on the stack if necessary.
+  auto &Subtarget = MF.getSubtarget<RISCVSubtarget>();
+  const RISCVTargetLowering *TLI = Subtarget.getTargetLowering();
+  bool NeedProbe = TLI->hasInlineStackProbe(MF);
+  uint64_t ProbeSize = TLI->getStackProbeSize(MF, getStackAlign());
+  if (StackSize != 0)
+    allocateStack(MBB, MBBI, MF, StackSize, RealStackSize, /*EmitCFI=*/true,
+                  NeedProbe, ProbeSize);
 
   // The frame pointer is callee-saved, and code has been generated for us to
   // save it to the stack. We need to skip over the storing of callee-saved
@@ -791,8 +842,9 @@ void RISCVFrameLowering::emitPrologue(MachineFunction &MF,
     assert(SecondSPAdjustAmount > 0 &&
            "SecondSPAdjustAmount should be greater than zero");
 
-    allocateStack(MBB, MBBI, MF, StackOffset::getFixed(-SecondSPAdjustAmount),
-                  getStackSizeWithRVVPadding(MF), !hasFP(MF));
+    allocateStack(MBB, MBBI, MF, SecondSPAdjustAmount,
+                  getStackSizeWithRVVPadding(MF), !hasFP(MF), NeedProbe,
+                  ProbeSize);
   }
 
   if (RVVStackSize) {
diff --git a/llvm/lib/Target/RISCV/RISCVFrameLowering.h b/llvm/lib/Target/RISCV/RISCVFrameLowering.h
@@ -79,8 +79,9 @@ class RISCVFrameLowering : public TargetFrameLowering {
   }
 
   void allocateStack(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
-                     MachineFunction &MF, StackOffset Offset,
-                     uint64_t RealStackSize, bool EmitCFI) const;
+                     MachineFunction &MF, uint64_t Offset,
+                     uint64_t RealStackSize, bool EmitCFI, bool NeedProbe,
+                     uint64_t ProbeSize) const;
 
 protected:
   const RISCVSubtarget &STI;
diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -22156,3 +22156,25 @@ namespace llvm::RISCVVIntrinsicsTable {
 #include "RISCVGenSearchableTables.inc"
 
 } // namespace llvm::RISCVVIntrinsicsTable
+
+bool RISCVTargetLowering::hasInlineStackProbe(const MachineFunction &MF) const {
+
+  // If the function specifically requests inline stack probes, emit them.
+  if (MF.getFunction().hasFnAttribute("probe-stack"))
+    return MF.getFunction().getFnAttribute("probe-stack").getValueAsString() ==
+           "inline-asm";
+
+  return false;
+}
+
+unsigned RISCVTargetLowering::getStackProbeSize(const MachineFunction &MF,
+                                                Align StackAlign) const {
+  // The default stack probe size is 4096 if the function has no
+  // stack-probe-size attribute.
+  const Function &Fn = MF.getFunction();
+  unsigned StackProbeSize =
+      Fn.getFnAttributeAsParsedInteger("stack-probe-size", 4096);
+  // Round down to the stack alignment.
+  StackProbeSize = alignDown(StackProbeSize, StackAlign.value());
+  return StackProbeSize ? StackProbeSize : StackAlign.value();
+}
diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.h b/llvm/lib/Target/RISCV/RISCVISelLowering.h
@@ -923,6 +923,12 @@ class RISCVTargetLowering : public TargetLowering {
                               MachineBasicBlock::instr_iterator &MBBI,
                               const TargetInstrInfo *TII) const override;
 
+  /// True if stack clash protection is enabled for this functions.
+  bool hasInlineStackProbe(const MachineFunction &MF) const override;
+
+  unsigned getStackProbeSize(const MachineFunction &MF,
+                             Align StackAlign) const;
+
 private:
   void analyzeInputArgs(MachineFunction &MF, CCState &CCInfo,
                         const SmallVectorImpl<ISD::InputArg> &Ins, bool IsRet,
diff --git a/llvm/lib/Target/RISCV/RISCVMachineFunctionInfo.cpp b/llvm/lib/Target/RISCV/RISCVMachineFunctionInfo.cpp
@@ -11,6 +11,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "RISCVMachineFunctionInfo.h"
+#include "llvm/IR/Module.h"
 
 using namespace llvm;
 
@@ -26,6 +27,35 @@ MachineFunctionInfo *RISCVMachineFunctionInfo::clone(
   return DestMF.cloneInfo<RISCVMachineFunctionInfo>(*this);
 }
 
+RISCVMachineFunctionInfo::RISCVMachineFunctionInfo(const Function &F,
+                                                   const RISCVSubtarget *STI) {
+
+  // The default stack probe size is 4096 if the function has no
+  // stack-probe-size attribute. This is a safe default because it is the
+  // smallest possible guard page size.
+  uint64_t ProbeSize = 4096;
+  if (F.hasFnAttribute("stack-probe-size"))
+    ProbeSize = F.getFnAttributeAsParsedInteger("stack-probe-size");
+  else if (const auto *PS = mdconst::extract_or_null<ConstantInt>(
+               F.getParent()->getModuleFlag("stack-probe-size")))
+    ProbeSize = PS->getZExtValue();
+  assert(int64_t(ProbeSize) > 0 && "Invalid stack probe size");
+
+  // Round down to the stack alignment.
+  uint64_t StackAlign =
+      STI->getFrameLowering()->getTransientStackAlign().value();
+  ProbeSize = std::max(StackAlign, alignDown(ProbeSize, StackAlign));
+  StringRef ProbeKind;
+  if (F.hasFnAttribute("probe-stack"))
+    ProbeKind = F.getFnAttribute("probe-stack").getValueAsString();
+  else if (const auto *PS = dyn_cast_or_null<MDString>(
+               F.getParent()->getModuleFlag("probe-stack")))
+    ProbeKind = PS->getString();
+  if (ProbeKind.size()) {
+    StackProbeSize = ProbeSize;
+  }
+}
+
 void yaml::RISCVMachineFunctionInfo::mappingImpl(yaml::IO &YamlIO) {
   MappingTraits<RISCVMachineFunctionInfo>::mapping(YamlIO, *this);
 }
diff --git a/llvm/lib/Target/RISCV/RISCVMachineFunctionInfo.h b/llvm/lib/Target/RISCV/RISCVMachineFunctionInfo.h
@@ -76,8 +76,10 @@ class RISCVMachineFunctionInfo : public MachineFunctionInfo {
   unsigned RVPushRegs = 0;
   int RVPushRlist = llvm::RISCVZC::RLISTENCODE::INVALID_RLIST;
 
+  int64_t StackProbeSize = 0;
+
 public:
-  RISCVMachineFunctionInfo(const Function &F, const TargetSubtargetInfo *STI) {}
+  RISCVMachineFunctionInfo(const Function &F, const RISCVSubtarget *STI);
 
   MachineFunctionInfo *
   clone(BumpPtrAllocator &Allocator, MachineFunction &DestMF,
@@ -157,6 +159,9 @@ class RISCVMachineFunctionInfo : public MachineFunctionInfo {
 
   bool isVectorCall() const { return IsVectorCall; }
   void setIsVectorCall() { IsVectorCall = true; }
+
+  bool hasStackProbing() const { return StackProbeSize != 0; }
+  int64_t getStackProbeSize() const { return StackProbeSize; }
 };
 
 } // end namespace llvm
diff --git a/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp b/llvm/lib/Target/RISCV/RISCVTargetMachine.cpp
@@ -266,8 +266,8 @@ RISCVTargetMachine::getSubtargetImpl(const Function &F) const {
 MachineFunctionInfo *RISCVTargetMachine::createMachineFunctionInfo(
     BumpPtrAllocator &Allocator, const Function &F,
     const TargetSubtargetInfo *STI) const {
-  return RISCVMachineFunctionInfo::create<RISCVMachineFunctionInfo>(Allocator,
-                                                                    F, STI);
+  return RISCVMachineFunctionInfo::create<RISCVMachineFunctionInfo>(
+      Allocator, F, static_cast<const RISCVSubtarget *>(STI));
 }
 
 TargetTransformInfo
diff --git a/llvm/test/CodeGen/RISCV/stack-clash-prologue-nounwind.ll b/llvm/test/CodeGen/RISCV/stack-clash-prologue-nounwind.ll
@@ -0,0 +1,72 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py
+; RUN: llc -mtriple=riscv64 -mattr=+m -O2 < %s \
+; RUN:   | FileCheck %s -check-prefix=RV64I
+; RUN: llc -mtriple=riscv32 -mattr=+m -O2 < %s \
+; RUN:   | FileCheck %s -check-prefix=RV32I
+
+; Tests copied from PowerPC.
+
+; Free probe
+define i8 @f0() #0 nounwind {
+; RV64I-LABEL: f0:
+; RV64I:       # %bb.0: # %entry
+; RV64I-NEXT:    addi sp, sp, -64
+; RV64I-NEXT:    li a0, 3
+; RV64I-NEXT:    sb a0, 0(sp)
+; RV64I-NEXT:    lbu a0, 0(sp)
+; RV64I-NEXT:    addi sp, sp, 64
+; RV64I-NEXT:    ret
+;
+; RV32I-LABEL: f0:
+; RV32I:       # %bb.0: # %entry
+; RV32I-NEXT:    addi sp, sp, -64
+; RV32I-NEXT:    li a0, 3
+; RV32I-NEXT:    sb a0, 0(sp)
+; RV32I-NEXT:    lbu a0, 0(sp)
+; RV32I-NEXT:    addi sp, sp, 64
+; RV32I-NEXT:    ret
+entry:
+  %a = alloca i8, i64 64
+  %b = getelementptr inbounds i8, ptr %a, i64 63
+  store volatile i8 3, ptr %a
+  %c = load volatile i8, ptr %a
+  ret i8 %c
+}
+
+define i8 @f1() #0 nounwind {
+; RV64I-LABEL: f1:
+; RV64I:       # %bb.0: # %entry
+; RV64I-NEXT:    lui a0, 1
+; RV64I-NEXT:    sub sp, sp, a0
+; RV64I-NEXT:    sd zero, 0(sp)
+; RV64I-NEXT:    addi sp, sp, -16
+; RV64I-NEXT:    li a0, 3
+; RV64I-NEXT:    sb a0, 16(sp)
+; RV64I-NEXT:    lbu a0, 16(sp)
+; RV64I-NEXT:    lui a1, 1
+; RV64I-NEXT:    addiw a1, a1, 16
+; RV64I-NEXT:    add sp, sp, a1
+; RV64I-NEXT:    ret
+;
+; RV32I-LABEL: f1:
+; RV32I:       # %bb.0: # %entry
+; RV32I-NEXT:    lui a0, 1
+; RV32I-NEXT:    sub sp, sp, a0
+; RV32I-NEXT:    sw zero, 0(sp)
+; RV32I-NEXT:    addi sp, sp, -16
+; RV32I-NEXT:    li a0, 3
+; RV32I-NEXT:    sb a0, 16(sp)
+; RV32I-NEXT:    lbu a0, 16(sp)
+; RV32I-NEXT:    lui a1, 1
+; RV32I-NEXT:    addi a1, a1, 16
+; RV32I-NEXT:    add sp, sp, a1
+; RV32I-NEXT:    ret
+entry:
+  %a = alloca i8, i64 4096
+  %b = getelementptr inbounds i8, ptr %a, i64 63
+  store volatile i8 3, ptr %a
+  %c = load volatile i8, ptr %a
+  ret i8 %c
+}
+
+attributes #0 = { "probe-stack"="inline-asm" }
diff --git a/llvm/test/CodeGen/RISCV/stack-clash-prologue.ll b/llvm/test/CodeGen/RISCV/stack-clash-prologue.ll
