[WebAssembly] Fix rethrow's index calculation

llvm/include/llvm/Target/TargetSelectionDAG.td

@@ -234,6 +234,10 @@ def SDTCatchret : SDTypeProfile<0, 2, [     // catchret
   SDTCisVT<0, OtherVT>, SDTCisVT<1, OtherVT>
 ]>;
 
+def SDTCleanupret : SDTypeProfile<0, 1, [   // cleanupret
+  SDTCisVT<0, OtherVT>
+]>;
+
 def SDTNone : SDTypeProfile<0, 0, []>;      // ret, trap
 
 def SDTUBSANTrap : SDTypeProfile<0, 1, []>;      // ubsantrap
@@ -697,7 +701,7 @@ def brind      : SDNode<"ISD::BRIND"      , SDTBrind,  [SDNPHasChain]>;
 def br         : SDNode<"ISD::BR"         , SDTBr,     [SDNPHasChain]>;
 def catchret   : SDNode<"ISD::CATCHRET"   , SDTCatchret,
                         [SDNPHasChain, SDNPSideEffect]>;
-def cleanupret : SDNode<"ISD::CLEANUPRET" , SDTNone,   [SDNPHasChain]>;
+def cleanupret : SDNode<"ISD::CLEANUPRET" , SDTCleanupret, [SDNPHasChain]>;
 
 def trap       : SDNode<"ISD::TRAP"       , SDTNone,
                         [SDNPHasChain, SDNPSideEffect]>;

llvm/lib/CodeGen/SelectionDAG/SelectionDAGBuilder.cpp

@@ -2151,8 +2151,10 @@ void SelectionDAGBuilder::visitCleanupRet(const CleanupReturnInst &I) {
   FuncInfo.MBB->normalizeSuccProbs();
 
   // Create the terminator node.
-  SDValue Ret =
-      DAG.getNode(ISD::CLEANUPRET, getCurSDLoc(), MVT::Other, getControlRoot());
+  MachineBasicBlock *CleanupPadMBB =
+      FuncInfo.getMBB(I.getCleanupPad()->getParent());
+  SDValue Ret = DAG.getNode(ISD::CLEANUPRET, getCurSDLoc(), MVT::Other,
+                            getControlRoot(), DAG.getBasicBlock(CleanupPadMBB));
   DAG.setRoot(Ret);
 }
 

llvm/lib/Target/AArch64/AArch64InstrInfo.td

@@ -5372,7 +5372,7 @@ let isPseudo = 1 in {
 //===----------------------------------------------------------------------===//
 let isTerminator = 1, hasSideEffects = 1, isBarrier = 1, hasCtrlDep = 1,
     isCodeGenOnly = 1, isReturn = 1, isEHScopeReturn = 1, isPseudo = 1 in {
-   def CLEANUPRET : Pseudo<(outs), (ins), [(cleanupret)]>, Sched<[]>;
+   def CLEANUPRET : Pseudo<(outs), (ins), [(cleanupret bb)]>, Sched<[]>;
    let usesCustomInserter = 1 in
      def CATCHRET : Pseudo<(outs), (ins am_brcond:$dst, am_brcond:$src), [(catchret bb:$dst, bb:$src)]>,
                     Sched<[]>;

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

@@ -93,9 +93,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);
@@ -1877,34 +1876,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;
@@ -1916,7 +1894,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.
@@ -1927,6 +1904,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()));
       }
@@ -1953,12 +1932,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;
@@ -1967,15 +1942,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));

llvm/lib/Target/WebAssembly/WebAssemblyISelDAGToDAG.cpp

@@ -248,6 +248,19 @@ void WebAssemblyDAGToDAGISel::Select(SDNode *Node) {
       ReplaceNode(Node, Throw);
       return;
     }
+    case Intrinsic::wasm_rethrow: {
+      // RETHROW's BB argument will be populated in LateEHPrepare. Just use a
+      // '0' as a placeholder for now.
+      MachineSDNode *Rethrow = CurDAG->getMachineNode(
+          WebAssembly::RETHROW, DL,
+          MVT::Other, // outchain type
+          {
+              CurDAG->getConstant(0, DL, MVT::i32), // placeholder
+              Node->getOperand(0)                   // inchain
+          });
+      ReplaceNode(Node, Rethrow);
+      return;
+    }
     }
     break;
   }

