[WebAssembly] Fix rethrow's index calculation (#114693)

So far we have assumed that we only rethrow the exception caught in the
innermost EH pad. This is true in code we directly generate, but after
inlining this may not be the case. For example, consider this code:
```ll
ehcleanup:
  %0 = cleanuppad ...
  call @destructor
  cleanupret from %0 unwind label %catch.dispatch
```

If `destructor` gets inlined into this function, the code can be like
```ll
ehcleanup:
  %0 = cleanuppad ...
  invoke @throwing_func
    to label %unreachale unwind label %catch.dispatch.i

catch.dispatch.i:
  catchswitch ... [ label %catch.start.i ]

catch.start.i:
  %1 = catchpad ...
  invoke @some_function
    to label %invoke.cont.i unwind label %terminate.i

invoke.cont.i:
  catchret from %1 to label %destructor.exit

destructor.exit:
  cleanupret from %0 unwind label %catch.dispatch
```

We lower a `cleanupret` into `rethrow`, which assumes it rethrows the
exception caught by the nearest dominating EH pad. But after the
inlining, the nearest dominating EH pad is not `ehcleanup` but
`catch.start.i`.

The problem exists in the same manner in the new (exnref) EH, because it
assumes the exception comes from the nearest EH pad and saves an exnref
from that EH pad and rethrows it (using `throw_ref`).

This problem can be fixed easily if `cleanupret` has the basic block
where its matching `cleanuppad` is. The bitcode instruction `cleanupret`
kind of has that info (it has a token from the `cleanuppad`), but that
info is lost when when we enter ISel, because `TargetSelectionDAG.td`'s
`cleanupret` node does not have any arguments:

https://github.com/llvm/llvm-project/blob/5091a359d9807db8f7d62375696f93fc34226969/llvm/include/llvm/Target/TargetSelectionDAG.td#L700
Note that `catchret` already has two basic block arguments, even though
neither of them means `catchpad`'s BB.

This PR adds the `cleanuppad`'s BB as an argument to `cleanupret` node
in ISel and uses it in the Wasm backend. Because this node is also used
in X86 backend we need to note its argument there too but nothing more
needs to change there as long as X86 doesn't need it.

---

- Details about changes in the Wasm backend:

After this PR, our pseudo `RETHROW` instruction takes a BB, which means
the EH pad whose exception it needs to rethrow. There are currently two
ways to generate a `RETHROW`: one is from `llvm.wasm.rethrow` intrinsic
and the other is from `CLEANUPRET` we discussed above. In case of
`llvm.wasm.rethrow`, we add a '0' as a placeholder argument when it is
lowered to a `RETHROW`, and change it to a BB in LateEHPrepare. As
written in the comments, this PR doesn't change how this BB is computed.
The BB argument will be converted to an immediate argument as with other
control flow instructions in CFGStackify.

In case of `CLEANUPRET`, it already has a BB argument pointing to an EH
pad, so it is just converted to a `RETHROW` with the same BB argument in
LateEHPrepare. This will also be lowered to an immediate in CFGStackify
with other control flow instructions.

---

Fixes #114600.

1 files changed, 7 insertions, 41 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
index 6c13beb..90b5945 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp
@@ -99,9 +99,8 @@ class WebAssemblyCFGStackify final : public MachineFunctionPass {
                           const MachineBasicBlock *MBB);
   unsigned getDelegateDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
                             const MachineBasicBlock *MBB);
-  unsigned
-  getRethrowDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
-                  const SmallVectorImpl<const MachineBasicBlock *> &EHPadStack);
+  unsigned getRethrowDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
+                           const MachineBasicBlock *EHPadToRethrow);
   void rewriteDepthImmediates(MachineFunction &MF);
   void fixEndsAtEndOfFunction(MachineFunction &MF);
   void cleanupFunctionData(MachineFunction &MF);
