1 files changed, 30 insertions, 0 deletions
diff --git a/llvm/lib/Target/RISCV/RISCVMachineFunctionInfo.cpp b/llvm/lib/Target/RISCV/RISCVMachineFunctionInfo.cpp
index d0c363042f51..a0d793176381 100644
--- 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);
 }