llvm/lib/Target/WebAssembly/WebAssemblyInstrControl.td

@@ -168,7 +168,8 @@ defm CATCH_ALL_REF : I<(outs EXNREF:$dst), (ins), (outs), (ins), []>;
 // Pseudo instructions: cleanupret / catchret
 let isTerminator = 1, hasSideEffects = 1, isBarrier = 1, hasCtrlDep = 1,
     isPseudo = 1, isEHScopeReturn = 1 in {
-  defm CLEANUPRET : NRI<(outs), (ins), [(cleanupret)], "cleanupret", 0>;
+  defm CLEANUPRET : NRI<(outs), (ins bb_op:$ehpad), [(cleanupret bb:$ehpad)],
+                        "cleanupret", 0>;
   defm CATCHRET : NRI<(outs), (ins bb_op:$dst, bb_op:$from),
                       [(catchret bb:$dst, bb:$from)], "catchret", 0>;
 } // isTerminator = 1, hasSideEffects = 1, isBarrier = 1, hasCtrlDep = 1,
@@ -180,11 +181,9 @@ let isTerminator = 1, hasSideEffects = 1, isBarrier = 1, hasCtrlDep = 1,
 // Rethrowing an exception: rethrow
 // The new exnref proposal also uses this instruction as an interim pseudo
 // instruction before we convert it to a THROW_REF.
+// $ehpad is the EH pad where the exception the rethrow has been caught.
 let isTerminator = 1, hasCtrlDep = 1, isBarrier = 1 in
