[WebAssembly] Remove TEEs when dests are unstackified

[aheejin]

Summary

When created in RegStackify pass, TEE has two destinations, where
op0 is stackified and op1 is not. But it is possible that
op0 becomes unstackified in fixUnwindMismatches function in
CFGStackify pass when a nested try-catch-end is introduced, violating
the invariant of TEEs destinations.

In this case we convert the TEE into two COPYs, which will
eventually be resolved in ExplicitLocals.

[dschuff]

This is indeed what I had in mind on the other CL.
... Do you think it makes sense/is actually better?

[aheejin]

Yeah I switched to this version to address your comment here. And yeah I think this is better that we don't violate assumptions on TEEs and keep ExplicitLocals to its core tasks, while we put all ugly things in fixUnwindMismatches, which is ugly anyway :$... Joking, but I think it's better that fixUnwindMismatches takes care of its own messes.

[aheejin]

• comment fix

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

void WebAssemblyCFGStackify::removeUnnecessaryInstrs(MachineFunction &MF) {
   }
   for (auto *MI : ToDelete) {
     if (MI->getOpcode() == WebAssembly::BLOCK)
       unregisterScope(MI);
     MI->eraseFromParent();
   }
 }
 
+// Get the appropriate copy opcode for the given register class.
+static unsigned getCopyOpcode(const TargetRegisterClass *RC) {
+  if (RC == &WebAssembly::I32RegClass)
+    return WebAssembly::COPY_I32;
+  if (RC == &WebAssembly::I64RegClass)
+    return WebAssembly::COPY_I64;
+  if (RC == &WebAssembly::F32RegClass)
+    return WebAssembly::COPY_F32;
+  if (RC == &WebAssembly::F64RegClass)
+    return WebAssembly::COPY_F64;
+  if (RC == &WebAssembly::V128RegClass)
+    return WebAssembly::COPY_V128;
+  if (RC == &WebAssembly::EXNREFRegClass)
+    return WebAssembly::COPY_EXNREF;
+  llvm_unreachable("Unexpected register class");
+}
+
 // 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) {
+                                         MachineRegisterInfo &MRI,
+                                         const WebAssemblyInstrInfo &TII) {
   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());
     }
   }
+
+  // In RegStackify, when a register definition is used multiple times,
+  //    Reg = INST ...
+  //    INST ..., Reg, ...
+  //    INST ..., Reg, ...
+  //    INST ..., Reg, ...
+  //
+  // we introduce a TEE, which has the following form:
+  //    DefReg = INST ...
+  //    TeeReg, Reg = TEE_... DefReg
+  //    INST ..., TeeReg, ...
+  //    INST ..., Reg, ...
+  //    INST ..., Reg, ...
+  // with DefReg and TeeReg stackified but Reg not stackified.
+  //
+  // But the invariant that TeeReg should be stackified can be violated while we
+  // unstackify registers in the split BB above. In this case, we convert TEEs
+  // into two COPYs. This COPY will be eventually eliminated in ExplicitLocals.
+  //    DefReg = INST ...
+  //    TeeReg = COPY DefReg
+  //    Reg = COPY DefReg
+  //    INST ..., TeeReg, ...
+  //    INST ..., Reg, ...
+  //    INST ..., Reg, ...
+  for (auto I = MBB.begin(), E = MBB.end(); I != E;) {
+    MachineInstr &MI = *I++;
+    if (!WebAssembly::isTee(MI.getOpcode()))
+      continue;
+    Register TeeReg = MI.getOperand(0).getReg();
+    Register Reg = MI.getOperand(1).getReg();
+    Register DefReg = MI.getOperand(2).getReg();
+    if (!MFI.isVRegStackified(TeeReg)) {
+      // Now we are not using TEE anymore, so unstackify DefReg too
+      MFI.unstackifyVReg(DefReg);
+      unsigned CopyOpc = getCopyOpcode(MRI.getRegClass(DefReg));
+      BuildMI(MBB, &MI, MI.getDebugLoc(), TII.get(CopyOpc), TeeReg)
+          .addReg(DefReg);
+      BuildMI(MBB, &MI, MI.getDebugLoc(), TII.get(CopyOpc), Reg).addReg(DefReg);
+      MI.eraseFromParent();
+    }
+  }
 }
 
 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
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);
+    unstackifyVRegsUsedInSplitBB(*EHPad, *BrDest, MFI, MRI, TII);
     // 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);
