[WebAssembly] Unstackify regs after fixing unwinding mismatches

[aheejin]

Summary

Fixing unwind mismatches for exception handling can result in splicing
existing BBs and moving some of instructions to new BBs. In this case
some of stackified def registers in the original BB can be used in the
split BB. For example, we have this BB and suppose %r0 is a stackified
register.

bb.1:
  %r0 = call @foo
  ... use %r0 ...

After fixing unwind mismatches in CFGStackify, bb.1 can be split and
some instructions can be moved to a newly created BB:

bb.1:
  %r0 = call @foo

bb.split (new):
  ... use %r0 ...

In this case we should make %r0 un-stackified, because its use is now in
another BB.

When spliting a BB, this CL unstackifies all def registers that have
uses in the new split BB.

[dschuff]

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

699

It might be more convenient/concise to just pass in MRI/MFI from the caller. (Hopefully the resulting code would be the same, but I guess in an unoptimized build, maybe faster too).

[aheejin]

• Address comments

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

void WebAssemblyCFGStackify::removeUnnecessaryInstrs(MachineFunction &MF) {
   }
   for (auto *MI : ToDelete) {
     if (MI->getOpcode() == WebAssembly::BLOCK)
       unregisterScope(MI);
     MI->eraseFromParent();
   }
 }
 
