[WebAssembly] Ensure terminate pads are a single BB

[aheejin]

Summary

This ensures every single terminate pad is a single BB in the form of:

%exn = catch $__cpp_exception
call @__clang_call_terminate(%exn)
unreachable

This is a preparation for HandleEHTerminatePads pass, which will be
added in a later CL and will run after CFGStackify. That pass duplicates
terminate pads with a catch_all instruction, and duplicating it
becomes simpler if we can ensure every terminate pad is a single BB.

[tlively]

llvm/test/CodeGen/WebAssembly/exception.ll

224

Where does this global.set come from? If I understand the code correctly, the call __clang_call_terminate should have been directly after the catch.

238

Is it possible for call __clang_call_terminate to have been duplicated as well so that it could appear in this basic block as well? If so, I don't think the code in this diff would handle deduplicating it.

[aheejin]

llvm/test/CodeGen/WebAssembly/exception.ll

224

It is added in restoreStackPointer in LateEHPrepare. When WebAssemblyFrameLowering::needsPrologForEH returns true, i.e., a function has calls, __stack_pointer has to be saved in the function prolog and restored after each catch.

238

Wouldn't the second __clang_call_terminate within the same EH pad be caught in this code and just skipped?

// If it is already the form we want, skip it
if (Call->getParent() == EHPad &&
    Call->getNextNode()->getOpcode() == WebAssembly::UNREACHABLE)
  continue;

[aheejin]

llvm/test/CodeGen/WebAssembly/exception.ll

238

Sorry, the current code turns out to have a problem. Because eraseDeadBBsAndChildren deletes all child BBs and instructions in there, when we encounter the second __clang_call_terminate call within the terminate pad, it is possible that instruction has been already deleted and its pointer is stale, and accessing it causes a segmentation fault. I fixed the code to only count one __clang_call_terminate per terminate pad. Also added a second __clang_call_terminate call in this BB to test the duplicated calls case you pointed out. PTAL.

[aheejin]

Fix duplicate __clang_call_terminate call case

[tlively]

llvm/test/CodeGen/WebAssembly/exception.ll

224

Makes sense, thanks!

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

 #endif
       break;
     }
     ++InsertPos;
   }
   return InsertPos;
 }
 
