[WebAssemblly] Fix rethrow's argument computation

[aheejin]

Summary

Previously we assumed rethrow's argument was always 0, but it turned
out rethrow follows the same rule with br or delegate:
https://github.com/WebAssembly/exception-handling/pull/137
https://github.com/WebAssembly/exception-handling/issues/146#issuecomment-777349038

Currently rethrows generated by our backend always rethrow the
exception caught by the innermost enclosing catch, so this adds a
function to compute that and replaces rethrow's argument with its
computed result.

This also renames EHPadStack in InstPrinter to TryStack, because
in CFGStackify we use EHPadStack to mean the range between
catch~end, while in InstPrinter we used it to mean the range
between try~catch, so choosing different names would look clearer.
Doesn't contain any functional changes in InstPrinter.

[aheejin]

This replaces D96532. D96532 didn't correctly handle this case:

try
  ...
catch
  try
    rethrow 1 ;; (1)
  catch
    rethrow 0 ;; (2)
  end
end

D96532 correctly computed the argument for (2) but not (1). Because (1) rethrows the exception caught by the outer catch, its argument shouldn't be 0, but D96532 computed it as 0. This CL correctly handles this case.

llvm/test/CodeGen/WebAssembly/exception.mir

49

These are just drive-by changes and not related to rethrow. (CATCHRET's second argument should be the landingpad BB but this was incorrect before, but because we don't really use the second argument this didn't matter.)

[dschuff]

In D96595#2559644, @aheejin wrote:

This replaces D96532. D96532 didn't correctly handle this case:

try
  ...
catch
  try
    rethrow 1 ;; (1)
  catch
    rethrow 0 ;; (2)
  end
end

OK, I think I get the indexing scheme now... I didn't think about 'try' having a label, so if (1) were 'rethrow 0' it would refer to the 'try' (and therefore be invalid)

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

1575

llvm/test/CodeGen/WebAssembly/exception.mir

108

this is "RETHROW 0" in the test input because it's just the placeholder from ISel, right?

[aheejin]

In D96595#2560362, @dschuff wrote:

In D96595#2559644, @aheejin wrote:

This replaces D96532. D96532 didn't correctly handle this case:

try
  ...
catch
  try
    rethrow 1 ;; (1)
  catch
    rethrow 0 ;; (2)
  end
end

OK, I think I get the indexing scheme now... I didn't think about 'try' having a label, so if (1) were 'rethrow 0' it would refer to the 'try' (and therefore be invalid)

Yes, if (1) were rethrow 0, it will refer to the inner try, so it is a validation failure.

I originally hadn't thought about try having a label as well, but as I wrote in https://github.com/WebAssembly/exception-handling/issues/146, it seems convenient and consistent to allow delegate have a label argument that is computed in the same way as branches, and rethrow seems no different from delegate if we look at the rethrow label rule written in https://github.com/WebAssembly/exception-handling/pull/137. I just copied that rule from the very first version of the spec and I myself was misunderstanding that rule, which I realized after reading Thibaud's comment in https://bugs.chromium.org/p/v8/issues/detail?id=11402.

llvm/test/CodeGen/WebAssembly/exception.mir

108

Yes.

[aheejin]

Address comments

[aheejin]

Remove a stray ; in comment

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.h

 
 namespace llvm {
 
 class MCSubtargetInfo;
 
 class WebAssemblyInstPrinter final : public MCInstPrinter {
   uint64_t ControlFlowCounter = 0;
   SmallVector<std::pair<uint64_t, bool>, 4> ControlFlowStack;
-  SmallVector<uint64_t, 4> EHPadStack;
+  SmallVector<uint64_t, 4> TryStack;
 
   enum EHInstKind { TRY, CATCH, CATCH_ALL };
   SmallVector<EHInstKind, 4> EHInstStack;
 
 public:
   WebAssemblyInstPrinter(const MCAsmInfo &MAI, const MCInstrInfo &MII,
                          const MCRegisterInfo &MRI);
 

llvm/lib/Target/WebAssembly/MCTargetDesc/WebAssemblyInstPrinter.cpp

if (CommentStream) {
     case WebAssembly::BLOCK:
     case WebAssembly::BLOCK_S:
       ControlFlowStack.push_back(std::make_pair(ControlFlowCounter++, false));
       return;
 
     case WebAssembly::TRY:
     case WebAssembly::TRY_S:
       ControlFlowStack.push_back(std::make_pair(ControlFlowCounter, false));
-      EHPadStack.push_back(ControlFlowCounter++);
+      TryStack.push_back(ControlFlowCounter++);
       EHInstStack.push_back(TRY);
       return;
 
     case WebAssembly::END_LOOP:
     case WebAssembly::END_LOOP_S:
       if (ControlFlowStack.empty()) {
         printAnnotation(OS, "End marker mismatch!");
       } else {
if (CommentStream) {
     case WebAssembly::CATCH_ALL_S:
       // There can be multiple catch instructions for one try instruction, so
       // we print a label only for the first 'catch' label.
       if (EHInstStack.empty()) {
         printAnnotation(OS, "try-catch mismatch!");
       } else if (EHInstStack.back() == CATCH_ALL) {
         printAnnotation(OS, "catch/catch_all cannot occur after catch_all");
       } else if (EHInstStack.back() == TRY) {
-        if (EHPadStack.empty()) {
+        if (TryStack.empty()) {
           printAnnotation(OS, "try-catch mismatch!");
         } else {
-          printAnnotation(OS,
-                          "catch" + utostr(EHPadStack.pop_back_val()) + ':');
+          printAnnotation(OS, "catch" + utostr(TryStack.pop_back_val()) + ':');
         }
         EHInstStack.pop_back();
         if (Opc == WebAssembly::CATCH || Opc == WebAssembly::CATCH_S) {
           EHInstStack.push_back(CATCH);
         } else {
           EHInstStack.push_back(CATCH_ALL);
         }
       }
       return;
 
     case WebAssembly::RETHROW:
     case WebAssembly::RETHROW_S:
       // 'rethrow' rethrows to the nearest enclosing catch scope, if any. If
       // there's no enclosing catch scope, it throws up to the caller.
-      if (EHPadStack.empty()) {
+      if (TryStack.empty()) {
         printAnnotation(OS, "to caller");
       } else {
-        printAnnotation(OS, "down to catch" + utostr(EHPadStack.back()));
+        printAnnotation(OS, "down to catch" + utostr(TryStack.back()));
       }
       return;
 
     case WebAssembly::DELEGATE:
     case WebAssembly::DELEGATE_S:
-      if (ControlFlowStack.empty() || EHPadStack.empty() ||
-          EHInstStack.empty()) {
+      if (ControlFlowStack.empty() || TryStack.empty() || EHInstStack.empty()) {
         printAnnotation(OS, "try-delegate mismatch!");
       } else {
         // 'delegate' is
         // 1. A marker for the end of block label
         // 2. A destination for throwing instructions
         // 3. An instruction that itself rethrows to another 'catch'
-        assert(ControlFlowStack.back().first == EHPadStack.back());
+        assert(ControlFlowStack.back().first == TryStack.back());
         std::string Label = "label/catch" +
                             utostr(ControlFlowStack.pop_back_val().first) +
                             ": ";
-        EHPadStack.pop_back();
+        TryStack.pop_back();
         EHInstStack.pop_back();
         uint64_t Depth = MI->getOperand(0).getImm();
         if (Depth >= ControlFlowStack.size()) {
           Label += "to caller";
         } else {
           const auto &Pair = ControlFlowStack.rbegin()[Depth];
           if (Pair.second)
             printAnnotation(OS, "delegate cannot target a loop");

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

class WebAssemblyCFGStackify final : public MachineFunctionPass {
 
   // Wrap-up
   using EndMarkerInfo =
       std::pair<const MachineBasicBlock *, const MachineInstr *>;
   unsigned getBranchDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
                           const MachineBasicBlock *MBB);
   unsigned getDelegateDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
                             const MachineBasicBlock *MBB);
+  unsigned
+  getRethrowDepth(const SmallVectorImpl<EndMarkerInfo> &Stack,
+                  const SmallVectorImpl<const MachineBasicBlock *> &EHPadStack);
   void rewriteDepthImmediates(MachineFunction &MF);
   void fixEndsAtEndOfFunction(MachineFunction &MF);
   void cleanupFunctionData(MachineFunction &MF);
 
   // For each BLOCK|LOOP|TRY, the corresponding END_(BLOCK|LOOP|TRY) or DELEGATE
   // (in case of TRY).
   DenseMap<const MachineInstr *, MachineInstr *> BeginToEnd;
   // For each END_(BLOCK|LOOP|TRY) or DELEGATE, the corresponding
for (auto X : reverse(Stack)) {
     if (X.first == EndTry->getParent() && X.second == EndTry)
       break;
     ++Depth;
   }
   assert(Depth < Stack.size() && "Delegate destination should be in scope");
   return Depth;
 }
 
+unsigned WebAssemblyCFGStackify::getRethrowDepth(
+    const SmallVectorImpl<EndMarkerInfo> &Stack,
+    const SmallVectorImpl<const MachineBasicBlock *> &EHPadStack) {
+  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)
+        break;
+    }
+    ++Depth;
+  }
+  assert(Depth < Stack.size() && "Rethrow destination should be in scope");
+  return Depth;
+}
+
 void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {
   // Now rewrite references to basic blocks to be depth immediates.
   SmallVector<EndMarkerInfo, 8> Stack;
+  SmallVector<const MachineBasicBlock *, 8> EHPadStack;
   for (auto &MBB : reverse(MF)) {
     for (auto I = MBB.rbegin(), E = MBB.rend(); I != E; ++I) {
       MachineInstr &MI = *I;
       switch (MI.getOpcode()) {
       case WebAssembly::BLOCK:
       case WebAssembly::TRY:
         assert(ScopeTops[Stack.back().first->getNumber()]->getNumber() <=
                    MBB.getNumber() &&
                "Block/try marker should be balanced");
         Stack.pop_back();
         break;
 
       case WebAssembly::LOOP:
         assert(Stack.back().first == &MBB && "Loop top should be balanced");
         Stack.pop_back();
         break;
 
       case WebAssembly::END_BLOCK:
         Stack.push_back(std::make_pair(&MBB, &MI));
         break;
 
-      case WebAssembly::END_TRY:
+      case WebAssembly::END_TRY: {
         // We handle DELEGATE in the default level, because DELEGATE has
         // immediate operands to rewrite.
         Stack.push_back(std::make_pair(&MBB, &MI));
+        auto *EHPad = TryToEHPad[EndToBegin[&MI]];
+        EHPadStack.push_back(EHPad);
         break;
+      }
 
       case WebAssembly::END_LOOP:
         Stack.push_back(std::make_pair(EndToBegin[&MI]->getParent(), &MI));
         break;
 
+      case WebAssembly::CATCH:
+      case WebAssembly::CATCH_ALL:
+        EHPadStack.pop_back();
+        break;
+
+      case WebAssembly::RETHROW:
+        MI.getOperand(0).setImm(getRethrowDepth(Stack, EHPadStack));
+        break;
+
       default:
         if (MI.isTerminator()) {
           // Rewrite MBB operands to be depth immediates.
           SmallVector<MachineOperand, 4> Ops(MI.operands());
           while (MI.getNumOperands() > 0)
             MI.RemoveOperand(MI.getNumOperands() - 1);
           for (auto MO : Ops) {
             if (MO.isMBB()) {

llvm/lib/Target/WebAssembly/WebAssemblyInstrControl.td

 // Throwing an exception: throw / rethrow
 let isTerminator = 1, hasCtrlDep = 1, isBarrier = 1 in {
 defm THROW : I<(outs), (ins event_op:$tag, variable_ops),
                (outs), (ins event_op:$tag),
                [(WebAssemblythrow (WebAssemblywrapper texternalsym:$tag))],
                "throw   \t$tag", "throw   \t$tag", 0x08>;
 defm RETHROW : NRI<(outs), (ins i32imm:$depth), [], "rethrow \t$depth", 0x09>;
 } // isTerminator = 1, hasCtrlDep = 1, isBarrier = 1
-// For C++ support, we only rethrow the latest exception, thus always setting
-// the depth to 0.
+// The depth argument will be computed in CFGStackify. We set it to 0 here for
+// now.
 def : Pat<(int_wasm_rethrow), (RETHROW 0)>;
 
 // Region within which an exception is caught: try / end_try
 let Uses = [VALUE_STACK], Defs = [VALUE_STACK] in {
 defm TRY     : NRI<(outs), (ins Signature:$sig), [], "try     \t$sig", 0x06>;
 defm END_TRY : NRI<(outs), (ins), [], "end_try", 0x0b>;
 } // Uses = [VALUE_STACK], Defs = [VALUE_STACK]
 

llvm/test/CodeGen/WebAssembly/cfg-stackify-eh.ll

 ; CHECK:             try
 ; CHECK:               call      foo
 ; CHECK:               br        2                     # 2: down to label[[L3:[0-9]+]]
 ; CHECK:             catch
 ; CHECK:               call      __cxa_end_catch
 ; CHECK:               rethrow   0                     # down to catch[[C0:[0-9]+]]
 ; CHECK:             end_try
 ; CHECK:           end_block                           # label[[L2]]:
-; CHECK:           rethrow   0                         # down to catch[[C0]]
+; CHECK:           rethrow   1                         # down to catch[[C0]]
 ; CHECK:         catch_all                             # catch[[C0]]:
 ; CHECK:           call      __cxa_end_catch
 ; CHECK:           rethrow   0                         # to caller
 ; CHECK:         end_try                               # label[[L3]]:
 ; CHECK:         call      __cxa_end_catch
 ; CHECK:         br        2                           # 2: down to label[[L1]]
 ; CHECK:       end_try
 ; CHECK:     end_block                                 # label[[L0]]:
-; CHECK:     rethrow   0                               # to caller
+; CHECK:     rethrow   1                               # to caller
 ; CHECK:   end_block                                   # label[[L1]]:
 ; CHECK:   call      __cxa_end_catch
 ; CHECK: end_try
 define void @test1() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   invoke void @foo()
           to label %try.cont11 unwind label %catch.dispatch
 

llvm/test/CodeGen/WebAssembly/exception.mir

 # RUN: llc -mtriple=wasm32-unknown-unknown -exception-model=wasm -mattr=+exception-handling -run-pass wasm-late-eh-prepare -run-pass wasm-cfg-stackify %s -o - | FileCheck %s
 
 --- |
   target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
   target triple = "wasm32-unknown-unknown"
 
   declare i32 @__gxx_wasm_personality_v0(...)
   declare void @foo()
   define void @eh_label_test() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
     ret void
   }
   define void @unreachable_ehpad_test() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
     ret void
   }
+  define void @rethrow_arg_test() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+    ret void
+  }
 ...
 
 ---
 # This tests 'try' and 'catch' instructions are correctly placed with respect to
 # EH_LABEL instructions.
 # CHECK-LABEL: name: eh_label_test
 name: eh_label_test
 liveins:
   - { reg: '$arguments' }
 body: |
   bb.0:
     ; TRY should be before EH_LABEL wrappers of throwing calls
+    ; CHECK:      bb.0
     ; CHECK:      TRY
     ; CHECK-NEXT: EH_LABEL
     ; CHECK-NEXT: CALL @foo
     ; CHECK-NEXT: EH_LABEL
     successors: %bb.1, %bb.2
     EH_LABEL <mcsymbol .Ltmp0>
     CALL @foo, implicit-def dead $arguments, implicit $sp32, implicit $sp64
     EH_LABEL <mcsymbol .Ltmp1>
     BR %bb.2, implicit-def dead $arguments
 
   bb.1 (landing-pad):
   ; predecessors: %bb.0
     successors: %bb.2
     ; CATCH_ALL should be after an EH_LABEL at the beginning of an EH pad
+    ; CHECK:      bb.1
     ; CHECK:      EH_LABEL
     ; CHECK-NEXT: CATCH_ALL
     EH_LABEL <mcsymbol .Ltmp2>
-    CATCHRET %bb.2, %bb.0, implicit-def dead $arguments
+    CATCHRET %bb.2, %bb.1, implicit-def dead $arguments
 
   bb.2:
   ; predecessors: %bb.0, %bb.1
     RETURN implicit-def dead $arguments
 ...
+
 ---
 # Unreachable EH pads should be removed by LateEHPrepare.
 # CHECK-LABEL: name: unreachable_ehpad_test
 name: unreachable_ehpad_test
 liveins:
   - { reg: '$arguments' }
 body: |
   ; CHECK: bb.0
   bb.0:
     successors: %bb.2
     BR %bb.2, implicit-def dead $arguments
 
   ; This EH pad is unreachable, so it should be removed by LateEHPrepare
   ; CHECK-NOT: bb.1 (landing-pad)
   bb.1 (landing-pad):
     successors: %bb.2
     EH_LABEL <mcsymbol .Ltmp2>
-    CATCHRET %bb.2, %bb.0, implicit-def dead $arguments
+    CATCHRET %bb.2, %bb.1, implicit-def dead $arguments
 
   ; CHECK: bb.2
   bb.2:
   ; predecessors: %bb.0, %bb.1
     RETURN implicit-def dead $arguments
 ...
+
+---
+# CHECK-LABEL: name: rethrow_arg_test
+name: rethrow_arg_test
+liveins:
+  - { reg: '$arguments' }
+body: |
+  bb.0:
+    successors: %bb.1, %bb.4
+    ; CHECK: bb.0
+    ; CHECK: TRY
+    EH_LABEL <mcsymbol .Ltmp0>
+    CALL @foo, implicit-def dead $arguments, implicit $sp32, implicit $sp64
+    EH_LABEL <mcsymbol .Ltmp1>
+    BR %bb.4, implicit-def dead $arguments
+
+  bb.1 (landing-pad):
+  ; predecessors: %bb.0
+    successors: %bb.2
+    ; CHECK: bb.1
+    ; CHECK: CATCH
+    ; CHECK: TRY
+    ; This RETHROW rethrows the exception caught by this BB's CATCH, but after
+    ; CFGStackify a TRY is placed between the CATCH and this RETHROW, so after
+    ; CFGStackify its immediate argument should become not 0, but 1.
+    ; CHECK: RETHROW 1
+    EH_LABEL <mcsymbol .Ltmp2>
+    %0:i32 = CATCH &__cpp_exception, implicit-def dead $arguments
+    RETHROW 0, implicit-def dead $arguments
+
+  bb.2 (landing-pad):
+  ; predecessors: %bb.1
+    successors: %bb.3
+    ; CHECK: bb.2
+    ; CHECK: CATCH
+    ; CHECK: RETHROW 0
+    EH_LABEL <mcsymbol .Ltmp5>
+    %1:i32 = CATCH &__cpp_exception, implicit-def dead $arguments
+    RETHROW 0, implicit-def dead $arguments
+    CATCHRET %bb.3, %bb.2, implicit-def dead $arguments
+
+  bb.3:
+  ; predecessors: %bb.2
+    successors: %bb.4
+    CATCHRET %bb.4, %bb.1, implicit-def dead $arguments
+
+  bb.4:
+  ; predecessors: %bb.0, %bb.3
+    ; CHECK: bb.4
+    ; CHECK: END_TRY
+    ; CHECK: END_TRY
+    RETURN implicit-def dead $arguments
+