Define behavior of "stack-probe-size" attribute when inlining.

[whitequark]

Summary

Also document the attribute, since "probe-stack" already is.

[majnemer]

docs/LangRef.rst

1515 ↗

(On Diff #103624)

I believe it is also respected by stack probes which don't utilize the probe-stack attribute.

lib/IR/Attributes.cpp

1649–1665 ↗

(On Diff #103624)

I think this should be in its own adjustCallerStackProbeSize (or something similar) merger.

[whitequark]

Addressed review comments

[majnemer]

LGTM

llvm/trunk/docs/LangRef.rst

``readonly``
     called with the same set of arguments and global state.  This means while it
     cannot unwind exceptions by calling the ``C++`` exception throwing methods
     (since they write to memory), there may be non-``C++`` mechanisms that throw
     exceptions without writing to LLVM visible memory.
 
     On an argument, this attribute indicates that the function does not write
     through this pointer argument, even though it may write to the memory that
     the pointer points to.
+``"stack-probe-size"``
+    This attribute controls the behavior of stack probes: either
+    the ``"probe-stack"`` attribute, or ABI-required stack probes, if any.
+    It defines the size of the guard region. It ensures that if the function
+    may use more stack space than the size of the guard region, stack probing
+    sequence will be emitted. It takes one required integer value, which
+    is 4096 by default.
+
+    If a function that has a ``"stack-probe-size"`` attribute is inlined into
+    a function with another ``"stack-probe-size"`` attribute, the resulting
+    function has the ``"stack-probe-size"`` attribute that has the lower
+    numeric value. If a function that has a ``"stack-probe-size"`` attribute is
+    inlined into a function that has no ``"stack-probe-size"`` attribute
+    at all, the resulting function has the ``"stack-probe-size"`` attribute
+    of the callee.
 ``writeonly``
     On a function, this attribute indicates that the function may write to but
     does not read from memory.
 
     On an argument, this attribute indicates that the function may write to but
     does not read through this pointer argument (even though it may read from
     the memory that the pointer points to).
 ``argmemonly``

llvm/trunk/include/llvm/IR/Attributes.td

 def : MergeRule<"setAND<LessPreciseFPMADAttr>">;
 def : MergeRule<"setAND<NoInfsFPMathAttr>">;
 def : MergeRule<"setAND<NoNansFPMathAttr>">;
 def : MergeRule<"setAND<UnsafeFPMathAttr>">;
 def : MergeRule<"setOR<NoImplicitFloatAttr>">;
 def : MergeRule<"setOR<NoJumpTablesAttr>">;
 def : MergeRule<"adjustCallerSSPLevel">;
 def : MergeRule<"adjustCallerStackProbes">;
+def : MergeRule<"adjustCallerStackProbeSize">;

llvm/trunk/lib/IR/Attributes.cpp

Caller.addFnAttr(Attribute::StackProtectStrong);
              !Caller.hasFnAttribute(Attribute::StackProtectReq) &&
              !Caller.hasFnAttribute(Attribute::StackProtectStrong))
     Caller.addFnAttr(Attribute::StackProtect);
 }
 
 /// \brief If the inlined function required stack probes, then ensure that
 /// the calling function has those too.
 static void adjustCallerStackProbes(Function &Caller, const Function &Callee) {
-  if (!Caller.hasFnAttribute("probe-stack") && Callee.hasFnAttribute("probe-stack"))
-    Caller.addFnAttr("probe-stack", Callee.getFnAttribute("probe-stack").getValueAsString());
+  if (!Caller.hasFnAttribute("probe-stack") &&
+      Callee.hasFnAttribute("probe-stack")) {
+    Caller.addFnAttr(Callee.getFnAttribute("probe-stack"));
+  }
+}
+
+/// \brief If the inlined function defines the size of guard region
+/// on the stack, then ensure that the calling function defines a guard region
+/// that is no larger.
+static void
+adjustCallerStackProbeSize(Function &Caller, const Function &Callee) {
+  if (Callee.hasFnAttribute("stack-probe-size")) {
+    uint64_t CalleeStackProbeSize;
+    Callee.getFnAttribute("stack-probe-size")
+          .getValueAsString()
+          .getAsInteger(0, CalleeStackProbeSize);
+    if (Caller.hasFnAttribute("stack-probe-size")) {
+      uint64_t CallerStackProbeSize;
+      Caller.getFnAttribute("stack-probe-size")
+            .getValueAsString()
+            .getAsInteger(0, CallerStackProbeSize);
+      if (CallerStackProbeSize > CalleeStackProbeSize) {
+        Caller.addFnAttr(Callee.getFnAttribute("stack-probe-size"));
+      }
+    } else {
+      Caller.addFnAttr(Callee.getFnAttribute("stack-probe-size"));
+    }
+  }
 }
 
 #define GET_ATTR_COMPAT_FUNC
 #include "AttributesCompatFunc.inc"
 
 bool AttributeFuncs::areInlineCompatible(const Function &Caller,
                                          const Function &Callee) {
   return hasCompatibleFnAttrs(Caller, Callee);
 }
 
 void AttributeFuncs::mergeAttributesForInlining(Function &Caller,
                                                 const Function &Callee) {
   mergeFnAttrs(Caller, Callee);
 }

llvm/trunk/test/Transforms/Inline/inline-stack-probe-size.ll

+; RUN: opt %s -inline -S | FileCheck %s
+
+define internal void @innerSmall() "stack-probe-size"="4096" {
+  ret void
+}
+
+define internal void @innerLarge() "stack-probe-size"="8192" {
+  ret void
+}
+
+define void @outerNoAttribute() {
+  call void @innerSmall()
+  ret void
+}
+
+define void @outerConflictingAttributeSmall() "stack-probe-size"="4096" {
+  call void @innerLarge()
+  ret void
+}
+
+define void @outerConflictingAttributeLarge() "stack-probe-size"="8192" {
+  call void @innerSmall()
+  ret void
+}
+
+; CHECK: define void @outerNoAttribute() #0
+; CHECK: define void @outerConflictingAttributeSmall() #0
+; CHECK: define void @outerConflictingAttributeLarge() #0
+; CHECK: attributes #0 = { "stack-probe-size"="4096" }