[WebAssembly] Don't write SP back when prolog is generated only for EH

[aheejin]

Summary

When we don't actually have stack-allocated variables but needs SP only
to support EH, we don't need to write SP back in the epilog, because we
don't bump down the stack pointer.

[dschuff]

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

73

__stack_pointer

75

can't leaf functions have try/catch blocks, even if there are no calls? I think the real issue is that leaf functions do not to restore the stack pointer on catch, because it won't have been updated by any callees.

lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp

122

Another possible way:
bool NeedsSPForLocalFrame = MFI.getStackSize() || !MFI.adjustsStack() || hasFP(MF);

which would avoid so many nots...

Actually, even better may be to turn NeedsSPForLocalFrame() into a function (that can be called by needsSP()). Then you can reuse it here instead of duplicating the logic?

[aheejin]

Address comments

[aheejin]

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

75

Right, good point.

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

for (auto &MBB : MF) {
     if (!MBB.isEHPad())
       continue;
     Changed = true;
 
     // Insert __stack_pointer restoring instructions at the beginning of each EH
     // pad, after the catch instruction. (Catch instructions may have been
     // reordered, and catch_all instructions have not been inserted yet, but
     // those cases are handled in LateEHPrepare).
+    //
+    // Here it is safe to assume that SP32 holds the latest value of
+    // __stack_pointer, because the only exception for this case is when a
+    // function uses the red zone, but that only happens with leaf functions,
+    // and we don't restore __stack_pointer in leaf functions anyway.
     auto InsertPos = MBB.begin();
     if (WebAssembly::isCatch(*MBB.begin()))
       InsertPos++;
     FrameLowering->writeSPToGlobal(WebAssembly::SP32, MF, MBB, InsertPos,
                                    MBB.begin()->getDebugLoc());
   }
   return Changed;
 }

lib/Target/WebAssembly/WebAssemblyFrameLowering.h

public:
   /// Write SP back to __stack_pointer global.
   void writeSPToGlobal(unsigned SrcReg, MachineFunction &MF,
                        MachineBasicBlock &MBB,
                        MachineBasicBlock::iterator &InsertStore,
                        const DebugLoc &DL) const;
 
 private:
   bool hasBP(const MachineFunction &MF) const;
+  bool needsSPForLocalFrame(const MachineFunction &MF) const;
   bool needsSP(const MachineFunction &MF) const;
   bool needsSPWriteback(const MachineFunction &MF) const;
 };
 
 }  // end namespace llvm
 
 #endif

lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp

 /// current function. This eliminates the need for add/sub sp brackets around
 /// call sites. Returns true if the call frame is included as part of the stack
 /// frame.
 bool WebAssemblyFrameLowering::hasReservedCallFrame(
     const MachineFunction &MF) const {
   return !MF.getFrameInfo().hasVarSizedObjects();
 }
 
