[WebAssembly] Fix bugs in BLOCK/TRY placement

Summary:
Before we placed all TRY/END_TRY markers before placing BLOCK/END_BLOCK
markers. This couldn't handle this case:
```
bb0:
  br bb2
bb1:          // nearest common dominator of bb3 and bb4
  br_if ... bb3
  br bb4
bb2:
  ...
bb3:
  call @foo   // unwinds to ehpad
bb4:
  call @bar   // unwinds to ehpad
ehpad:
  catch
  ...
```

When we placed TRY markers, we placed it in bb1 because it is the
nearest common dominator of bb3 and bb4. But because bb0 jumps to bb2,
when we placed block markers, we ended up with interleaved scopes like
```
block
try
end_block
catch
end_try
```
which was not correct.

This patch fixes the bug by placing BLOCK and TRY markers in one pass
while iterating BBs in a function. This also adds some more routines to
`placeTryMarkers`, because we now have to assume that there can be
previously placed BLOCK and END_BLOCK.

Reviewers: dschuff

Subscribers: sunfish, sbc100, jgravelle-google, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D59739

llvm-svn: 357007

Author: aheejin

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

@@ -193,10 +193,7 @@ void WebAssemblyCFGStackify::unregisterScope(MachineInstr *Begin) {
 
 /// Insert a BLOCK marker for branches to MBB (if needed).
 void WebAssemblyCFGStackify::placeBlockMarker(MachineBasicBlock &MBB) {
-  // This should have been handled in placeTryMarker.
-  if (MBB.isEHPad())
-    return;
-
+  assert(!MBB.isEHPad());
   MachineFunction &MF = *MBB.getParent();
   auto &MDT = getAnalysis<MachineDominatorTree>();
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
@@ -253,24 +250,23 @@ void WebAssemblyCFGStackify::placeBlockMarker(MachineBasicBlock &MBB) {
   // Instructions that should go after the BLOCK.
   SmallPtrSet<const MachineInstr *, 4> AfterSet;
   for (const auto &MI : *Header) {
-    // If there is a previously placed LOOP/TRY marker and the bottom block of
-    // the loop/exception is above MBB, it should be after the BLOCK, because
-    // the loop/exception is nested in this block. Otherwise it should be before
-    // the BLOCK.
-    if (MI.getOpcode() == WebAssembly::LOOP ||
-        MI.getOpcode() == WebAssembly::TRY) {
-      auto *BottomBB = BeginToEnd[&MI]->getParent()->getPrevNode();
-      if (MBB.getNumber() > BottomBB->getNumber())
+    // If there is a previously placed LOOP marker and the bottom block of the
+    // loop is above MBB, it should be after the BLOCK, because the loop is
+    // nested in this BLOCK. Otherwise it should be before the BLOCK.
+    if (MI.getOpcode() == WebAssembly::LOOP) {
+      auto *LoopBottom = BeginToEnd[&MI]->getParent()->getPrevNode();
+      if (MBB.getNumber() > LoopBottom->getNumber())
         AfterSet.insert(&MI);
 #ifndef NDEBUG
       else
         BeforeSet.insert(&MI);
 #endif
     }
 
-    // All previously inserted BLOCK markers should be after the BLOCK because
-    // they are all nested blocks.
-    if (MI.getOpcode() == WebAssembly::BLOCK)
+    // All previously inserted BLOCK/TRY markers should be after the BLOCK
+    // because they are all nested blocks.
+    if (MI.getOpcode() == WebAssembly::BLOCK ||
+        MI.getOpcode() == WebAssembly::TRY)
       AfterSet.insert(&MI);
 
 #ifndef NDEBUG
@@ -428,9 +424,7 @@ void WebAssemblyCFGStackify::placeLoopMarker(MachineBasicBlock &MBB) {
 }
 
 void WebAssemblyCFGStackify::placeTryMarker(MachineBasicBlock &MBB) {
-  if (!MBB.isEHPad())
-    return;
-
+  assert(MBB.isEHPad());
   MachineFunction &MF = *MBB.getParent();
   auto &MDT = getAnalysis<MachineDominatorTree>();
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
@@ -486,31 +480,34 @@ void WebAssemblyCFGStackify::placeTryMarker(MachineBasicBlock &MBB) {
 
   // Decide where in Header to put the TRY.
 
-  // Instructions that should go before the BLOCK.
+  // Instructions that should go before the TRY.
   SmallPtrSet<const MachineInstr *, 4> BeforeSet;
-  // Instructions that should go after the BLOCK.
+  // Instructions that should go after the TRY.
   SmallPtrSet<const MachineInstr *, 4> AfterSet;
   for (const auto &MI : *Header) {
-    // If there is a previously placed LOOP marker and the bottom block of
-    // the loop is above MBB, the LOOP should be after the TRY, because the
-    // loop is nested in this try. Otherwise it should be before the TRY.
+    // If there is a previously placed LOOP marker and the bottom block of the
+    // loop is above MBB, it should be after the TRY, because the loop is nested
+    // in this TRY. Otherwise it should be before the TRY.
     if (MI.getOpcode() == WebAssembly::LOOP) {
-      if (MBB.getNumber() > Bottom->getNumber())
+      auto *LoopBottom = BeginToEnd[&MI]->getParent()->getPrevNode();
+      if (MBB.getNumber() > LoopBottom->getNumber())
         AfterSet.insert(&MI);
 #ifndef NDEBUG
       else
         BeforeSet.insert(&MI);
 #endif
     }
 
-    // All previously inserted TRY markers should be after the TRY because they
-    // are all nested trys.
-    if (MI.getOpcode() == WebAssembly::TRY)
+    // All previously inserted BLOCK/TRY markers should be after the TRY because
+    // they are all nested trys.
+    if (MI.getOpcode() == WebAssembly::BLOCK ||
+        MI.getOpcode() == WebAssembly::TRY)
       AfterSet.insert(&MI);
 
 #ifndef NDEBUG
-    // All END_(LOOP/TRY) markers should be before the TRY.
-    if (MI.getOpcode() == WebAssembly::END_LOOP ||
+    // All END_(BLOCK/LOOP/TRY) markers should be before the TRY.
+    if (MI.getOpcode() == WebAssembly::END_BLOCK ||
+        MI.getOpcode() == WebAssembly::END_LOOP ||
         MI.getOpcode() == WebAssembly::END_TRY)
       BeforeSet.insert(&MI);
 #endif
@@ -566,8 +563,9 @@ void WebAssemblyCFGStackify::placeTryMarker(MachineBasicBlock &MBB) {
   AfterSet.clear();
   for (const auto &MI : *Cont) {
 #ifndef NDEBUG
-    // END_TRY should precede existing LOOP markers.
-    if (MI.getOpcode() == WebAssembly::LOOP)
+    // END_TRY should precede existing LOOP and BLOCK markers.
+    if (MI.getOpcode() == WebAssembly::LOOP ||
+        MI.getOpcode() == WebAssembly::BLOCK)
       AfterSet.insert(&MI);
 
     // All END_TRY markers placed earlier belong to exceptions that contains
@@ -588,6 +586,8 @@ void WebAssemblyCFGStackify::placeTryMarker(MachineBasicBlock &MBB) {
         AfterSet.insert(&MI);
 #endif
     }
+
+    // It is not possible for an END_BLOCK to be already in this block.
   }
 
   // Mark the end of the TRY.
@@ -774,15 +774,19 @@ void WebAssemblyCFGStackify::placeMarkers(MachineFunction &MF) {
   // Place the LOOP for MBB if MBB is the header of a loop.
   for (auto &MBB : MF)
     placeLoopMarker(MBB);
-  // Place the TRY for MBB if MBB is the EH pad of an exception.
+
   const MCAsmInfo *MCAI = MF.getTarget().getMCAsmInfo();
-  if (MCAI->getExceptionHandlingType() == ExceptionHandling::Wasm &&
-      MF.getFunction().hasPersonalityFn())
-    for (auto &MBB : MF)
-      placeTryMarker(MBB);
-  // Place the BLOCK for MBB if MBB is branched to from above.
-  for (auto &MBB : MF)
-    placeBlockMarker(MBB);
+  for (auto &MBB : MF) {
+    if (MBB.isEHPad()) {
+      // Place the TRY for MBB if MBB is the EH pad of an exception.
+      if (MCAI->getExceptionHandlingType() == ExceptionHandling::Wasm &&
+          MF.getFunction().hasPersonalityFn())
+        placeTryMarker(MBB);
+    } else {
+      // Place the BLOCK for MBB if MBB is branched to from above.
+      placeBlockMarker(MBB);
+    }
+  }
 }
 
 void WebAssemblyCFGStackify::rewriteDepthImmediates(MachineFunction &MF) {

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

@@ -1,4 +1,5 @@
 ; RUN: llc < %s -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -disable-block-placement -verify-machineinstrs -fast-isel=false -machine-sink-split-probability-threshold=0 -cgp-freq-ratio-to-skip-merge=1000 -exception-model=wasm -mattr=+exception-handling | FileCheck %s
+; RUN: llc < %s -O0 -disable-wasm-fallthrough-return-opt -wasm-disable-explicit-locals -wasm-keep-registers -verify-machineinstrs -exception-model=wasm -mattr=+exception-handling | FileCheck %s --check-prefix=NOOPT
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
@@ -70,7 +71,7 @@ rethrow:                                          ; preds = %catch.fallthrough
   call void @llvm.wasm.rethrow.in.catch() [ "funclet"(token %1) ]
   unreachable
 
-try.cont:                                         ; preds = %entry, %catch, %catch2
+try.cont:                                         ; preds = %catch, %catch2, %entry
   ret void
 }
 
@@ -173,15 +174,15 @@ rethrow5:                                         ; preds = %catch.start3
   invoke void @llvm.wasm.rethrow.in.catch() [ "funclet"(token %9) ]
           to label %unreachable unwind label %ehcleanup9
 
-try.cont:                                         ; preds = %catch, %invoke.cont8
+try.cont:                                         ; preds = %invoke.cont8, %catch
   call void @__cxa_end_catch() [ "funclet"(token %1) ]
   catchret from %1 to label %try.cont11
 
 rethrow:                                          ; preds = %catch.start
   call void @llvm.wasm.rethrow.in.catch() [ "funclet"(token %1) ]
   unreachable
 
-try.cont11:                                       ; preds = %entry, %try.cont
+try.cont11:                                       ; preds = %try.cont, %entry
   ret void
 
 ehcleanup:                                        ; preds = %catch6
@@ -286,6 +287,54 @@ terminate:                                        ; preds = %ehcleanup
   unreachable
 }
 
+; Tests if block and try markers are correctly placed. Even if two predecessors
+; of the EH pad are bb2 and bb3 and their nearest common dominator is bb1, the
+; TRY marker should be placed at bb0 because there's a branch from bb0 to bb2,
+; and scopes cannot be interleaved.
+
+; NOOPT-LABEL: test3
+; NOOPT: try
+; NOOPT:   block
+; NOOPT:     block
+; NOOPT:       block
+; NOOPT:       end_block
+; NOOPT:     end_block
+; NOOPT:     call      foo
+; NOOPT:   end_block
+; NOOPT:   call      bar
+; NOOPT: catch     {{.*}}
+; NOOPT: end_try
+define void @test3() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+bb0:
+  br i1 undef, label %bb1, label %bb2
+
+bb1:                                              ; preds = %bb0
+  br i1 undef, label %bb3, label %bb4
+
+bb2:                                              ; preds = %bb0
+  br label %try.cont
+
+bb3:                                              ; preds = %bb1
+  invoke void @foo()
+          to label %try.cont unwind label %catch.dispatch
+
+bb4:                                              ; preds = %bb1
+  invoke void @bar()
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %bb4, %bb3
+  %0 = catchswitch within none [label %catch.start] unwind to caller
+
+catch.start:                                      ; preds = %catch.dispatch
+  %1 = catchpad within %0 [i8* null]
+  %2 = call i8* @llvm.wasm.get.exception(token %1)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  catchret from %1 to label %try.cont
+
+try.cont:                                         ; preds = %catch.start, %bb4, %bb3, %bb2
+  ret void
+}
+
 declare void @foo()
 declare void @bar()
 declare i32 @__gxx_wasm_personality_v0(...)