[WebAssembly] Fix phi handling for Wasm SjLj (#99730)

aheejin · web-flow · commit 2bf71b8bc851 · 2024-07-23T16:06:00.000-07:00
In Wasm SjLj, longjmpable `call`s that in functions that call `setjmp` are converted into `invoke`s. Those `invoke`s are meant to unwind to `catch.dispatch.longjmp` to figure out which `setjmp` those `longjmp` buffers belong to: https://github.com/llvm/llvm-project/blob/fada9227325b3eaa0bdc09a486f29a7f08b7b3fb/llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp#L250-L260 But in case a longjmpable call is within another `catchpad` or `cleanuppad` scope, to maintain the nested scope structure, we should make them unwind to the scope's next unwind destination and not directly to `catch.dispatch.longjmp`: https://github.com/llvm/llvm-project/blob/fada9227325b3eaa0bdc09a486f29a7f08b7b3fb/llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp#L1698-L1727 In this case the longjmps will eventually unwind to `catch.dispatch.longjmp` and be handled there. In this case, it is possible that the unwind destination (which is an existing `catchpad` or `cleanuppad`) may already have `phi`s. And because the unwind destinations get new predecessors because of the newly created `invoke`s, those `phi`s need to have new entries for those new predecessors. This adds new preds as new incoming blocks to those `phi`s, and we use a separate `SSAUpdater` to calculate the correct incoming values to those blocks. I have assumed `SSAUpdaterBulk` used in `rebuildSSA` would take care of these things, but apparently it doesn't. It takes available defs and adds `phi`s in the defs' dominance frontiers, i.e., where each def's dominance ends, and rewrites other uses based on the newly added `phi`s. But it doesn't add entries to existing `phi`s, and the case in this bug may not even involve dominance frontiers; this bug is simply about existing `phis`s that have gained new preds need new entries for them. It is kind of surprising that this bug was only reported recently, given that this pass has not been changed much in years. Fixes #97496 and fixes emscripten-core/emscripten#22170.
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyLowerEmscriptenEHSjLj.cpp
@@ -777,6 +777,8 @@ void WebAssemblyLowerEmscriptenEHSjLj::rebuildSSA(Function &F) {
   SSAUpdaterBulk SSA;
   for (BasicBlock &BB : F) {
     for (Instruction &I : BB) {
+      if (I.getType()->isVoidTy())
+        continue;
       unsigned VarID = SSA.AddVariable(I.getName(), I.getType());
       // If a value is defined by an invoke instruction, it is only available in
       // its normal destination and not in its unwind destination.
@@ -1687,6 +1689,8 @@ void WebAssemblyLowerEmscriptenEHSjLj::handleLongjmpableCallsForWasmSjLj(
     }
   }
 
+  SmallDenseMap<BasicBlock *, SmallPtrSet<BasicBlock *, 4>, 4>
+      UnwindDestToNewPreds;
   for (auto *CI : LongjmpableCalls) {
     // Even if the callee function has attribute 'nounwind', which is true for
     // all C functions, it can longjmp, which means it can throw a Wasm
@@ -1724,6 +1728,11 @@ void WebAssemblyLowerEmscriptenEHSjLj::handleLongjmpableCallsForWasmSjLj(
     }
     if (!UnwindDest)
       UnwindDest = CatchDispatchLongjmpBB;
+    // Because we are changing a longjmpable call to an invoke, its unwind
+    // destination can be an existing EH pad that already have phis, and the BB
+    // with the newly created invoke will become a new predecessor of that EH
+    // pad. In this case we need to add the new predecessor to those phis.
+    UnwindDestToNewPreds[UnwindDest].insert(CI->getParent());
     changeToInvokeAndSplitBasicBlock(CI, UnwindDest);
   }
 
@@ -1752,4 +1761,45 @@ void WebAssemblyLowerEmscriptenEHSjLj::handleLongjmpableCallsForWasmSjLj(
 
   for (Instruction *I : ToErase)
     I->eraseFromParent();
+
+  // Add entries for new predecessors to phis in unwind destinations. We use
+  // 'undef' as a placeholder value. We should make sure the phis have a valid
+  // set of predecessors before running SSAUpdater, because SSAUpdater
+  // internally can use existing phis to gather predecessor info rather than
+  // scanning the actual CFG (See FindPredecessorBlocks in SSAUpdater.cpp for
+  // details).
+  for (auto &[UnwindDest, NewPreds] : UnwindDestToNewPreds) {
+    for (PHINode &PN : UnwindDest->phis()) {
+      for (auto *NewPred : NewPreds) {
+        assert(PN.getBasicBlockIndex(NewPred) == -1);
+        PN.addIncoming(UndefValue::get(PN.getType()), NewPred);
+      }
+    }
+  }
+
+  // For unwind destinations for newly added invokes to longjmpable functions,
+  // calculate incoming values for the newly added predecessors using
+  // SSAUpdater. We add existing values in the phis to SSAUpdater as available
+  // values and let it calculate what the value should be at the end of new
+  // incoming blocks.
+  for (auto &[UnwindDest, NewPreds] : UnwindDestToNewPreds) {
+    for (PHINode &PN : UnwindDest->phis()) {
+      SSAUpdater SSA;
+      SSA.Initialize(PN.getType(), PN.getName());
+      for (unsigned Idx = 0, E = PN.getNumIncomingValues(); Idx != E; ++Idx) {
+        if (NewPreds.contains(PN.getIncomingBlock(Idx)))
+          continue;
+        Value *V = PN.getIncomingValue(Idx);
+        if (auto *II = dyn_cast<InvokeInst>(V))
+          SSA.AddAvailableValue(II->getNormalDest(), II);
+        else if (auto *I = dyn_cast<Instruction>(V))
+          SSA.AddAvailableValue(I->getParent(), I);
+        else
+          SSA.AddAvailableValue(PN.getIncomingBlock(Idx), V);
+      }
+      for (auto *NewPred : NewPreds)
+        PN.setIncomingValueForBlock(NewPred, SSA.GetValueAtEndOfBlock(NewPred));
+      assert(PN.isComplete());
+    }
+  }
 }
diff --git a/llvm/test/CodeGen/WebAssembly/lower-wasm-ehsjlj-phi.ll b/llvm/test/CodeGen/WebAssembly/lower-wasm-ehsjlj-phi.ll
@@ -0,0 +1,126 @@
+; RUN: opt < %s -wasm-lower-em-ehsjlj -wasm-enable-eh -wasm-enable-sjlj -S | FileCheck %s
+
+target triple = "wasm32-unknown-emscripten"
+
+%struct.__jmp_buf_tag = type { [6 x i32], i32, [32 x i32] }
+@buf = internal global [1 x %struct.__jmp_buf_tag] zeroinitializer, align 16
+
+; When longjmpable calls are coverted into invokes in Wasm SjLj transformation
+; and their unwind destination is an existing catchpad or cleanuppad due to
+; maintain the scope structure, the new pred BBs created by invokes and the
+; correct incoming values should be added the existing phis in those unwind
+; destinations.
+
+; When longjmpable calls are within a cleanuppad.
+define void @longjmpable_invoke_phi0() personality ptr @__gxx_wasm_personality_v0 {
+; CHECK-LABEL: @longjmpable_invoke_phi0
+entry:
+  %val.entry = call i32 @llvm.wasm.memory.size.i32(i32 0)
+  %0 = call i32 @setjmp(ptr @buf) #2
+  invoke void @foo()
+          to label %bb1 unwind label %ehcleanup1
+
+bb1:                                              ; preds = %entry
+  ; We use llvm.wasm.memory.size intrinsic just to get/use an i32 value. The
+  ; reason we use an intrinsic here is to make it not longjmpable. If this can
+  ; longjmp, the result will be more complicated and hard to check.
+  %val.bb1 = call i32 @llvm.wasm.memory.size.i32(i32 0)
+  invoke void @foo()
+          to label %bb2 unwind label %ehcleanup0
+
+bb2:                                              ; preds = %bb1
+  unreachable
+
+ehcleanup0:                                       ; preds = %bb1
+  %1 = cleanuppad within none []
+  call void @longjmpable() [ "funclet"(token %1) ]
+; CHECK:      ehcleanup0
+; CHECK:        invoke void @longjmpable
+; CHECK-NEXT:           to label %.noexc unwind label %ehcleanup1
+  invoke void @foo() [ "funclet"(token %1) ]
+          to label %bb3 unwind label %ehcleanup1
+
+bb3:                                              ; preds = %ehcleanup0
+  %val.bb3 = call i32 @llvm.wasm.memory.size.i32(i32 0)
+  call void @longjmpable() [ "funclet"(token %1) ]
+; CHECK:      bb3:
+; CHECK:        invoke void @longjmpable
+; CHECK-NEXT:           to label %.noexc1 unwind label %ehcleanup1
+  cleanupret from %1 unwind label %ehcleanup1
+
+ehcleanup1:                                       ; preds = %bb3, %ehcleanup0, %entry
+  %phi = phi i32 [ %val.entry, %entry ], [ %val.bb1, %ehcleanup0 ], [ %val.bb3, %bb3 ]
+; CHECK:      ehcleanup1:
+; CHECK-NEXT:   %phi = phi i32 [ %val.entry2, %entry.split.split ], [ %val.bb1, %.noexc ], [ %val.bb3, %.noexc1 ], [ %val.bb1, %ehcleanup0 ], [ %val.bb3, %bb3 ]
+  %2 = cleanuppad within none []
+  %3 = call i32 @llvm.wasm.memory.size.i32(i32 %phi)
+  cleanupret from %2 unwind to caller
+}
+
+; When longjmpable calls are within a catchpad.
+define void @longjmpable_invoke_phi1() personality ptr @__gxx_wasm_personality_v0 {
+; CHECK-LABEL: @longjmpable_invoke_phi1
+entry:
+  %val.entry = call i32 @llvm.wasm.memory.size.i32(i32 0)
+  %0 = call i32 @setjmp(ptr @buf) #2
+  invoke void @foo()
+          to label %bb1 unwind label %ehcleanup
+
+bb1:                                              ; preds = %entry
+  %val.bb1 = call i32 @llvm.wasm.memory.size.i32(i32 0)
+  invoke void @foo()
+          to label %bb2 unwind label %catch.dispatch
+
+bb2:                                              ; preds = %bb1
+  unreachable
+
+catch.dispatch:                                   ; preds = %bb1
+  %1 = catchswitch within none [label %catch.start] unwind label %ehcleanup
+
+catch.start:                                      ; preds = %catch.dispatch
+  %2 = catchpad within %1 [ptr null]
+  %3 = call ptr @llvm.wasm.get.exception(token %2)
+  %4 = call i32 @llvm.wasm.get.ehselector(token %2)
+  call void @longjmpable() [ "funclet"(token %2) ]
+; CHECK:      catch.start:
+; CHECK:        invoke void @longjmpable
+; CHECK-NEXT:           to label %.noexc unwind label %ehcleanup
+  invoke void @foo() [ "funclet"(token %2) ]
+          to label %bb3 unwind label %ehcleanup
+
+bb3:                                              ; preds = %catch.start
+  %val.bb3 = call i32 @llvm.wasm.memory.size.i32(i32 0)
+  call void @longjmpable() [ "funclet"(token %2) ]
+; CHECK:      bb3:
+; CHECK:        invoke void @longjmpable
+; CHECK-NEXT:           to label %.noexc1 unwind label %ehcleanup
+  invoke void @foo() [ "funclet"(token %2) ]
+          to label %bb4 unwind label %ehcleanup
+
+bb4:                                              ; preds = %bb3
+  unreachable
+
+ehcleanup:                                        ; preds = %bb3, %catch.start, %catch.dispatch, %entry
+  %phi = phi i32 [ %val.entry, %entry ], [ %val.bb1, %catch.dispatch ], [ %val.bb1, %catch.start ], [ %val.bb3, %bb3 ]
+; CHECK:      ehcleanup:
+; CHECK-NEXT:   %phi = phi i32 [ %val.entry2, %entry.split.split ], [ %val.bb1, %catch.dispatch ], [ %val.bb1, %.noexc ], [ %val.bb3, %.noexc1 ], [ %val.bb1, %catch.start ], [ %val.bb3, %bb3 ]
+  %5 = cleanuppad within none []
+  %6 = call i32 @llvm.wasm.memory.size.i32(i32 %phi)
+  cleanupret from %5 unwind to caller
+}
+
+declare i32 @setjmp(ptr)
+declare i32 @__gxx_wasm_personality_v0(...)
+declare void @foo()
+declare void @longjmpable()
+declare void @use_i32(i32)
+; Function Attrs: nocallback nofree nosync nounwind willreturn
+declare i32 @llvm.wasm.get.ehselector(token) #0
+; Function Attrs: nocallback nofree nosync nounwind willreturn
+declare ptr @llvm.wasm.get.exception(token) #0
+; Function Attrs: nocallback nofree nosync nounwind willreturn memory(read)
+declare i32 @llvm.wasm.memory.size.i32(i32) #1
+
+attributes #0 = { nocallback nofree nosync nounwind willreturn }
+attributes #1 = { nocallback nofree nosync nounwind willreturn memory(read) }
+attributes #2 = { returns_twice }