+      unstackifyVRegsUsedInSplitBB(*MBB, *NestedCont, MFI, MRI, TII);
       registerTryScope(NestedTry, NestedEndTry, NestedEHPad);
 
       // Fix predecessor-successor relationship.
       NestedCont->transferSuccessors(MBB);
       if (EHPad) {
         NestedCont->removeSuccessor(EHPad);
         // If EHPad does not have any predecessors left after removing
         // NextedCont predecessor, remove its successor too, because this EHPad

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

 ; REQUIRES: asserts
 ; 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 -disable-wasm-fallthrough-return-opt -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
 ; 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
 ; 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 -wasm-disable-ehpad-sort | FileCheck %s --check-prefix=NOSORT
+; RUN: llc < %s -disable-wasm-fallthrough-return-opt -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 -wasm-disable-ehpad-sort | FileCheck %s --check-prefix=NOSORT-LOCALS
 ; 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 -wasm-disable-ehpad-sort -stats 2>&1 | FileCheck %s --check-prefix=NOSORT-STAT
 
 target datalayout = "e-m:e-p:32:32-i64:64-n32:64-S128"
 target triple = "wasm32-unknown-unknown"
 
 @_ZTIi = external constant i8*
 @_ZTId = external constant i8*
 
catch.start0:                                     ; preds = %catch.dispatch0
   %3 = call i8* @llvm.wasm.get.exception(token %2)
   %j = call i32 @llvm.wasm.get.ehselector(token %2)
   catchret from %2 to label %try.cont
 
 try.cont:                                         ; preds = %catch.start0
   ret i32 0
 }
 
+; Tests the case when TEE stackifies a register in RegStackify but it gets
+; unstackified in fixUnwindMismatches in CFGStackify.
+
+; NOSORT-LOCALS-LABEL: test8
+define void @test8(i32 %x) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+bb0:
+  invoke void @foo()
+          to label %bb1 unwind label %catch.dispatch0
+
+bb1:                                              ; preds = %bb0
+  %t = add i32 %x, 4
+  ; This %addr is used in multiple places, so tee is introduced in RegStackify,
+  ; which stackifies the use of %addr in store instruction. A tee has two dest
+  ; registers, the first of which is stackified and the second is not.
+  ; But when we introduce a nested try-catch in fixUnwindMismatches in
+  ; CFGStackify, it is possible that we end up unstackifying the first dest
+  ; register. In that case, we convert that tee into a copy.
+  %addr = inttoptr i32 %t to i32*
+  %load = load i32, i32* %addr
+  %call = call i32 @baz()
+  %add = add i32 %load, %call
+  store i32 %add, i32* %addr
+  ret void
+; NOSORT-LOCALS:       i32.add
+; NOSORT-LOCALS-NOT:   local.tee
+; NOSORT-LOCALS-NEXT:  local.set
+
+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)
+  %3 = call i32 @llvm.wasm.get.ehselector(token %1)
+  catchret from %1 to label %try.cont
+
+try.cont:                                         ; preds = %catch.start0
+  ret void
+}
+
 ; If not for the unwind destination mismatch, the LOOP marker here would have an
 ; i32 signature. But because we add a rethrow instruction at the end of the
 ; appendix block, now the LOOP marker does not have a signature (= has a void
 ; signature). Here the two calls two 'bar' are supposed to throw up to the
 ; caller, but incorrectly unwind to 'catch19' after linearizing the CFG.
 
-; NOSORT-LABEL: test8
+; NOSORT-LABEL: test9
 ; NOSORT: block
 ; NOSORT-NOT: loop      i32
-; NOSORT:   loop                                       # label40:
+; NOSORT:   loop                                       # label42:
 ; NOSORT:     try
 ; NOSORT:       call      foo
 ; --- Nested try/catch/end_try starts
 ; NOSORT:       try
 ; NOSORT:         call      bar
 ; NOSORT:         call      bar
 ; NOSORT:       catch     $[[REG:[0-9]+]]=
-; NOSORT:         br        1                          # 1: down to label41
+; NOSORT:         br        1                          # 1: down to label43
 ; NOSORT:       end_try
 ; --- Nested try/catch/end_try ends
 ; NOSORT:       return    {{.*}}
-; NOSORT:     catch     $drop=                         # catch21:
-; NOSORT:       br        1                            # 1: up to label40
-; NOSORT:     end_try                                  # label41:
+; NOSORT:     catch     $drop=                         # catch23:
+; NOSORT:       br        1                            # 1: up to label42
+; NOSORT:     end_try                                  # label43:
 ; NOSORT:   end_loop
 ; NOSORT: end_block
 ; NOSORT: rethrow   $[[REG]]                           # to caller
 