+// When MBB is split into MBB and Split, we should unstackify defs in MBB that
+// have their uses in Split.
+static void unstackifyVRegsUsedInSplitBB(MachineBasicBlock &MBB,
+                                         MachineBasicBlock &Split,
+                                         WebAssemblyFunctionInfo &MFI,
+                                         MachineRegisterInfo &MRI) {
+  for (auto &MI : Split) {
+    for (auto &MO : MI.explicit_uses()) {
+      if (!MO.isReg() || Register::isPhysicalRegister(MO.getReg()))
+        continue;
+      if (MachineInstr *Def = MRI.getUniqueVRegDef(MO.getReg()))
+        if (Def->getParent() == &MBB)
+          MFI.unstackifyVReg(MO.getReg());
+    }
+  }
+}
+
 bool WebAssemblyCFGStackify::fixUnwindMismatches(MachineFunction &MF) {
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+  auto &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
   MachineRegisterInfo &MRI = MF.getRegInfo();
 
   // Linearizing the control flow by placing TRY / END_TRY markers can create
   // mismatches in unwind destinations. There are two kinds of mismatches we
   // try to solve here.
 
   // 1. When an instruction may throw, but the EH pad it will unwind to can be
   //    different from the original CFG.
for (auto &P : UnwindDestToTryRanges) {
     MachineInstr *EndTry = BeginToEnd[BeginTry];
     MachineBasicBlock *Cont = EndTry->getParent();
     auto *BrDest = MF.CreateMachineBasicBlock();
     MF.insert(std::next(EHPad->getIterator()), BrDest);
     // Hoist up the existing 'end_try'.
     BrDest->insert(BrDest->end(), EndTry->removeFromParent());
     // Take out the handler body from EH pad to the new branch destination BB.
     BrDest->splice(BrDest->end(), EHPad, SplitPos, EHPad->end());
+    unstackifyVRegsUsedInSplitBB(*EHPad, *BrDest, MFI, MRI);
     // Fix predecessor-successor relationship.
     BrDest->transferSuccessors(EHPad);
     EHPad->addSuccessor(BrDest);
 
     // All try ranges that were supposed to unwind to this EH pad now have to
     // branch to this new branch dest BB.
     for (auto Range : UnwindDestToTryRanges[EHPad])
       BrDestToTryRanges[BrDest].push_back(Range);
for (auto Range : TryRanges) {
       MF.insert(std::next(NestedEHPad->getIterator()), NestedCont);
       MachineInstr *NestedEndTry =
           BuildMI(*NestedCont, NestedCont->begin(), RangeEnd->getDebugLoc(),
                   TII.get(WebAssembly::END_TRY));
       // In case MBB has more instructions after the try range, move them to the
       // new nested continuation BB.
       NestedCont->splice(NestedCont->end(), MBB,
                          std::next(RangeEnd->getIterator()), MBB->end());
+      unstackifyVRegsUsedInSplitBB(*MBB, *NestedCont, MFI, MRI);
       registerTryScope(NestedTry, NestedEndTry, NestedEHPad);
 
       // Fix predecessor-successor relationship.
       NestedCont->transferSuccessors(MBB);
       if (EHPad)
         NestedCont->removeSuccessor(EHPad);
       MBB->addSuccessor(NestedEHPad);
       MBB->addSuccessor(NestedCont);

llvm/lib/Target/WebAssembly/WebAssemblyMachineFunctionInfo.h

public:
 
   void stackifyVReg(unsigned VReg) {
     assert(MF.getRegInfo().getUniqueVRegDef(VReg));
     auto I = Register::virtReg2Index(VReg);
     if (I >= VRegStackified.size())
       VRegStackified.resize(I + 1);
     VRegStackified.set(I);
   }
+  void unstackifyVReg(unsigned VReg) {
+    auto I = Register::virtReg2Index(VReg);
+    if (I < VRegStackified.size())
+      VRegStackified.reset(I);
+  }
   bool isVRegStackified(unsigned VReg) const {
     auto I = Register::virtReg2Index(VReg);
     if (I >= VRegStackified.size())
       return false;
     return VRegStackified.test(I);
   }
 
   void initWARegs();

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

 }
 
 ; Two 'call bar''s original unwind destination was the caller, but after control
 ; flow linearization, their unwind destination incorrectly becomes 'catch17'. We
 ; fix this by wrapping the call with a nested try/catch/end_try and branching to
 ; the right destination (label4), from which we rethrow the exception to the
 ; caller.
 
+; And the return value of 'baz' should NOT be stackified because the BB is split
+; during fixing unwind mismatches.
+
 ; NOSORT-LABEL: test6
 ; NOSORT:   try
 ; NOSORT:     call      foo
 ; --- Nested try/catch/end_try starts
 ; NOSORT:     try
 ; NOSORT:       call      bar
-; NOSORT:       call      bar
+; NOSORT:       i32.call  ${{[0-9]+}}=, baz
+; NOSORT-NOT:   i32.call  $push{{.*}}=, baz
 ; NOSORT:     catch     $[[REG:[0-9]+]]=
 ; NOSORT:       br        1                            # 1: down to label35
 ; NOSORT:     end_try
 ; --- Nested try/catch/end_try ends
 ; NOSORT:     return
 ; NOSORT:   catch     $drop=                           # catch17:
 ; NOSORT:     return
 ; NOSORT:   end_try                                    # label35:
 ; NOSORT:   rethrow   $[[REG]]                         # to caller
 
 define void @test6() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 bb0:
   invoke void @foo()
           to label %bb1 unwind label %catch.dispatch0
 
 bb1:                                              ; preds = %bb0
   call void @bar()
-  call void @bar()
+  %call = call i32 @baz()
+  call void @nothrow(i32 %call) #0
   ret void
 
 catch.dispatch0:                                  ; preds = %bb0
   %0 = catchswitch within none [label %catch.start0] unwind to caller
 
 catch.start0:                                     ; preds = %catch.dispatch0
   %1 = catchpad within %0 [i8* null]
   %2 = call i8* @llvm.wasm.get.exception(token %1)
catch.start1:                                     ; preds = %catch.dispatch1
   catchret from %6 to label %try.cont
 
 try.cont:                                         ; preds = %catch.start1, %catch.start0, %bb1
   ret void
 }
 
 declare void @foo()
 declare void @bar()
+declare i32 @baz()
+; Function Attrs: nounwind
+declare void @nothrow(i32) #0
 declare i32 @__gxx_wasm_personality_v0(...)
 declare i8* @llvm.wasm.get.exception(token)
 declare i32 @llvm.wasm.get.ehselector(token)
 declare void @llvm.wasm.rethrow.in.catch()
 declare i32 @llvm.eh.typeid.for(i8*)
 declare i8* @__cxa_begin_catch(i8*)
 declare void @__cxa_end_catch()
 declare void @__clang_call_terminate(i8*)
 declare void @_ZSt9terminatev()
+
+attributes #0 = { nounwind }