[WebAssembly] Restore __stack_pointer after catch instructions

[aheejin]

Summary

After the stack is unwound due to a thrown exception, the
__stack_pointer global can point to an invalid address. This inserts
instructions that restore __stack_pointer global.

[jgravelle-google]

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

58

This is kinda awkwardly indented, was this clang-formatted?
Might be worth extracting a temporary variable here for the getExceptionHandlingType() call, i.e. auto EHType = MF.getTarget().getMCAsmInfo()->getExceptionHandlingType();, which should make the line-breaks more consistent.

... Reading further down this patch, can we just call needsPrologForEH here?

[dschuff]

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

73

This comment is a little confusing. Looking at the code, I guess it means we just always insert the fixup at the top of the BB, (but after the catch if it hasn't been moved)?
What prevents the SP restoration from being moved after a call in the catch block?

lib/Target/WebAssembly/WebAssemblyFrameLowering.h

49

Since SP is now a global and not a memory address anymore, we should change this function name. I realize that's pre-existing for this change, so it can be in a different CL (either before or after this one, whatever is convenient).

[aheejin]

Address comments

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

58

This is kinda awkwardly indented, was this clang-formatted?

Yes it was.

... Reading further down this patch, can we just call needsPrologForEH here?

Good idea. Done.

Might be worth extracting a temporary variable here for the getExceptionHandlingType() call, i.e. auto EHType = MF.getTarget().getMCAsmInfo()->getExceptionHandlingType();, which should make the line-breaks more consistent.

Done in WebAssemblyFrameLowering::needsPrologForEH.

73

This comment is a little confusing. Looking at the code, I guess it means we just always insert the fixup at the top of the BB, (but after the catch if it hasn't been moved)?

Yes that's what it tries to do. Do you have any suggestions on making it less confusing? What I tried to say is, at this point, we have catch by this point but not catch_all (it will be generated later in LateEHPrepare). And even if we have catch, it may not be staying at the top. But here we don't want to go out of the way to find the catch instruction and generate these instructions after it because it will be done in LateEHPrepare anyway.

What prevents the SP restoration from being moved after a call in the catch block?

I'm not exactly sure about this. What prevents normal epilogues generated by WebAssemblyFrameLowering::emitEpilogue being moved after a call within the same BB?

lib/Target/WebAssembly/WebAssemblyFrameLowering.h

49

Done in D51046, which should land before this one.

[dschuff]

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

73

Stepping back, why can catch instructions get reordered? It seems odd that they should be if we model them as having side effects or control flow or whatever.

Stepping back, if we move catches later, and insert catch_all later, is there a reason we even need to put the catch in the BB early? Why not just add them when we add catch_all? Is it just because we ISel it from catchpad? I guess if we could avoid reordering, then it wouldn't be so confusing.

I'm not exactly sure about this. What prevents normal epilogues generated by WebAssemblyFrameLowering::emitEpilogue being moved after a call within the same BB?

I think actually what happens is that calls are modeled as modifying SP, so that is ordered wrt anything else that modifies SP. So maybe that should just work for us too

[aheejin]

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

73

Stepping back, why can catch instructions get reordered? It seems odd that they should be if we model them as having side effects or control flow or whatever.

It is currently marked as hasCtrlDep, but I still have noticed cases that catches are reordered (mostly something like i32.const is hoisted before catch). I added hasSideEffects in D50919 just to be safer. But even with that I don't think we can be sure instructions like i32.const will not be reordered before catch, or something like this won't happen:

ehpad:
  br bb2

bb3:
  catch

So LateEHPrepare even considers cases like the latter and make sure catch will be the first instruction of an EH pad.

Stepping back, if we move catches later, and insert catch_all later, is there a reason we even need to put the catch in the BB early? Why not just add them when we add catch_all? Is it just because we ISel it from catchpad? I guess if we could avoid reordering, then it wouldn't be so confusing.

The reason we select catch in the isel stage is because unlike catch_all, it produces a value. Instructions below catch use the value produced by the catch.

[dschuff]

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

73

So LateEHPrepare even considers cases like the latter and make sure catch will be the first instruction of an EH pad.

Hm, interesting.
In the latter case, I guess we'd end up with control flow inside the catch block, which is OK even when there are instructions in the first BB.

So getting back to this code, I guess it's not actually necessary to move the insert point below catch, and it makes the explanation funny, but OTOH it probably results in more readable IR in the common case.
I guess the comment could be something like:
"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)."
That's not much better than what you had, but it does state the simple case first, followed by the extra explanation, which IMO is a little simpler to read.

[aheejin]

Fix comments

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

73

Yes, it's not strictly necessary to move the insert point below catch. I did it because as you said it's gonna make code more readable in common cases. Fixed comments; I agree this is simpler and easier to understand. Thanks.

lib/Target/WebAssembly/CMakeLists.txt

 add_llvm_target(WebAssemblyCodeGen
   WebAssemblyAddMissingPrototypes.cpp
   WebAssemblyArgumentMove.cpp
   WebAssemblyAsmPrinter.cpp
   WebAssemblyCallIndirectFixup.cpp
   WebAssemblyCFGStackify.cpp
   WebAssemblyCFGSort.cpp
   WebAssemblyLateEHPrepare.cpp
+  WebAssemblyEHRestoreStackPointer.cpp
   WebAssemblyExceptionInfo.cpp
   WebAssemblyExplicitLocals.cpp
   WebAssemblyFastISel.cpp
   WebAssemblyFixIrreducibleControlFlow.cpp
   WebAssemblyFixFunctionBitcasts.cpp
   WebAssemblyFrameLowering.cpp
   WebAssemblyISelDAGToDAG.cpp
   WebAssemblyISelLowering.cpp

lib/Target/WebAssembly/WebAssembly.h

 
 // ISel and immediate followup passes.
 FunctionPass *createWebAssemblyISelDag(WebAssemblyTargetMachine &TM,
                                        CodeGenOpt::Level OptLevel);
 FunctionPass *createWebAssemblyArgumentMove();
 FunctionPass *createWebAssemblySetP2AlignOperands();
 
 // Late passes.
+FunctionPass *createWebAssemblyEHRestoreStackPointer();
 FunctionPass *createWebAssemblyReplacePhysRegs();
 FunctionPass *createWebAssemblyPrepareForLiveIntervals();
 FunctionPass *createWebAssemblyOptimizeLiveIntervals();
 FunctionPass *createWebAssemblyStoreResults();
 FunctionPass *createWebAssemblyRegStackify();
 FunctionPass *createWebAssemblyRegColoring();
 FunctionPass *createWebAssemblyExplicitLocals();
 FunctionPass *createWebAssemblyFixIrreducibleControlFlow();
 FunctionPass *createWebAssemblyLateEHPrepare();
 FunctionPass *createWebAssemblyCFGSort();
 FunctionPass *createWebAssemblyCFGStackify();
 FunctionPass *createWebAssemblyLowerBrUnless();
 FunctionPass *createWebAssemblyRegNumbering();
 FunctionPass *createWebAssemblyPeephole();
 FunctionPass *createWebAssemblyCallIndirectFixup();
 
 // PassRegistry initialization declarations.
 void initializeWebAssemblyAddMissingPrototypesPass(PassRegistry &);
 void initializeWebAssemblyLowerEmscriptenEHSjLjPass(PassRegistry &);
 void initializeLowerGlobalDtorsPass(PassRegistry &);
 void initializeFixFunctionBitcastsPass(PassRegistry &);
 void initializeOptimizeReturnedPass(PassRegistry &);
 void initializeWebAssemblyArgumentMovePass(PassRegistry &);
 void initializeWebAssemblySetP2AlignOperandsPass(PassRegistry &);
+void initializeWebAssemblyEHRestoreStackPointerPass(PassRegistry &);
 void initializeWebAssemblyReplacePhysRegsPass(PassRegistry &);
 void initializeWebAssemblyPrepareForLiveIntervalsPass(PassRegistry &);
 void initializeWebAssemblyOptimizeLiveIntervalsPass(PassRegistry &);
 void initializeWebAssemblyStoreResultsPass(PassRegistry &);
 void initializeWebAssemblyRegStackifyPass(PassRegistry &);
 void initializeWebAssemblyRegColoringPass(PassRegistry &);
 void initializeWebAssemblyExplicitLocalsPass(PassRegistry &);
 void initializeWebAssemblyFixIrreducibleControlFlowPass(PassRegistry &);

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

+//===-- WebAssemblyEHRestoreStackPointer.cpp - __stack_pointer restoration ===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+///
+/// \file
+/// After the stack is unwound due to a thrown exception, the __stack_pointer
+/// global can point to an invalid address. This inserts instructions that
+/// restore __stack_pointer global.
+///
+//===----------------------------------------------------------------------===//
+
+#include "MCTargetDesc/WebAssemblyMCTargetDesc.h"
+#include "WebAssembly.h"
+#include "WebAssemblySubtarget.h"
+#include "WebAssemblyUtilities.h"
+#include "llvm/CodeGen/MachineFrameInfo.h"
+#include "llvm/MC/MCAsmInfo.h"
+using namespace llvm;
+
+#define DEBUG_TYPE "wasm-eh-restore-stack-pointer"
+
+namespace {
+class WebAssemblyEHRestoreStackPointer final : public MachineFunctionPass {
+public:
+  static char ID; // Pass identification, replacement for typeid
+  WebAssemblyEHRestoreStackPointer() : MachineFunctionPass(ID) {}
+
+  StringRef getPassName() const override {
+    return "WebAssembly Restore Stack Pointer for Exception Handling";
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesCFG();
+    // TODO more?
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+};
+} // end anonymous namespace
+
+char WebAssemblyEHRestoreStackPointer::ID = 0;
+INITIALIZE_PASS(WebAssemblyEHRestoreStackPointer, DEBUG_TYPE,
+                "Restore Stack Pointer for Exception Handling", true, false)
+
+FunctionPass *llvm::createWebAssemblyEHRestoreStackPointer() {
+  return new WebAssemblyEHRestoreStackPointer();
+}
+
+bool WebAssemblyEHRestoreStackPointer::runOnMachineFunction(
+    MachineFunction &MF) {
+  if (MF.getTarget().getMCAsmInfo()->getExceptionHandlingType() !=
+          ExceptionHandling::Wasm ||
+      !MF.getFunction().hasPersonalityFn() || !MF.getFrameInfo().hasCalls())
+    return false;
+  bool Changed = false;
+
+  for (auto &MBB : MF) {
+    if (!MBB.isEHPad())
+      continue;
+    Changed = true;
+
+    // Insert __stack_pointer restoring instructions at the beginning of each EH
+    // pad. At this point we don't have catch_all instructions yet, which are
+    // added in LateEHPrepare, and we have catch instructions supposedly at the
+    // top of an EH pad but it is possible they have been reordered with other
+    // instructions. We hoist them to the top of EH pads in LateEHPrepare, so
+    // here we just insert instructions after a catch if it stays at the top.
+    auto InsertPos = MBB.begin();
+    if (WebAssembly::isCatch(*MBB.begin()))
+      InsertPos++;
+    const auto *FrameLowering = static_cast<const WebAssemblyFrameLowering *>(
+        MF.getSubtarget().getFrameLowering());
+    FrameLowering->writeSPToMemory(WebAssembly::SP32, MF, MBB, InsertPos,
+                                   MBB.begin()->getDebugLoc());
+  }
+  return Changed;
+}

lib/Target/WebAssembly/WebAssemblyFrameLowering.h

public:
 
   /// These methods insert prolog and epilog code into the function.
   void emitPrologue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
   void emitEpilogue(MachineFunction &MF, MachineBasicBlock &MBB) const override;
 
   bool hasFP(const MachineFunction &MF) const override;
   bool hasReservedCallFrame(const MachineFunction &MF) const override;
 
- private:
+  /// Write SP back to __stack_pointer global.
+  void writeSPToMemory(unsigned SrcReg, MachineFunction &MF,
+                       MachineBasicBlock &MBB,
+                       MachineBasicBlock::iterator &InsertStore,
+                       const DebugLoc &DL) const;
+
   bool hasBP(const MachineFunction &MF) const;
   bool needsSP(const MachineFunction &MF, const MachineFrameInfo &MFI) const;
   bool needsSPWriteback(const MachineFunction &MF,
                         const MachineFrameInfo &MFI) const;
 };
 
 }  // end namespace llvm
 
 #endif

lib/Target/WebAssembly/WebAssemblyFrameLowering.cpp

 #include "WebAssemblySubtarget.h"
 #include "WebAssemblyTargetMachine.h"
 #include "WebAssemblyUtilities.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineModuleInfoImpls.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/MC/MCAsmInfo.h"
 #include "llvm/Support/Debug.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "wasm-frame-info"
 
 // TODO: wasm64
 // TODO: Emit TargetOpcode::CFI_INSTRUCTION instructions
 

 /// 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();
 }
 
+// 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.
+static bool needsPrologForEH(const MachineFunction &MF) {
+  return MF.getTarget().getMCAsmInfo()->getExceptionHandlingType() ==
+             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 MachineFrameInfo &MFI) const {
-  return MFI.getStackSize() || MFI.adjustsStack() || hasFP(MF);
+  return MFI.getStackSize() || MFI.adjustsStack() || hasFP(MF) ||
+         needsPrologForEH(MF);
 }
 
 /// Returns true if the local user-space stack pointer needs to be written back
 /// to memory 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 MachineFrameInfo &MFI) const {
   assert(needsSP(MF, MFI));
   return MFI.getStackSize() > RedZoneSize || MFI.hasCalls() ||
          MF.getFunction().hasFnAttribute(Attribute::NoRedZone);
 }
 