-define i32 @test8(i32* %p) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+define i32 @test9(i32* %p) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   store volatile i32 0, i32* %p
   br label %loop
 
 loop:                                             ; preds = %try.cont, %entry
   store volatile i32 1, i32* %p
   invoke void @foo()
           to label %bb unwind label %catch.dispatch

 }
 
 ; When we have both kinds of EH pad unwind mismatches:
 ; - A may-throw instruction unwinds to an incorrect EH pad after linearizing the
 ;   CFG, when it is supposed to unwind to another EH pad.
 ; - A may-throw instruction unwinds to an incorrect EH pad after linearizing the
 ;   CFG, when it is supposed to unwind to the caller.
 
-; NOSORT-LABEL: test9
+; NOSORT-LABEL: test10
 ; NOSORT: block
 ; NOSORT:   block
 ; NOSORT:     try
 ; NOSORT:       try
 ; NOSORT:         call      foo
 ; --- Nested try/catch/end_try starts
 ; NOSORT:         try
 ; NOSORT:           call      bar
 ; NOSORT:         catch     $[[REG0:[0-9]+]]=
-; NOSORT:           br        2                        # 2: down to label45
+; NOSORT:           br        2                        # 2: down to label47
 ; NOSORT:         end_try
 ; --- Nested try/catch/end_try ends
-; NOSORT:         br        2                          # 2: down to label44
+; NOSORT:         br        2                          # 2: down to label46
 ; NOSORT:       catch     {{.*}}
 ; NOSORT:         block     i32
-; NOSORT:           br_on_exn   0, {{.*}}              # 0: down to label48
+; NOSORT:           br_on_exn   0, {{.*}}              # 0: down to label50
 ; --- Nested try/catch/end_try starts
 ; NOSORT:           try
-; NOSORT:             rethrow   {{.*}}                 # down to catch26
-; NOSORT:           catch     $[[REG1:[0-9]+]]=        # catch26:
-; NOSORT:             br        5                      # 5: down to label43
+; NOSORT:             rethrow   {{.*}}                 # down to catch28
+; NOSORT:           catch     $[[REG1:[0-9]+]]=        # catch28:
+; NOSORT:             br        5                      # 5: down to label45
 ; NOSORT:           end_try
 ; --- Nested try/catch/end_try ends
-; NOSORT:         end_block                            # label48:
+; NOSORT:         end_block                            # label50:
 ; NOSORT:         call      $drop=, __cxa_begin_catch
 ; --- Nested try/catch/end_try starts
 ; NOSORT:         try
 ; NOSORT:           call      __cxa_end_catch
 ; NOSORT:         catch     $[[REG1]]=
-; NOSORT:           br        4                        # 4: down to label43
+; NOSORT:           br        4                        # 4: down to label45
 ; NOSORT:         end_try
 ; --- Nested try/catch/end_try ends
-; NOSORT:         br        2                          # 2: down to label44
+; NOSORT:         br        2                          # 2: down to label46
 ; NOSORT:       end_try
 ; NOSORT:     catch     $[[REG0]]=
-; NOSORT:     end_try                                  # label45:
+; NOSORT:     end_try                                  # label47:
 ; NOSORT:     call      $drop=, __cxa_begin_catch
 ; NOSORT:     call      __cxa_end_catch
-; NOSORT:   end_block                                  # label44:
+; NOSORT:   end_block                                  # label46:
 ; NOSORT:   return
-; NOSORT: end_block                                    # label43:
+; NOSORT: end_block                                    # label45:
 ; NOSORT: rethrow   $[[REG1]]                          # to caller
-define void @test9() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+define void @test10() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 bb0:
   invoke void @foo()
           to label %bb1 unwind label %catch.dispatch0
 
 bb1:                                              ; preds = %bb0
   invoke void @bar()
           to label %try.cont unwind label %catch.dispatch1
 

 ; NOOPT: block
 ; NOOPT:   try
 ; NOOPT:     call      foo
 ; NOOPT:   catch
 ; NOOPT:   end_try
 ; NOOPT:   call      foo
 ; NOOPT: end_block
 ; NOOPT: return