-// Find a catch instruction and its destination register within an EH pad.
-static MachineInstr *findCatch(MachineBasicBlock *EHPad, Register &ExnReg) {
-  assert(EHPad->isEHPad());
-  MachineInstr *Catch = nullptr;
-  for (auto &MI : *EHPad) {
-    if (WebAssembly::isCatch(MI.getOpcode())) {
-      Catch = &MI;
-      ExnReg = Catch->getOperand(0).getReg();
-      break;
-    }
-  }
-  assert(Catch && "EH pad does not have a catch");
-  assert(ExnReg != 0 && "Invalid register");
-  return Catch;
-}
-
-static MachineInstr *findCatch(MachineBasicBlock *EHPad) {
-  Register Dummy;
-  return findCatch(EHPad, Dummy);
-}
-
 void WebAssemblyCFGStackify::registerScope(MachineInstr *Begin,
                                            MachineInstr *End) {
   BeginToEnd[Begin] = End;
   EndToBegin[End] = Begin;
 }
 
 void WebAssemblyCFGStackify::registerTryScope(MachineInstr *Begin,
                                               MachineInstr *End,
while (It != MBB->rend()) {
         continue;
       switch (MI.getOpcode()) {
       case WebAssembly::END_TRY: {
         // If a 'try''s return type is fixed, both its try body and catch body
         // should satisfy the return type, so we need to search 'end'
         // instructions before its corresponding 'catch' too.
         auto *EHPad = TryToEHPad.lookup(EndToBegin[&MI]);
         assert(EHPad);
-        Worklist.push_back(std::next(findCatch(EHPad)->getReverseIterator()));
+        auto NextIt =
+            std::next(WebAssembly::findCatch(EHPad)->getReverseIterator());
+        if (NextIt != EHPad->rend())
+          Worklist.push_back(NextIt);
         LLVM_FALLTHROUGH;
       }
       case WebAssembly::END_BLOCK:
       case WebAssembly::END_LOOP:
         EndToBegin[&MI]->getOperand(0).setImm(int32_t(RetType));
         continue;
       default:
         // Something other than an `end`. We're done for this BB.

llvm/lib/Target/WebAssembly/WebAssemblyLateEHPrepare.cpp

class WebAssemblyLateEHPrepare final : public MachineFunctionPass {
 
   bool runOnMachineFunction(MachineFunction &MF) override;
   bool removeUnreachableEHPads(MachineFunction &MF);
   void recordCatchRetBBs(MachineFunction &MF);
   bool hoistCatches(MachineFunction &MF);
   bool addCatchAlls(MachineFunction &MF);
   bool replaceFuncletReturns(MachineFunction &MF);
   bool removeUnnecessaryUnreachables(MachineFunction &MF);
+  bool ensureSingleBBTermPads(MachineFunction &MF);
   bool restoreStackPointer(MachineFunction &MF);
 
   MachineBasicBlock *getMatchingEHPad(MachineInstr *MI);
   SmallSet<MachineBasicBlock *, 8> CatchRetBBs;
 
 public:
   static char ID; // Pass identification, replacement for typeid
   WebAssemblyLateEHPrepare() : MachineFunctionPass(ID) {}
bool WebAssemblyLateEHPrepare::runOnMachineFunction(MachineFunction &MF) {
 
   bool Changed = false;
   if (MF.getFunction().hasPersonalityFn()) {
     Changed |= removeUnreachableEHPads(MF);
     recordCatchRetBBs(MF);
     Changed |= hoistCatches(MF);
     Changed |= addCatchAlls(MF);
     Changed |= replaceFuncletReturns(MF);
+    Changed |= ensureSingleBBTermPads(MF);
   }
   Changed |= removeUnnecessaryUnreachables(MF);
   if (MF.getFunction().hasPersonalityFn())
     Changed |= restoreStackPointer(MF);
   return Changed;
 }
 
 // Remove unreachable EH pads and its children. If they remain, CFG
for (auto &MI : MBB) {
           MBB.removeSuccessor(Succ);
       eraseDeadBBsAndChildren(Succs);
     }
   }
 
   return Changed;
 }
 
+// Clang-generated terminate pads are an single-BB EH pad in the form of
+// termpad:
+//   %exn = catch $__cpp_exception
+//   call @__clang_call_terminate(%exn)
+//   unreachable
+// (There can be local.set and local.gets before the call if we didn't run
+// RegStackify)
+// But code transformations can change or add more control flow, so the call to
+// __clang_call_terminate() function may not be in the original EH pad anymore.
+// This ensures every terminate pad is a single BB in the form illustrated
+// above.
+//
+// This is preparation work for the HandleEHTerminatePads pass later, which
+// duplicates terminate pads both for 'catch' and 'catch_all'. Refer to
+// WebAssemblyHandleEHTerminatePads.cpp for details.
+bool WebAssemblyLateEHPrepare::ensureSingleBBTermPads(MachineFunction &MF) {
+  const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+
+  // Find calls to __clang_call_terminate()
+  SmallVector<MachineInstr *, 8> ClangCallTerminateCalls;
+  for (auto &MBB : MF) {
+    for (auto &MI : MBB) {
+      if (MI.isCall()) {
+        const MachineOperand &CalleeOp = MI.getOperand(0);
+        if (CalleeOp.isGlobal() && CalleeOp.getGlobal()->getName() ==
+                                       WebAssembly::ClangCallTerminateFn)
+          ClangCallTerminateCalls.push_back(&MI);
+      }
+    }
+  }
+
+  bool Changed = false;
+  for (auto *Call : ClangCallTerminateCalls) {
+    MachineBasicBlock *EHPad = getMatchingEHPad(Call);
+    assert(EHPad && "No matching EH pad for catch");
+
+    // If it is already the form we want, skip it
+    if (Call->getParent() == EHPad &&
+        Call->getNextNode()->getOpcode() == WebAssembly::UNREACHABLE)
+      continue;
+
+    // In case the __clang_call_terminate() call is not in its matching EH pad,
+    // move the call to the end of EH pad and add an unreachable instruction
+    // after that. Delete all successors and their children if any, because here
+    // the program terminates.
+    Changed = true;
+    // This runs after hoistCatches(), so catch instruction should be at the top
+    MachineInstr *Catch = WebAssembly::findCatch(EHPad);
+    assert(Catch && "EH pad does not have a catch instruction");
+    // Takes the result register of the catch instruction as argument. There may
+    // have been some other local.set/local.gets in between, but at this point
+    // we don't care.
+    Call->getOperand(1).setReg(Catch->getOperand(0).getReg());
+    auto InsertPos = std::next(MachineBasicBlock::iterator(Catch));
+    EHPad->insert(InsertPos, Call->removeFromParent());
+    BuildMI(*EHPad, InsertPos, Call->getDebugLoc(),
+            TII.get(WebAssembly::UNREACHABLE));
+    EHPad->erase(InsertPos, EHPad->end());
+    SmallVector<MachineBasicBlock *, 8> Succs(EHPad->succ_begin(),
+                                              EHPad->succ_end());
+    for (auto *Succ : Succs)
+      EHPad->removeSuccessor(Succ);
+    eraseDeadBBsAndChildren(Succs);
+  }
+  return Changed;
+}
+
 // After the stack is unwound due to a thrown exception, the __stack_pointer
 // global can point to an invalid address. This inserts instructions that
 // restore __stack_pointer global.
 bool WebAssemblyLateEHPrepare::restoreStackPointer(MachineFunction &MF) {
   const auto *FrameLowering = static_cast<const WebAssemblyFrameLowering *>(
       MF.getSubtarget().getFrameLowering());
   if (!FrameLowering->needsPrologForEH(MF))
     return false;

llvm/lib/Target/WebAssembly/WebAssemblyUtilities.h

 ///
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_LIB_TARGET_WEBASSEMBLY_WEBASSEMBLYUTILITIES_H
 #define LLVM_LIB_TARGET_WEBASSEMBLY_WEBASSEMBLYUTILITIES_H
 
 namespace llvm {
 
+class MachineBasicBlock;
 class MachineInstr;
 class MachineOperand;
 class WebAssemblyFunctionInfo;
 
 namespace WebAssembly {
 
 bool isChild(const MachineInstr &MI, const WebAssemblyFunctionInfo &MFI);
 bool mayThrow(const MachineInstr &MI);
 
 // Exception-related function names
 extern const char *const ClangCallTerminateFn;
 extern const char *const CxaBeginCatchFn;
 extern const char *const CxaRethrowFn;
 extern const char *const StdTerminateFn;
 extern const char *const PersonalityWrapperFn;
 
 /// Returns the operand number of a callee, assuming the argument is a call
 /// instruction.
 const MachineOperand &getCalleeOp(const MachineInstr &MI);
 
+// Find a catch instruction from an EH pad. Returns null if no catch instruction
+// found or the catch is in an invalid location.
+MachineInstr *findCatch(MachineBasicBlock *EHPad);
+
 } // end namespace WebAssembly
 
 } // end namespace llvm
 
 #endif

llvm/lib/Target/WebAssembly/WebAssemblyUtilities.cpp

const MachineOperand &WebAssembly::getCalleeOp(const MachineInstr &MI) {
   case WebAssembly::CALL_INDIRECT_S:
   case WebAssembly::RET_CALL_INDIRECT:
   case WebAssembly::RET_CALL_INDIRECT_S:
     return MI.getOperand(MI.getNumOperands() - 1);
   default:
     llvm_unreachable("Not a call instruction");
   }
 }
+
+// Find a catch instruction from an EH pad.
+MachineInstr *WebAssembly::findCatch(MachineBasicBlock *EHPad) {
+  assert(EHPad->isEHPad());
+  auto Pos = EHPad->begin();
+  // Skip any label or debug instructions. Also skip 'end' marker instructions
+  // that may exist after marker placement in CFGStackify.
+  while (Pos != EHPad->end() &&
+         (Pos->isLabel() || Pos->isDebugInstr() || isMarker(Pos->getOpcode())))
+    Pos++;
+  if (Pos != EHPad->end() && WebAssembly::isCatch(Pos->getOpcode()))
+    return &*Pos;
+  return nullptr;
+}

llvm/test/CodeGen/WebAssembly/exception.ll

 
 terminate:                                        ; preds = %ehcleanup
   %6 = cleanuppad within %5 []
   %7 = call i8* @llvm.wasm.get.exception(token %6)
   call void @__clang_call_terminate(i8* %7) [ "funclet"(token %6) ]
   unreachable
 }
 
+; Tests a case when there are multiple BBs within a terminate pad. This kind of
+; structure is not generated by clang, but can generated by code
+; transformations. After LateEHPrepare, there should be a single 'terminate' BB
+; with these instructions:
+
+; %exn = catch $__cpp_exception
+; call @__clang_call_terminate(%exn)
+; unreachable
+
+; CHECK-LABEL: test_split_terminatepad
+define void @test_split_terminatepad(i1 %arg) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+entry:
+  invoke void @foo()
+          to label %try.cont unwind label %catch.dispatch
+
+catch.dispatch:                                   ; preds = %entry
+  %0 = catchswitch within none [label %catch.start] unwind to caller
+
+; CHECK:      catch
+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)
+  %4 = call i8* @__cxa_begin_catch(i8* %2) [ "funclet"(token %1) ]
+  invoke void @foo() [ "funclet"(token %1) ]
+          to label %invoke.cont1 unwind label %ehcleanup
+
+invoke.cont1:                                     ; preds = %catch.start
+  call void @__cxa_end_catch() [ "funclet"(token %1) ]
+  catchret from %1 to label %try.cont
+
+try.cont:                                         ; preds = %invoke.cont1, %entry
+  ret void
+
+; CHECK:      catch_all
+ehcleanup:                                        ; preds = %catch.start
+  %5 = cleanuppad within %1 []
+  invoke void @__cxa_end_catch() [ "funclet"(token %5) ]
+          to label %invoke.cont2 unwind label %terminate
+
+invoke.cont2:                                     ; preds = %ehcleanup
+  cleanupret from %5 unwind to caller
+
+; This weird structure of split terminate pads are not generated by clang, but
+; we cannot guarantee this kind of multi-BB terminate pads cannot be generated
+; by code transformations. This structure is manually created for this test.
+; CHECK:      catch     $[[EXN:[0-9]+]]=, __cpp_exception
+; CHECK-NEXT: global.set  __stack_pointer
+; CHECK-NEXT: call  __clang_call_terminate, $[[EXN]]
+; CHECK-NEXT: unreachable
+
+terminate:                                        ; preds = %ehcleanup
+  %6 = cleanuppad within %5 []
+  br i1 %arg, label %terminate.split1, label %terminate.split2
+
+terminate.split1:
+  %7 = call i8* @llvm.wasm.get.exception(token %6)
+  call void @__clang_call_terminate(i8* %7) [ "funclet"(token %6) ]
+  unreachable
+
+terminate.split2:
+  unreachable
+}
+
 ; Tests prologues and epilogues are not generated within EH scopes.
 ; They should not be treated as funclets; BBs starting with a catch instruction
 ; should not have a prologue, and BBs ending with a catchret/cleanupret should
 ; not have an epilogue. This is separate from __stack_pointer restoring
 ; instructions after a catch instruction.
 ;
 ; void bar(int) noexcept;
 ; void test_no_prolog_epilog_in_ehpad() {