+// Returns true if this function needs a local user-space stack pointer for its
+// local frame (not for exception handling).
+bool WebAssemblyFrameLowering::needsSPForLocalFrame(
+    const MachineFunction &MF) const {
+  auto &MFI = MF.getFrameInfo();
+  return MFI.getStackSize() || MFI.adjustsStack() || hasFP(MF);
+}
+
 // In function with EH pads, we need to make a copy of the value of
 // __stack_pointer global in SP32 register, in order to use it when restoring
 // __stack_pointer after an exception is caught.
 bool WebAssemblyFrameLowering::needsPrologForEH(
     const MachineFunction &MF) const {
   auto EHType = MF.getTarget().getMCAsmInfo()->getExceptionHandlingType();
   return EHType == ExceptionHandling::Wasm &&
          MF.getFunction().hasPersonalityFn() && MF.getFrameInfo().hasCalls();
 }
 
 /// Returns true if this function needs a local user-space stack pointer.
 /// Unlike a machine stack pointer, the wasm user stack pointer is a global
 /// variable, so it is loaded into a register in the prolog.
 bool WebAssemblyFrameLowering::needsSP(const MachineFunction &MF) const {
-  auto &MFI = MF.getFrameInfo();
-  return MFI.getStackSize() || MFI.adjustsStack() || hasFP(MF) ||
-         needsPrologForEH(MF);
+  return needsSPForLocalFrame(MF) || needsPrologForEH(MF);
 }
 
 /// Returns true if the local user-space stack pointer needs to be written back
 /// to __stack_pointer global by this function (this is not meaningful if
 /// needsSP is false). If false, the stack red zone can be used and only a local
 /// SP is needed.
 bool WebAssemblyFrameLowering::needsSPWriteback(
     const MachineFunction &MF) const {
   auto &MFI = MF.getFrameInfo();
   assert(needsSP(MF));
-  return MFI.getStackSize() > RedZoneSize || MFI.hasCalls() ||
-         MF.getFunction().hasFnAttribute(Attribute::NoRedZone);
+  // When we don't need a local stack pointer for its local frame but only to
+  // support EH, we don't need to write SP back in the epilog, because we don't
+  // bump down the stack pointer in the prolog. We need to write SP back in the
+  // epilog only if
+  // 1. We need SP not only for EH support but also because we actually use
+  // stack or we have a frame address taken.
+  // 2. We cannot use the red zone.
+  bool CanUseRedZone = MFI.getStackSize() <= RedZoneSize && !MFI.hasCalls() &&
+                       !MF.getFunction().hasFnAttribute(Attribute::NoRedZone);
+  return needsSPForLocalFrame(MF) && !CanUseRedZone;
 }
 
 void WebAssemblyFrameLowering::writeSPToGlobal(
     unsigned SrcReg, MachineFunction &MF, MachineBasicBlock &MBB,
     MachineBasicBlock::iterator &InsertStore, const DebugLoc &DL) const {
   const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
 
   const char *ES = "__stack_pointer";

test/CodeGen/WebAssembly/exception.ll

try.cont:                                         ; preds = %entry, %invoke.cont1
   ret void
 
 ehcleanup:                                        ; preds = %catch
   %8 = cleanuppad within %1 []
   call void @__cxa_end_catch() [ "funclet"(token %8) ]
   cleanupret from %8 unwind to caller
 }
 
+; When a function does not have stack-allocated objects, it does not need to
+; store SP back to __stack_pointer global at the epilog.
+
+; CHECK-LABEL: no_sp_writeback
+; CHECK:  try
+; CHECK:  call foo@FUNCTION
+; CHECK:  end_try
+; CHECK-NOT:  set_global  __stack_pointer@GLOBAL
+; CHECK:  return
+define void @no_sp_writeback() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  invoke void @foo()
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %entry
+  %0 = catchswitch within none [label %catch.start] unwind to caller
+
+catch.start:                                      ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null]
+  %2 = call i8* @llvm.wasm.get.exception(token %1)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  %4 = call i8* @__cxa_begin_catch(i8* %2) #2 [ "funclet"(token %1) ]
+  call void @__cxa_end_catch() [ "funclet"(token %1) ]
+  catchret from %1 to label %try.cont
+
+try.cont:                                         ; preds = %entry, %catch.start
+  ret void
+}
+
 declare void @foo()
 declare void @bar(i32*)
 declare i32 @__gxx_wasm_personality_v0(...)
 declare i8* @llvm.wasm.get.exception(token)
 declare i32 @llvm.wasm.get.ehselector(token)
 declare i32 @llvm.eh.typeid.for(i8*)
 declare i8* @__cxa_begin_catch(i8*)
 declare void @__cxa_end_catch()
 declare void @__cxa_rethrow()
 declare void @__clang_call_terminate(i8*)
 declare void @_ZSt9terminatev()
 declare %struct.Cleanup* @_ZN7CleanupD1Ev(%struct.Cleanup* returned)