-define void @test10(i32 %arg) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+define void @test11(i32 %arg) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   %tobool = icmp ne i32 %arg, 0
   br i1 %tobool, label %if.then, label %if.end
 
 catch.dispatch:                                   ; preds = %if.then
   %0 = catchswitch within none [label %catch.start] unwind to caller
 
 catch.start:                                      ; preds = %catch.dispatch

 
 ; Intrinsics like memcpy, memmove, and memset don't throw and are lowered into
 ; calls to external symbols (not global addresses) in instruction selection,
 ; which will be eventually lowered to library function calls.
 ; Because this test runs with -wasm-disable-ehpad-sort, these library calls in
 ; invoke.cont BB fall within try~end_try, but they shouldn't cause crashes or
 ; unwinding destination mismatches in CFGStackify.
 
-; NOSORT-LABEL: test11
+; NOSORT-LABEL: test12
 ; NOSORT: try
 ; NOSORT:   call  foo
 ; NOSORT:   call {{.*}} memcpy
 ; NOSORT:   call {{.*}} memmove
 ; NOSORT:   call {{.*}} memset
 ; NOSORT:   return
 ; NOSORT: catch
 ; NOSORT:   rethrow
 ; NOSORT: end_try
-define void @test11(i8* %a, i8* %b) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+define void @test12(i8* %a, i8* %b) personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   %o = alloca %class.Object, align 1
   invoke void @foo()
           to label %invoke.cont unwind label %ehcleanup
 
 invoke.cont:                                      ; preds = %entry
   call void @llvm.memcpy.p0i8.p0i8.i32(i8* %a, i8* %b, i32 100, i1 false)
   call void @llvm.memmove.p0i8.p0i8.i32(i8* %a, i8* %b, i32 100, i1 false)
   call void @llvm.memset.p0i8.i32(i8* %a, i8 0, i32 100, i1 false)
   %call = call %class.Object* @_ZN6ObjectD2Ev(%class.Object* %o) #1
   ret void
 
 ehcleanup:                                        ; preds = %entry
   %0 = cleanuppad within none []
   %call2 = call %class.Object* @_ZN6ObjectD2Ev(%class.Object* %o) #1 [ "funclet"(token %0) ]
   cleanupret from %0 unwind to caller
 }
 
 ; Tests if 'try' marker is placed correctly. In this test, 'try' should be
 ; placed before the call to 'nothrow_i32' and not between the call to
 ; 'nothrow_i32' and 'fun', because the return value of 'nothrow_i32' is
 ; stackified and pushed onto the stack to be consumed by the call to 'fun'.
 
-; CHECK-LABEL: test12
+; CHECK-LABEL: test13
 ; CHECK: try
 ; CHECK: call      $push{{.*}}=, nothrow_i32
 ; CHECK: call      fun, $pop{{.*}}
-define void @test12() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+define void @test13() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   %call = call i32 @nothrow_i32()
   invoke void @fun(i32 %call)
           to label %invoke.cont unwind label %terminate
 
 invoke.cont:                                      ; preds = %entry
   ret void
 
 terminate:                                        ; preds = %entry
   %0 = cleanuppad within none []
   %1 = tail call i8* @llvm.wasm.get.exception(token %0)
   call void @__clang_call_terminate(i8* %1) [ "funclet"(token %0) ]
   unreachable
 }
 
 %class.MyClass = type { i32 }
 
 ; This crashed on debug mode (= when NDEBUG is not defined) when the logic for
 ; computing the innermost region was not correct, in which a loop region
 ; contains an exception region. This should pass CFGSort without crashing.
-define void @test13() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
+define void @test14() personality i8* bitcast (i32 (...)* @__gxx_wasm_personality_v0 to i8*) {
 entry:
   %e = alloca %class.MyClass, align 4
   br label %for.cond
 
 for.cond:                                         ; preds = %for.inc, %entry
   %i.0 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
   %cmp = icmp slt i32 %i.0, 9
   br i1 %cmp, label %for.body, label %for.end
catch.start:                                      ; preds = %catch.dispatch
   call void @__cxa_end_catch() [ "funclet"(token %1) ]
   catchret from %1 to label %try.cont
 
 try.cont:                                         ; preds = %catch.start, %for.end
   ret void
 }
 
 ; Check if the unwind destination mismatch stats are correct
-; NOSORT-STAT: 16 wasm-cfg-stackify    - Number of EH pad unwind mismatches found
+; NOSORT-STAT: 17 wasm-cfg-stackify    - Number of EH pad unwind mismatches found
 
 declare void @foo()
 declare void @bar()
 declare i32 @baz()
 declare i32 @qux(i32)
 declare void @quux(i32)
 declare void @fun(i32)
 ; Function Attrs: nounwind