[WebAssembly] Restore stack pointer right after catch instruction

aheejin · aheejin · commit 0bb986501173 · 2019-01-30T22:44:45.000Z
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 r352598 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. Reviewers: dschuff Subscribers: mgorny, sbc100, jgravelle-google, sunfish, llvm-commits Differential Revision: https://reviews.llvm.org/D57421 llvm-svn: 352681
diff --git a/llvm/lib/Target/WebAssembly/CMakeLists.txt b/llvm/lib/Target/WebAssembly/CMakeLists.txt
@@ -21,7 +21,6 @@ add_llvm_target(WebAssemblyCodeGen
   WebAssemblyCFGSort.cpp
   WebAssemblyDebugValueManager.cpp
   WebAssemblyLateEHPrepare.cpp
-  WebAssemblyEHRestoreStackPointer.cpp
   WebAssemblyExceptionInfo.cpp
   WebAssemblyExplicitLocals.cpp
   WebAssemblyFastISel.cpp
diff --git a/llvm/lib/Target/WebAssembly/WebAssembly.h b/llvm/lib/Target/WebAssembly/WebAssembly.h
@@ -38,7 +38,6 @@ FunctionPass *createWebAssemblyArgumentMove();
 FunctionPass *createWebAssemblySetP2AlignOperands();
 
 // Late passes.
-FunctionPass *createWebAssemblyEHRestoreStackPointer();
 FunctionPass *createWebAssemblyReplacePhysRegs();
 FunctionPass *createWebAssemblyPrepareForLiveIntervals();
 FunctionPass *createWebAssemblyOptimizeLiveIntervals();
@@ -63,7 +62,6 @@ void initializeFixFunctionBitcastsPass(PassRegistry &);
 void initializeOptimizeReturnedPass(PassRegistry &);
 void initializeWebAssemblyArgumentMovePass(PassRegistry &);
 void initializeWebAssemblySetP2AlignOperandsPass(PassRegistry &);
-void initializeWebAssemblyEHRestoreStackPointerPass(PassRegistry &);
 void initializeWebAssemblyReplacePhysRegsPass(PassRegistry &);
 void initializeWebAssemblyPrepareForLiveIntervalsPass(PassRegistry &);
 void initializeWebAssemblyOptimizeLiveIntervalsPass(PassRegistry &);
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyEHRestoreStackPointer.cpp
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp
@@ -34,6 +34,7 @@ class WebAssemblyLateEHPrepare final : public MachineFunctionPass {
   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
@@ -113,6 +114,7 @@ bool WebAssemblyLateEHPrepare::runOnMachineFunction(MachineFunction &MF) {
   Changed |= replaceFuncletReturns(MF);
   Changed |= addCatches(MF);
   Changed |= addExceptionExtraction(MF);
+  Changed |= restoreStackPointer(MF);
   return Changed;
 }
 
@@ -330,3 +332,33 @@ bool WebAssemblyLateEHPrepare::addExceptionExtraction(MachineFunction &MF) {
 
   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. 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;
+}
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyTargetMachine.cpp
@@ -57,7 +57,6 @@ extern "C" void LLVMInitializeWebAssemblyTarget() {
   initializeOptimizeReturnedPass(PR);
   initializeWebAssemblyArgumentMovePass(PR);
   initializeWebAssemblySetP2AlignOperandsPass(PR);
-  initializeWebAssemblyEHRestoreStackPointerPass(PR);
   initializeWebAssemblyReplacePhysRegsPass(PR);
   initializeWebAssemblyPrepareForLiveIntervalsPass(PR);
   initializeWebAssemblyOptimizeLiveIntervalsPass(PR);
@@ -284,14 +283,6 @@ void WebAssemblyPassConfig::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.
@@ -304,6 +295,11 @@ void WebAssemblyPassConfig::addPreEmitPass() {
   // 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.
diff --git a/llvm/test/CodeGen/WebAssembly/exception.ll b/llvm/test/CodeGen/WebAssembly/exception.ll
@@ -30,12 +30,12 @@ define void @test_rethrow(i8* %p) {
 ; 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]]