-defm RETHROW : NRI<(outs), (ins i32imm:$depth), [], "rethrow \t$depth", 0x09>;
-// The depth argument will be computed in CFGStackify. We set it to 0 here for
-// now.
-def : Pat<(int_wasm_rethrow), (RETHROW 0)>;
+defm RETHROW : NRI<(outs), (ins bb_op:$ehpad), [], "rethrow \t$ehpad", 0x09>;
 
 // Region within which an exception is caught: try / end_try
 let Uses = [VALUE_STACK], Defs = [VALUE_STACK] in {

llvm/lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp

@@ -251,15 +251,40 @@ bool WebAssemblyLateEHPrepare::replaceFuncletReturns(MachineFunction &MF) {
       Changed = true;
       break;
     }
+    case WebAssembly::RETHROW:
+      // These RETHROWs here were lowered from llvm.wasm.rethrow() intrinsics,
+      // generated in Clang for when an exception is not caught by the given
+      // type (e.g. catch (int)).
+      //
+      // RETHROW's BB argument is the EH pad the exception to rethrow has been
+      // caught. (Until this point, RETHROW has just a '0' as a placeholder
+      // argument.) For these llvm.wasm.rethrow()s, we can safely assume the
+      // exception comes from the nearest dominating EH pad, because catch.start
+      // EH pad is structured like this:
+      //
+      // catch.start:
+      //   catchpad ...
+      //   %matches = compare ehselector with typeid
+      //   br i1 %matches, label %catch, label %rethrow
+      //
+      // rethrow:
+      //   ;; rethrows the exception caught in 'catch.start'
+      //   call @llvm.wasm.rethrow()
+      while (TI->getNumOperands() > 0)
+        TI->removeOperand(TI->getNumOperands() - 1);
+      TI->addOperand(MachineOperand::CreateMBB(getMatchingEHPad(TI)));
+      Changed = true;
+      break;
     case WebAssembly::CLEANUPRET: {
-      // Replace a cleanupret with a rethrow. For C++ support, currently
-      // rethrow's immediate argument is always 0 (= the latest exception).
+      // CLEANUPRETs have the EH pad BB the exception to rethrow has been caught
+      // as an argument. Use it and change the instruction opcode to 'RETHROW'
+      // to make rethrowing instructions consistent.
       //
-      // Even when -wasm-enable-exnref is true, we use a RETHROW here for the
-      // moment. This will be converted to a THROW_REF in
-      // addCatchRefsAndThrowRefs.
+      // This is because we cannot safely assume that it is always the nearest
+      // dominating EH pad, in case there are code transformations such as
+      // inlining.
       BuildMI(MBB, TI, TI->getDebugLoc(), TII.get(WebAssembly::RETHROW))
-          .addImm(0);
+          .addMBB(TI->getOperand(0).getMBB());
       TI->eraseFromParent();
       Changed = true;
       break;
@@ -272,21 +297,17 @@ bool WebAssemblyLateEHPrepare::replaceFuncletReturns(MachineFunction &MF) {
 // Add CATCH_REF and CATCH_ALL_REF pseudo instructions to EH pads, and convert
 // RETHROWs to THROW_REFs.
 bool WebAssemblyLateEHPrepare::addCatchRefsAndThrowRefs(MachineFunction &MF) {
-  bool Changed = false;
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
   auto &MRI = MF.getRegInfo();
   DenseMap<MachineBasicBlock *, SmallVector<MachineInstr *, 2>> EHPadToRethrows;
 
   // Create a map of <EH pad, a vector of RETHROWs rethrowing its exception>
-  for (auto &MBB : MF) {
-    for (auto &MI : MBB) {
-      if (MI.getOpcode() == WebAssembly::RETHROW) {
-        Changed = true;
-        auto *EHPad = getMatchingEHPad(&MI);
-        EHPadToRethrows[EHPad].push_back(&MI);
-      }
-    }
-  }
+  for (auto &MBB : MF)
+    for (auto &MI : MBB)
+      if (MI.getOpcode() == WebAssembly::RETHROW)
+        EHPadToRethrows[MI.getOperand(0).getMBB()].push_back(&MI);
+  if (EHPadToRethrows.empty())
+    return false;
 
   // Convert CATCH into CATCH_REF and CATCH_ALL into CATCH_ALL_REF, when the
   // caught exception is rethrown. And convert RETHROWs to THROW_REFs.
@@ -325,7 +346,7 @@ bool WebAssemblyLateEHPrepare::addCatchRefsAndThrowRefs(MachineFunction &MF) {
     }
   }
 
-  return Changed;
+  return true;
 }
 
 // Remove unnecessary unreachables after a throw/rethrow/throw_ref.

llvm/lib/Target/X86/X86InstrCompiler.td

@@ -195,7 +195,8 @@ def EH_RETURN64   : I<0xC3, RawFrm, (outs), (ins GR64:$addr),
 
 let isTerminator = 1, hasSideEffects = 1, isBarrier = 1, hasCtrlDep = 1,
     isCodeGenOnly = 1, isReturn = 1, isEHScopeReturn = 1 in {
-  def CLEANUPRET : I<0, Pseudo, (outs), (ins), "# CLEANUPRET", [(cleanupret)]>;
+  def CLEANUPRET : I<0, Pseudo, (outs), (ins), "# CLEANUPRET",
+                     [(cleanupret bb)]>;
 
   // CATCHRET needs a custom inserter for SEH.
   let usesCustomInserter = 1 in

llvm/test/CodeGen/WebAssembly/cfg-stackify-eh-legacy.mir

@@ -40,15 +40,16 @@ body: |
     ; CHECK: RETHROW 1
     EH_LABEL <mcsymbol .Ltmp2>
     %0:i32 = CATCH_LEGACY &__cpp_exception, implicit-def dead $arguments
-    RETHROW 0, implicit-def dead $arguments
+    RETHROW %bb.1, implicit-def dead $arguments
 
   bb.2 (landing-pad):
+    successors: %bb.3
     ; CHECK: bb.2 (landing-pad):
     ; CHECK: CATCH_LEGACY
     ; CHECK: RETHROW 0
     EH_LABEL <mcsymbol .Ltmp3>
     %1:i32 = CATCH_LEGACY &__cpp_exception, implicit-def dead $arguments
-    RETHROW 0, implicit-def dead $arguments
+    RETHROW %bb.2, implicit-def dead $arguments
 
   bb.3:
     ; CHECK: bb.3:
@@ -105,12 +106,14 @@ body: |
     RETURN %0:i32, implicit-def dead $arguments
 
   bb.3 (landing-pad):
+    successors:
     EH_LABEL <mcsymbol .Ltmp4>
     %0:i32 = CATCH_LEGACY &__cpp_exception, implicit-def dead $arguments
-    RETHROW 0, implicit-def dead $arguments
+    RETHROW %bb.3, implicit-def dead $arguments
 
   bb.4 (landing-pad):
+    successors:
     EH_LABEL <mcsymbol .Ltmp5>
     %1:i32 = CATCH_LEGACY &__cpp_exception, implicit-def dead $arguments
-    RETHROW 0, implicit-def dead $arguments
+    RETHROW %bb.4, implicit-def dead $arguments
 ...