-static void writeSPToMemory(unsigned SrcReg, MachineFunction &MF,
-                            MachineBasicBlock &MBB,
-                            MachineBasicBlock::iterator &InsertStore,
-                            const DebugLoc &DL) {
+void WebAssemblyFrameLowering::writeSPToMemory(
+    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";
   auto *SPSymbol = MF.createExternalSymbolName(ES);
   BuildMI(MBB, InsertStore, DL, TII->get(WebAssembly::SET_GLOBAL_I32))
       .addExternalSymbol(SPSymbol, WebAssemblyII::MO_SYMBOL_GLOBAL)
       .addReg(SrcReg);
 }

lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp

extern "C" void LLVMInitializeWebAssemblyTarget() {
   auto &PR = *PassRegistry::getPassRegistry();
   initializeWebAssemblyAddMissingPrototypesPass(PR);
   initializeWebAssemblyLowerEmscriptenEHSjLjPass(PR);
   initializeLowerGlobalDtorsPass(PR);
   initializeFixFunctionBitcastsPass(PR);
   initializeOptimizeReturnedPass(PR);
   initializeWebAssemblyArgumentMovePass(PR);
   initializeWebAssemblySetP2AlignOperandsPass(PR);
+  initializeWebAssemblyEHRestoreStackPointerPass(PR);
   initializeWebAssemblyReplacePhysRegsPass(PR);
   initializeWebAssemblyPrepareForLiveIntervalsPass(PR);
   initializeWebAssemblyOptimizeLiveIntervalsPass(PR);
   initializeWebAssemblyStoreResultsPass(PR);
   initializeWebAssemblyRegStackifyPass(PR);
   initializeWebAssemblyRegColoringPass(PR);
   initializeWebAssemblyExplicitLocalsPass(PR);
   initializeWebAssemblyFixIrreducibleControlFlowPass(PR);
void WebAssemblyPassConfig::addPostRegAlloc() {
   disablePass(&ShrinkWrapID);
 
   TargetPassConfig::addPostRegAlloc();
 }
 
 void WebAssemblyPassConfig::addPreEmitPass() {
   TargetPassConfig::addPreEmitPass();
 
+  // Restore __stack_pointer global after an exception is thrown.
+  addPass(createWebAssemblyEHRestoreStackPointer());
+
   // Now that we have a prologue and epilogue and all frame indices are
   // rewritten, eliminate SP and FP. This allows them to be stackified,
   // colored, and numbered with the rest of the registers.
   addPass(createWebAssemblyReplacePhysRegs());
 
   // Rewrite pseudo call_indirect instructions as real instructions.
   // This needs to run before register stackification, because we change the
   // order of the arguments.

test/CodeGen/WebAssembly/exception.ll

 ; CHECK-NEXT: i32.const $push0=, 0
 ; CHECK-NEXT: throw 0, $pop0
 define void @test_throw() {
   call void @llvm.wasm.throw(i32 0, i8* null)
   ret void
 }
 
 ; CHECK-LABEL: test_catch_rethrow:
+; CHECK:   get_global  $push{{.+}}=, __stack_pointer@GLOBAL
 ; CHECK:   try
 ; CHECK:   call      foo@FUNCTION
 ; CHECK:   i32.catch     $push{{.+}}=, 0
+; CHECK:   set_global  __stack_pointer@GLOBAL
 ; CHECK-DAG:   i32.store  __wasm_lpad_context
 ; CHECK-DAG:   i32.store  __wasm_lpad_context+4
 ; CHECK:   i32.call  $push{{.+}}=, _Unwind_CallPersonality@FUNCTION
 ; CHECK:   i32.call  $push{{.+}}=, __cxa_begin_catch@FUNCTION
 ; CHECK:   call      __cxa_end_catch@FUNCTION
 ; CHECK:   call      __cxa_rethrow@FUNCTION
 ; CHECK-NEXT:   rethrow
 ; CHECK:   end_try

 try.cont:                                         ; preds = %entry, %catch
   ret void
 }
 
 ; CHECK-LABEL: test_cleanup:
 ; CHECK:   try
 ; CHECK:   call      foo@FUNCTION
 ; CHECK:   catch_all
+; CHECK:   set_global  __stack_pointer@GLOBAL
 ; CHECK:   i32.call  $push{{.+}}=, _ZN7CleanupD1Ev@FUNCTION
 ; CHECK:   rethrow
 ; CHECK:   end_try
 define void @test_cleanup() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   %c = alloca %struct.Cleanup, align 1
   invoke void @foo()
           to label %invoke.cont unwind label %ehcleanup

 ; not have an epilogue. This is separate from __stack_pointer restoring
 ; instructions after a catch instruction.
 
 ; CHECK-LABEL: test_no_prolog_epilog_in_ehpad
 ; CHECK:  try
 ; CHECK:  call      foo@FUNCTION
 ; CHECK:  i32.catch
 ; CHECK-NOT:  get_global  $push{{.+}}=, __stack_pointer@GLOBAL
+; CHECK:  set_global  __stack_pointer@GLOBAL
 ; CHECK:  try
 ; CHECK:  call      foo@FUNCTION
 ; CHECK:  catch_all
 ; CHECK-NOT:  get_global  $push{{.+}}=, __stack_pointer@GLOBAL
+; CHECK:  set_global  __stack_pointer@GLOBAL
 ; CHECK:  call      __cxa_end_catch@FUNCTION
 ; CHECK-NOT:  set_global  __stack_pointer@GLOBAL, $pop{{.+}}
 ; CHECK:  end_try
 ; CHECK-NOT:  set_global  __stack_pointer@GLOBAL, $pop{{.+}}
 ; CHECK:  end_try
 define void @test_no_prolog_epilog_in_ehpad() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   %stack_var = alloca i32, align 4