@@ -2458,34 +2457,13 @@ unsigned WebAssemblyCFGStackify::getDelegateDepth(
 
 unsigned WebAssemblyCFGStackify::getRethrowDepth(
     const SmallVectorImpl<EndMarkerInfo> &Stack,
-    const SmallVectorImpl<const MachineBasicBlock *> &EHPadStack) {
+    const MachineBasicBlock *EHPadToRethrow) {
   unsigned Depth = 0;
-  // In our current implementation, rethrows always rethrow the exception caught
-  // by the innermost enclosing catch. This means while traversing Stack in the
-  // reverse direction, when we encounter END_TRY, we should check if the
-  // END_TRY corresponds to the current innermost EH pad. For example:
-  // try
-  //   ...
-  // catch         ;; (a)
-  //   try
-  //     rethrow 1 ;; (b)
-  //   catch       ;; (c)
-  //     rethrow 0 ;; (d)
-  //   end         ;; (e)
-  // end           ;; (f)
-  //
-  // When we are at 'rethrow' (d), while reversely traversing Stack the first
-  // 'end' we encounter is the 'end' (e), which corresponds to the 'catch' (c).
-  // And 'rethrow' (d) rethrows the exception caught by 'catch' (c), so we stop
-  // there and the depth should be 0. But when we are at 'rethrow' (b), it
-  // rethrows the exception caught by 'catch' (a), so when traversing Stack
-  // reversely, we should skip the 'end' (e) and choose 'end' (f), which
-  // corresponds to 'catch' (a).
   for (auto X : reverse(Stack)) {
     const MachineInstr *End = X.second;
     if (End->getOpcode() == WebAssembly::END_TRY) {
       auto *EHPad = TryToEHPad[EndToBegin[End]];
-      if (EHPadStack.back() == EHPad)
+      if (EHPadToRethrow == EHPad)
         break;
     }
     ++Depth;
@@ -2497,7 +2475,6 @@ unsigned WebAssemblyCFGStackify::getRethrowDepth(
 void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
   // Now rewrite references to basic blocks to be depth immediates.
   SmallVector<EndMarkerInfo, 8> Stack;
-  SmallVector<const MachineBasicBlock *, 8> EHPadStack;
 
   auto RewriteOperands = [&](MachineInstr &MI) {
     // Rewrite MBB operands to be depth immediates.
@@ -2508,6 +2485,8 @@ void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
       if (MO.isMBB()) {
         if (MI.getOpcode() == WebAssembly::DELEGATE)
           MO = MachineOperand::CreateImm(getDelegateDepth(Stack, MO.getMBB()));
+        else if (MI.getOpcode() == WebAssembly::RETHROW)
+          MO = MachineOperand::CreateImm(getRethrowDepth(Stack, MO.getMBB()));
         else
           MO = MachineOperand::CreateImm(getBranchDepth(Stack, MO.getMBB()));
       }
@@ -2539,12 +2518,8 @@ void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
         Stack.pop_back();
         break;
 
-      case WebAssembly::END_TRY: {
-        auto *EHPad = TryToEHPad[EndToBegin[&MI]];
-        EHPadStack.push_back(EHPad);
-        [[fallthrough]];
-      }
       case WebAssembly::END_BLOCK:
+      case WebAssembly::END_TRY:
       case WebAssembly::END_TRY_TABLE:
         Stack.push_back(std::make_pair(&MBB, &MI));
         break;
@@ -2553,15 +2528,6 @@ void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
         Stack.push_back(std::make_pair(EndToBegin[&MI]->getParent(), &MI));
         break;
 
-      case WebAssembly::CATCH_LEGACY:
-      case WebAssembly::CATCH_ALL_LEGACY:
-        EHPadStack.pop_back();
-        break;
-
-      case WebAssembly::RETHROW:
-        MI.getOperand(0).setImm(getRethrowDepth(Stack, EHPadStack));
-        break;
-
       case WebAssembly::DELEGATE:
         RewriteOperands(MI);
         Stack.push_back(std::make_pair(&MBB, &MI));