[CodeGen] Clear the registers the stack clear leaves data in

llvm/include/llvm/IR/Function.h

@@ -517,6 +517,21 @@ class LLVM_ABI Function : public GlobalObject, public ilist_node<Function> {
     return AttributeSets.getParamNoFPClass(ArgNo);
   }
 
+  /// The "zeroize-stack" attribute: the function's undertaking to clear its own
+  /// stack frame before it returns. Named once here so the enforcement sites
+  /// spread across CodeGen, the Verifier and the IPO passes cannot drift, and a
+  /// misspelling fails to compile rather than silently dropping the guarantee.
+  static constexpr StringRef ZeroizeStackAttrName = "zeroize-stack";
+
+  /// Return true if the function carries the "zeroize-stack" attribute.
+  bool hasZeroizeStack() const { return hasFnAttribute(ZeroizeStackAttrName); }
+
+  /// Return the "zeroize-stack" mode, or an empty string if the attribute is
+  /// absent. An empty or unrecognized value means the widest mode; see LangRef.
+  StringRef getZeroizeStackMode() const {
+    return getFnAttribute(ZeroizeStackAttrName).getValueAsString();
+  }
+
   /// Determine if the function is presplit coroutine.
   bool isPresplitCoroutine() const {
     return hasFnAttribute(Attribute::PresplitCoroutine);

llvm/lib/CodeGen/Analysis.cpp

@@ -538,6 +538,16 @@ static bool nextRealType(SmallVectorImpl<Type *> &SubTypes,
 /// This function only tests target-independent requirements.
 bool llvm::isInTailCallPosition(const CallBase &Call, const TargetMachine &TM,
                                 bool ReturnsFirstArg) {
+  // A tail call replaces the caller's frame and jumps away, so a function that
+  // promised to clear its frame before returning never gets to. Suppress it at
+  // this one target-independent point, which SelectionDAGBuilder, FastISel,
+  // GlobalISel, and memcpy/memmove/memset folding all reach. musttail is
+  // rejected in the Verifier (a caller cannot drop it); suppressing it here too
+  // is a backstop for unverified IR, failing closed into the backend's musttail
+  // error rather than a protected function that keeps its frame.
+  if (Call.getCaller()->hasZeroizeStack())
+    return false;
+
   const BasicBlock *ExitBB = Call.getParent();
   const Instruction *Term = ExitBB->getTerminator();
   const ReturnInst *Ret = dyn_cast<ReturnInst>(Term);

llvm/lib/CodeGen/GlobalISel/LegalizerHelper.cpp

@@ -534,6 +534,13 @@ static bool isLibCallInTailPosition(const CallLowering::ArgInfo &Result,
   MachineBasicBlock &MBB = *MI.getParent();
   const Function &F = MBB.getParent()->getFunction();
 
+  // A protected function cannot clear its frame after a tail call. This is the
+  // GlobalISel analog of the check in TargetLowering::isInTailCallPosition: a
+  // legalizer libcall folded into a tail call replaces the frame just the same,
+  // and is refused for the same reason.
+  if (F.hasZeroizeStack())
+    return false;
+
   // Conservatively require the attributes of the call to match those of
   // the return. Ignore NoAlias and NonNull because they don't affect the
   // call sequence.

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

@@ -66,6 +66,12 @@ bool TargetLowering::isInTailCallPosition(SelectionDAG &DAG, SDNode *Node,
   if (F.getFnAttribute("disable-tail-calls").getValueAsBool())
     return false;
 
+  // A protected function cannot clear its frame after a tail call. This is a
+  // legalizer libcall being folded into one, refused for the same reason as a
+  // tail call in the IR (see isInTailCallPosition in Analysis.cpp).
+  if (F.hasZeroizeStack())
+    return false;
+
   // Conservatively require the attributes of the call to match those of
   // the return. Ignore following attributes because they don't affect the
   // call sequence.

llvm/lib/IR/Attributes.cpp

@@ -2621,12 +2621,12 @@ static bool checkZeroizeStack(const Function &Caller, const Function &Callee) {
   // frame would have sat below the stack pointer at the caller's return, where
   // no clear reaches it, and inlining turns those bytes into frame bytes the
   // caller does clear.
-  if (!Callee.hasFnAttribute("zeroize-stack"))
+  if (!Callee.hasZeroizeStack())
     return true;
 
   // Otherwise the caller has to carry the attribute too, or there is no clear
   // for the callee's frame bytes to be folded into.
-  if (!Caller.hasFnAttribute("zeroize-stack"))
+  if (!Caller.hasZeroizeStack())
     return false;
 
   // Both are protected, so the caller must not ask for less of its frame than
@@ -2637,10 +2637,8 @@ static bool checkZeroizeStack(const Function &Caller, const Function &Callee) {
   // An absent, an empty, and an unrecognized value all mean "used", so a value
   // this consumer cannot interpret clears more of the frame than it must,
   // never less. See llvm/test/Transforms/Inline/zeroize-stack.ll.
-  return Caller.getFnAttribute("zeroize-stack").getValueAsString() !=
-             ZeroizeStackNarrowestMode ||
-         Callee.getFnAttribute("zeroize-stack").getValueAsString() ==
-             ZeroizeStackNarrowestMode;
+  return Caller.getZeroizeStackMode() != ZeroizeStackNarrowestMode ||
+         Callee.getZeroizeStackMode() == ZeroizeStackNarrowestMode;
 }
 
 template<typename AttrClass>

llvm/lib/IR/Verifier.cpp

@@ -3115,6 +3115,15 @@ void Verifier::visitFunction(const Function &F) {
     for (const Argument &Arg : F.args())
       Check(Arg.use_empty(), "cannot use argument of naked function", &Arg);
 
+  // CoroSplit lowers a presplit coroutine into resume/destroy clones that hand
+  // off with musttail calls (symmetric transfer, and the async coro.end) and
+  // copies the coroutine's function attributes onto those clones. A protected
+  // coroutine would become a protected function holding a musttail call, which
+  // verifyMustTailCall rejects, so reject it here before the split for the same
+  // reason: the frame the attribute must clear is handed off and never cleared.
+  Check(!F.isPresplitCoroutine() || !F.hasZeroizeStack(),
+        "cannot use the \"zeroize-stack\" attribute on a coroutine", &F);
+
   // Check that this function meets the restrictions on this calling convention.
   // Sometimes varargs is used for perfectly forwarding thunks, so some of these
   // restrictions can be lifted.
@@ -4214,6 +4223,16 @@ void Verifier::verifyMustTailCall(CallInst &CI) {
   Check(!CI.isInlineAsm(), "cannot use musttail call with inline asm", &CI);
 
   Function *F = CI.getParent()->getParent();
+
+  // "zeroize-stack" clears the frame before returning; musttail replaces the
+  // frame and never returns here to clear it. An ordinary tail call is an
+  // optimization and is suppressed, but musttail is a requirement the caller
+  // cannot drop, so the two together describe a function that cannot exist.
+  Check(!F->hasZeroizeStack(),
+        "cannot use musttail call in a function with the \"zeroize-stack\" "
+        "attribute",
+        &CI);
+
   FunctionType *CallerTy = F->getFunctionType();
   FunctionType *CalleeTy = CI.getFunctionType();
   Check(CallerTy->isVarArg() == CalleeTy->isVarArg(),

llvm/lib/Transforms/IPO/MergeFunctions.cpp

@@ -459,6 +459,12 @@ static bool hasDistinctMetadataIntrinsic(const Function &F) {
 static bool isEligibleForMerging(Function &F) {
   return !F.isDeclaration() && !F.hasAvailableExternallyLinkage() &&
          !F.hasFnAttribute(Attribute::NoIPA) &&
+         // Merging turns the function into a thunk that tail-calls the merged
+         // body and keeps its attributes; a "zeroize-stack" thunk would promise
+         // to clear a frame it no longer owns, and a musttail thunk (swifttailcc)
+         // cannot clear one at all. Keep protected functions whole, as inlining
+         // already does.
+         !F.hasZeroizeStack() &&
          !hasDistinctMetadataIntrinsic(F);
 }
 

llvm/test/CodeGen/AArch64/zeroize-tailcall-gisel.ll

@@ -0,0 +1,26 @@
+; A libcall the legalizer generates has no call in the IR behind it, so it is
+; not caught where an IR call is. GlobalISel forms these libcalls on its own
+; path, separate from SelectionDAG, and decides the tail call in the legalizer
+; rather than in TargetLowering::isInTailCallPosition. The suppression is asked
+; there too: a protected function keeps the call and returns through its own
+; epilogue where an unprotected one branches away.
+
+; RUN: llc -mtriple=aarch64-unknown-linux-gnu -global-isel %s -o - 2>/dev/null | FileCheck %s
+
+; The "zeroize-stack" function also reports that no target clears the frame yet.
+; That report is a warning, so llc still succeeds, and stderr is discarded here.
+
+; CHECK-LABEL: protected_libcall:
+; CHECK:         bl fmod
+; CHECK:         ret
+define double @protected_libcall(double %a, double %b) "zeroize-stack"="used" {
+  %r = frem double %a, %b
+  ret double %r
+}
+
+; CHECK-LABEL: unprotected_libcall:
+; CHECK:         b fmod
+define double @unprotected_libcall(double %a, double %b) {
+  %r = frem double %a, %b
+  ret double %r
+}

llvm/test/CodeGen/ARM/zeroize-fallback.ll

@@ -0,0 +1,52 @@
+; The fallbacks are not written in terms of any one target's instructions, and
+; two of them are visible on a target that cannot clear anything at all: which
+; exits are in scope is decided before any target is asked, and an unreadable
+; mode is resolved before the target is asked too.
+;
+; trailofbits/vspells-ct-internal-notes#24.
+
+; Both runs are under "not", because the widened mode reaches a refusal this
+; target has to give and llc exits non-zero for it. That refusal is the second
+; half of what is being tested.
+; RUN: not llc -mtriple=armv7-unknown-linux-gnueabi -pei-print-clearing-sequence %s -o /dev/null 2>&1 | FileCheck --check-prefix=SEQ %s
+; RUN: not llc -mtriple=armv7-unknown-linux-gnueabi %s -o /dev/null 2>&1 | FileCheck --check-prefix=DIAG %s
+
+@g = external global i32
+
+declare void @llvm.trap()
+
+; A supervisor call written as inline assembly ends the block, and whether
+; control comes back from it is not something the compiler can decide. It is
+; in scope here for the same reason it is on x86-64.
+; SEQ-LABEL: clearing sequence for function 'opaque_asm':
+; SEQ-NEXT:  %bb.0 unknown: clear-stack=not-requested clear-registers=not-requested clear-flags=unimplemented
+; SEQ-NEXT:  end clearing sequence for function 'opaque_asm'
+define void @opaque_asm(i32 %a, i32 %b) {
+  %s = add i32 %a, %b
+  store i32 %s, ptr @g
+  call void asm sideeffect "svc #0", "~{memory}"()
+  unreachable
+}
+
+; A trap is a trap on every target that marks one, and stays out of scope.
+; SEQ-LABEL: clearing sequence for function 'traps':
+; SEQ-NEXT:  end clearing sequence for function 'traps'
+define void @traps() {
+  call void @llvm.trap()
+  unreachable
+}
+
+; An unreadable mode is not "skip". ARM cannot clear registers, so what the
+; widened mode reaches here is the target's refusal, which is reported; what it
+; does not do is quietly resolve to clearing nothing and say nothing.
+; DIAG: error: {{.*}}in function unrecognized_mode i32 (i32): "zero-call-used-regs" is not supported by this target
+define i32 @unrecognized_mode(i32 %x) "zero-call-used-regs"="a-mode-from-the-future" {
+  ret i32 %x
+}
+
+; A mode that says to skip is read and honored, on this target as on any other,
+; so it reaches no refusal.
+; DIAG-NOT: in function skips_explicitly
+define i32 @skips_explicitly(i32 %x) "zero-call-used-regs"="skip" {
+  ret i32 %x
+}

llvm/test/CodeGen/ARM/zeroize-scratch-regs.ll

@@ -0,0 +1,35 @@
+; The register clear is what finishes the stack clear's work: it destroys the
+; registers the stack clear read the frame through. A target that cannot clear
+; registers therefore cannot clear the frame either, and has to say so rather
+; than emit the half of the sequence it can do. ARM implements neither, so it
+; is where that can be pinned.
+;
+; As in the X86 test, -pei-stack-clear-scratch-regs stands in for the step that
+; clears the frame, which no target implements
+; (trailofbits/vspells-ct-internal-notes#26).
+
+; RUN: not llc -mtriple=armv7-unknown-linux-gnueabi -pei-stack-clear-scratch-regs=r4 < %s -o /dev/null 2>&1 | FileCheck %s
+
+; The function asked for its frame to be cleared and said nothing about its
+; registers, so the register clear it gets is one it did not ask for. It is
+; still a register clear, and this target cannot do one, so the request to
+; clear the frame cannot be discharged.
+; CHECK: error: {{.*}}in function stack_only i32 (i32): clearing the stack needs the registers it uses to be cleared afterwards, which is not supported by this target
+define i32 @stack_only(i32 %x) "zeroize-stack"="used" {
+  ret i32 %x
+}
+
+; A function that did ask for its registers to be cleared is refused on its own
+; terms, by the query that has always answered that request, rather than being
+; refused twice or reported as something it did not ask for.
+; CHECK: error: {{.*}}in function asked_for_both i32 (i32): "zero-call-used-regs" is not supported by this target
+; CHECK-NOT: in function asked_for_both {{.*}}clearing the stack needs
+define i32 @asked_for_both(i32 %x) "zeroize-stack"="used" "zero-call-used-regs"="used-gpr" {
+  ret i32 %x
+}
+
+; A function that asked for neither is not dragged into any of this.
+; CHECK-NOT: in function untouched
+define i32 @untouched(i32 %x) {
+  ret i32 %x
+}

llvm/test/CodeGen/ARM/zeroize-tailcall.ll

@@ -0,0 +1,43 @@
+; The suppression is decided in target-independent code, so it reaches a target
+; that forms its tail calls differently. ARM turns a call in tail position into
+; a branch; a protected function keeps the call and returns through its own
+; epilogue instead.
+
+; RUN: llc -mtriple=armv7-unknown-linux-gnueabi %s -o - 2>/dev/null | FileCheck %s
+
+; The "zeroize-stack" functions also report that no target clears the frame yet.
+; That report is a warning, so llc still succeeds, and stderr is discarded here.
+
+declare i32 @callee(i32)
+
+; CHECK-LABEL: protected:
+; CHECK:         bl callee
+; CHECK:         pop {r11, pc}
+define i32 @protected(i32 %x) "zeroize-stack"="used" {
+  %r = tail call i32 @callee(i32 %x)
+  ret i32 %r
+}
+
+; CHECK-LABEL: unprotected:
+; CHECK:         b callee
+define i32 @unprotected(i32 %x) {
+  %r = tail call i32 @callee(i32 %x)
+  ret i32 %r
+}
+
+; A libcall the legalizer generates has no call in the IR behind it, and is
+; refused at the second place the tail-call question is asked.
+; CHECK-LABEL: protected_libcall:
+; CHECK:         bl fmod
+; CHECK:         pop {r11, pc}
+define double @protected_libcall(double %a, double %b) "zeroize-stack"="used" {
+  %r = frem double %a, %b
+  ret double %r
+}
+
+; CHECK-LABEL: unprotected_libcall:
+; CHECK:         b fmod
+define double @unprotected_libcall(double %a, double %b) {
+  %r = frem double %a, %b
+  ret double %r
+}

llvm/test/CodeGen/RISCV/zero-call-used-regs-fp.ll

@@ -33,6 +33,7 @@ define double @used(double noundef %a, float noundef %b) "zero-call-used-regs"="
 ; 32-BITS-F-NEXT:    .cfi_def_cfa_offset 0
 ; 32-BITS-F-NEXT:    li a2, 0
 ; 32-BITS-F-NEXT:    li a3, 0
+; 32-BITS-F-NEXT:    fmv.w.x fa0, zero
 ; 32-BITS-F-NEXT:    ret
 ;
 ; 32-BITS-D-LABEL: used:
@@ -70,6 +71,7 @@ define double @used(double noundef %a, float noundef %b) "zero-call-used-regs"="
 ; 64-BITS-F-NEXT:    addi sp, sp, 16
 ; 64-BITS-F-NEXT:    .cfi_def_cfa_offset 0
 ; 64-BITS-F-NEXT:    li a1, 0
+; 64-BITS-F-NEXT:    fmv.w.x fa0, zero
 ; 64-BITS-F-NEXT:    ret
 ;
 ; 64-BITS-D-LABEL: used:
@@ -187,6 +189,8 @@ define double @used_arg_double(double noundef %a, double noundef %b) "zero-call-
 ; 32-BITS-F-NEXT:    .cfi_restore ra
 ; 32-BITS-F-NEXT:    addi sp, sp, 16
 ; 32-BITS-F-NEXT:    .cfi_def_cfa_offset 0
+; 32-BITS-F-NEXT:    li a2, 0
+; 32-BITS-F-NEXT:    li a3, 0
 ; 32-BITS-F-NEXT:    ret
 ;
 ; 32-BITS-D-LABEL: used_arg_double:
@@ -212,6 +216,7 @@ define double @used_arg_double(double noundef %a, double noundef %b) "zero-call-
 ; 64-BITS-F-NEXT:    .cfi_restore ra
 ; 64-BITS-F-NEXT:    addi sp, sp, 16
 ; 64-BITS-F-NEXT:    .cfi_def_cfa_offset 0
+; 64-BITS-F-NEXT:    li a1, 0
 ; 64-BITS-F-NEXT:    ret
 ;
 ; 64-BITS-D-LABEL: used_arg_double: