[WebAssembly] Restore stack pointer right after catch instruction

[aheejin]

Summary

After the staack is unwound due to a thrown exxception,
__stack_pointer global can point to an invalid address. So
a global.set to restore __stack_pointer should be inserted right
after catch instruction.

But after D57134 the global.set instruction is inserted not right
after catch but after block - br-on-exn - end_block -
extract_exception sequence. This CL fixes it.

While doing that, we can actually move ReplacePhysRegs pass after
LateEHPrepare and merge EHRestoreStackPointer pass into LateEHPrepare,
and now placing global.set to __stack_pointer right after catch is
much easier. Otherwise it is hard to guarantee that global.set is
still right after catch and not touched with other transformations, in
which case we have to do something to hoist it.

[dschuff]

lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp

353

Is this comment still accurate? We don't need so many catch_all instructions now, right?

[aheejin]

• Fix outdated comments

lib/Target/WebAssembly/CMakeLists.txt

add_llvm_target(WebAssemblyCodeGen
   WebAssemblyAddMissingPrototypes.cpp
   WebAssemblyArgumentMove.cpp
   WebAssemblyAsmPrinter.cpp
   WebAssemblyCallIndirectFixup.cpp
   WebAssemblyCFGStackify.cpp
   WebAssemblyCFGSort.cpp
   WebAssemblyDebugValueManager.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 *createWebAssemblyMemIntrinsicResults();
 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 initializeWebAssemblyMemIntrinsicResultsPass(PassRegistry &);
 void initializeWebAssemblyRegStackifyPass(PassRegistry &);
 void initializeWebAssemblyRegColoringPass(PassRegistry &);
 void initializeWebAssemblyExplicitLocalsPass(PassRegistry &);
 void initializeWebAssemblyFixIrreducibleControlFlowPass(PassRegistry &);

lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp

-//===-- WebAssemblyEHRestoreStackPointer.cpp - __stack_pointer restoration ===//
-//
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
-// See https://llvm.org/LICENSE.txt for license information.
-// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
-//
-//===----------------------------------------------------------------------===//
-///
-/// \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();
-    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) {
-  LLVM_DEBUG(dbgs() << "********** EH Restore Stack Pointer **********\n"
-                       "********** Function: "
-                    << MF.getName() << '\n');
-
-  const auto *FrameLowering = static_cast<const WebAssemblyFrameLowering *>(
-      MF.getSubtarget().getFrameLowering());
-  if (!FrameLowering->needsPrologForEH(MF))
-    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, 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 (MBB.begin()->getOpcode() == WebAssembly::CATCH)
-      InsertPos++;
-    FrameLowering->writeSPToGlobal(WebAssembly::SP32, MF, MBB, InsertPos,
-                                   MBB.begin()->getDebugLoc());
-  }
-  return Changed;
-}

lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp

StringRef getPassName() const override {
     return "WebAssembly Late Prepare Exception";
   }
 
   bool runOnMachineFunction(MachineFunction &MF) override;
   bool removeUnnecessaryUnreachables(MachineFunction &MF);
   bool replaceFuncletReturns(MachineFunction &MF);
   bool addCatches(MachineFunction &MF);
   bool addExceptionExtraction(MachineFunction &MF);
+  bool restoreStackPointer(MachineFunction &MF);
 
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyLateEHPrepare() : MachineFunctionPass(ID) {}
 };
 } // end anonymous namespace
 
 char WebAssemblyLateEHPrepare::ID = 0;
bool WebAssemblyLateEHPrepare::runOnMachineFunction(MachineFunction &MF) {
 
   bool Changed = false;
   Changed |= removeUnnecessaryUnreachables(MF);
   if (!MF.getFunction().hasPersonalityFn())
     return Changed;
   Changed |= replaceFuncletReturns(MF);
   Changed |= addCatches(MF);
   Changed |= addExceptionExtraction(MF);
+  Changed |= restoreStackPointer(MF);
   return Changed;
 }
 
 bool WebAssemblyLateEHPrepare::removeUnnecessaryUnreachables(
     MachineFunction &MF) {
   bool Changed = false;
   for (auto &MBB : MF) {
     for (auto &MI : MBB) {
if (TerminatePads.count(EHPad)) {
       BuildMI(ElseMBB, DL, TII.get(WebAssembly::RETHROW));
       if (EHInfo->hasEHPadUnwindDest(EHPad))
         EHInfo->setThrowUnwindDest(ElseMBB, EHInfo->getEHPadUnwindDest(EHPad));
     }
   }
 
   return true;
 }
+
+// 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.
+bool WebAssemblyLateEHPrepare::restoreStackPointer(MachineFunction &MF) {
+  const auto *FrameLowering = static_cast<const WebAssemblyFrameLowering *>(
+      MF.getSubtarget().getFrameLowering());
+  if (!FrameLowering->needsPrologForEH(MF))
+    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, 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 (MBB.begin()->getOpcode() == WebAssembly::CATCH)
+      InsertPos++;
+    FrameLowering->writeSPToGlobal(WebAssembly::SP32, MF, MBB, InsertPos,
+                                   MBB.begin()->getDebugLoc());
+  }
+  return Changed;
+}

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);
   initializeWebAssemblyMemIntrinsicResultsPass(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.
   addPass(createWebAssemblyCallIndirectFixup());
 
   // Eliminate multiple-entry loops.
   addPass(createWebAssemblyFixIrreducibleControlFlow());
 
   // Do various transformations for exception handling.
   // Every CFG-changing optimizations should come before this.
   addPass(createWebAssemblyLateEHPrepare());
 
+  // 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());
+
   // Preparations and optimizations related to register stackification.
   if (getOptLevel() != CodeGenOpt::None) {
     // LiveIntervals isn't commonly run this late. Re-establish preconditions.
     addPass(createWebAssemblyPrepareForLiveIntervals());
 
     // Depend on LiveIntervals and perform some optimizations on it.
     addPass(createWebAssemblyOptimizeLiveIntervals());
 

test/CodeGen/WebAssembly/exception.ll

define void @test_rethrow(i8* %p) {
   ret void
 }
 
 ; CHECK-LABEL: test_catch_rethrow:
 ; CHECK:   global.get  ${{.+}}=, __stack_pointer@GLOBAL
 ; CHECK:   try
 ; CHECK:   call      foo@FUNCTION
 ; CHECK:   catch     $[[EXCEPT_REF:[0-9]+]]=
+; CHECK:   global.set  __stack_pointer@GLOBAL
 ; CHECK:   block i32
 ; CHECK:   br_on_exn 0, __cpp_exception@EVENT, $[[EXCEPT_REF]]
 ; CHECK:   rethrow
 ; CHECK:   end_block
 ; CHECK:   extract_exception $[[EXN:[0-9]+]]=
-; CHECK:   global.set  __stack_pointer@GLOBAL
 ; CHECK-DAG:   i32.store  __wasm_lpad_context
 ; CHECK-DAG:   i32.store  __wasm_lpad_context+4
 ; CHECK:   i32.call  $drop=, _Unwind_CallPersonality@FUNCTION, $[[EXN]]
 ; CHECK:   i32.call  $drop=, __cxa_begin_catch@FUNCTION
 ; CHECK:   call      __cxa_end_catch@FUNCTION
 ; CHECK:   call      __cxa_rethrow@FUNCTION
 ; CHECK:   end_try
 define void @test_catch_rethrow() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {