[WebAssembly] Stackify induction variable increment instructions.

This handles instructions where the defined register is also used, as in
"x = x + 1".

llvm-svn: 269830

3 files changed, 109 insertions, 42 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp
index f64c2d0941d..b47de787999 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp
@@ -211,7 +211,9 @@ static bool ShouldRematerialize(const MachineInstr *Def, AliasAnalysis &AA,
          TII->isTriviallyReMaterializable(Def, &AA);
 }
 
-/// Identify the definition for this register at this point.
+// Identify the definition for this register at this point. This is a
+// generalization of MachineRegisterInfo::getUniqueVRegDef that uses
+// LiveIntervals to handle complex cases.
 static MachineInstr *GetVRegDef(unsigned Reg, const MachineInstr *Insert,
                                 const MachineRegisterInfo &MRI,
                                 const LiveIntervals &LIS)
@@ -228,6 +230,34 @@ static MachineInstr *GetVRegDef(unsigned Reg, const MachineInstr *Insert,
   return nullptr;
 }
 
+// Test whether Reg, as defined at Def, has exactly one use. This is a
+// generalization of MachineRegisterInfo::hasOneUse that uses LiveIntervals
+// to handle complex cases.
+static bool HasOneUse(unsigned Reg, MachineInstr *Def,
+                      MachineRegisterInfo &MRI, MachineDominatorTree &MDT,
+                      LiveIntervals &LIS) {
+  // Most registers are in SSA form here so we try a quick MRI query first.
+  if (MRI.hasOneUse(Reg))
+    return true;
+
+  bool HasOne = false;
+  const LiveInterval &LI = LIS.getInterval(Reg);
+  const VNInfo *DefVNI = LI.getVNInfoAt(
+      LIS.getInstructionIndex(*Def).getRegSlot());
+  assert(DefVNI);
+  for (auto I : MRI.use_operands(Reg)) {
+    const auto &Result = LI.Query(LIS.getInstructionIndex(*I.getParent()));
+    if (Result.valueIn() == DefVNI) {
+      if (!Result.isKill())
+        return false;
+      if (HasOne)
+        return false;
+      HasOne = true;
+    }
+  }
+  return HasOne;
+}
+
 // Test whether it's safe to move Def to just before Insert.
 // TODO: Compute memory dependencies in a way that doesn't require always
 // walking the block.
@@ -263,10 +293,15 @@ static bool IsSafeToMove(const MachineInstr *Def, const MachineInstr *Insert,
 
     // Ask LiveIntervals whether moving this virtual register use or def to
     // Insert will change which value numbers are seen.
+    // 
+    // If the operand is a use of a register that is also defined in the same
+    // instruction, test that the newly defined value reaches the insert point,
+    // since the operand will be moving along with the def.
     const LiveInterval &LI = LIS.getInterval(Reg);
     VNInfo *DefVNI =
-        MO.isDef() ? LI.getVNInfoAt(LIS.getInstructionIndex(*Def).getRegSlot())
-                   : LI.getVNInfoBefore(LIS.getInstructionIndex(*Def));
+        (MO.isDef() || Def->definesRegister(Reg)) ?
+        LI.getVNInfoAt(LIS.getInstructionIndex(*Def).getRegSlot()) :
+        LI.getVNInfoBefore(LIS.getInstructionIndex(*Def));
     assert(DefVNI && "Instruction input missing value number");
     VNInfo *InsVNI = LI.getVNInfoBefore(LIS.getInstructionIndex(*Insert));
     if (InsVNI && DefVNI != InsVNI)
@@ -321,14 +356,7 @@ static bool OneUseDominatesOtherUses(unsigned Reg, const MachineOperand &OneUse,
       continue;
 
     const MachineInstr *OneUseInst = OneUse.getParent();
-    if (UseInst->getOpcode() == TargetOpcode::PHI) {
-      // Test that the PHI use, which happens on the CFG edge rather than
-      // within the PHI's own block, is dominated by the one selected use.
-      const MachineBasicBlock *Pred =
-          UseInst->getOperand(&Use - &UseInst->getOperand(0) + 1).getMBB();
-      if (!MDT.dominates(&MBB, Pred))
-        return false;
-    } else if (UseInst == OneUseInst) {
+    if (UseInst == OneUseInst) {
       // Another use in the same instruction. We need to ensure that the one
       // selected use happens "before" it.
       if (&OneUse > &Use)
@@ -376,9 +404,13 @@ static MachineInstr *MoveForSingleUse(unsigned Reg, MachineOperand& Op,
   MBB.splice(Insert, &MBB, Def);
   LIS.handleMove(*Def);
 
-  if (MRI.hasOneDef(Reg)) {
+  if (MRI.hasOneDef(Reg) && MRI.hasOneUse(Reg)) {
+    // No one else is using this register for anything so we can just stackify
+    // it in place.
     MFI.stackifyVReg(Reg);
   } else {
+    // The register may have unrelated uses or defs; create a new register for
+    // just our one def and use so that we can stackify it.
     unsigned NewReg = MRI.createVirtualRegister(MRI.getRegClass(Reg));
     Def->getOperand(0).setReg(NewReg);
     Op.setReg(NewReg);
@@ -456,7 +488,7 @@ RematerializeCheapDef(unsigned Reg, MachineOperand &Op, MachineInstr *Def,
 /// to this:
 ///
 ///    DefReg = INST ...     // Def (to become the new Insert)
-///    TeeReg, NewReg = TEE_LOCAL_... DefReg
+///    TeeReg, Reg = TEE_LOCAL_... DefReg
 ///    INST ..., TeeReg, ... // Insert
 ///    INST ..., NewReg, ...
 ///    INST ..., NewReg, ...
@@ -469,29 +501,42 @@ static MachineInstr *MoveAndTeeForMultiUse(
     MachineRegisterInfo &MRI, const WebAssemblyInstrInfo *TII) {
   DEBUG(dbgs() << "Move and tee for multi-use:"; Def->dump());
 
+  // Move Def into place.
   MBB.splice(Insert, &MBB, Def);
   LIS.handleMove(*Def);
+
+  // Create the Tee and attach the registers.
   const auto *RegClass = MRI.getRegClass(Reg);
-  unsigned NewReg = MRI.createVirtualRegister(RegClass);
   unsigned TeeReg = MRI.createVirtualRegister(RegClass);
   unsigned DefReg = MRI.createVirtualRegister(RegClass);
   MachineOperand &DefMO = Def->getOperand(0);
-  MRI.replaceRegWith(Reg, NewReg);
   MachineInstr *Tee = BuildMI(MBB, Insert, Insert->getDebugLoc(),
                               TII->get(GetTeeLocalOpcode(RegClass)), TeeReg)
-                          .addReg(NewReg, RegState::Define)
+                          .addReg(Reg, RegState::Define)
                           .addReg(DefReg, getUndefRegState(DefMO.isDead()));
   Op.setReg(TeeReg);
   DefMO.setReg(DefReg);
-  LIS.InsertMachineInstrInMaps(*Tee);
-  LIS.removeInterval(Reg);
-  LIS.createAndComputeVirtRegInterval(NewReg);
+  SlotIndex TeeIdx = LIS.InsertMachineInstrInMaps(*Tee).getRegSlot();
+  SlotIndex DefIdx = LIS.getInstructionIndex(*Def).getRegSlot();
+
+  // Tell LiveIntervals we moved the original vreg def from Def to Tee.
+  LiveInterval &LI = LIS.getInterval(Reg);
+  LiveInterval::iterator I = LI.FindSegmentContaining(DefIdx);
+  VNInfo *ValNo = LI.getVNInfoAt(DefIdx);
+  I->start = TeeIdx;
+  ValNo->def = TeeIdx;
+  ShrinkToUses(LI, LIS);
+
+  // Finish stackifying the new regs.
   LIS.createAndComputeVirtRegInterval(TeeReg);
   LIS.createAndComputeVirtRegInterval(DefReg);
   MFI.stackifyVReg(DefReg);
   MFI.stackifyVReg(TeeReg);
   ImposeStackOrdering(Def);
   ImposeStackOrdering(Tee);
+
+  DEBUG(dbgs() << " - Replaced register: "; Def->dump());
+  DEBUG(dbgs() << " - Tee instruction: "; Tee->dump());
   return Def;
 }
 
@@ -636,10 +681,6 @@ bool WebAssemblyRegStackify::runOnMachineFunction(MachineFunction &MF) {
     // iterating over it and the end iterator may change.
     for (auto MII = MBB.rbegin(); MII != MBB.rend(); ++MII) {
       MachineInstr *Insert = &*MII;
-      // Don't nest anything inside a phi.
-      if (Insert->getOpcode() == TargetOpcode::PHI)
-        break;
-
       // Don't nest anything inside an inline asm, because we don't have
       // constraints for $push inputs.
       if (Insert->getOpcode() == TargetOpcode::INLINEASM)
@@ -678,10 +719,6 @@ bool WebAssemblyRegStackify::runOnMachineFunction(MachineFunction &MF) {
         if (Def->getOpcode() == TargetOpcode::INLINEASM)
           continue;
 
-        // Don't nest PHIs inside of anything.
-        if (Def->getOpcode() == TargetOpcode::PHI)
-          continue;
-
         // Argument instructions represent live-in registers and not real
         // instructions.
         if (Def->getOpcode() == WebAssembly::ARGUMENT_I32 ||
@@ -699,7 +736,7 @@ bool WebAssemblyRegStackify::runOnMachineFunction(MachineFunction &MF) {
         bool SameBlock = Def->getParent() == &MBB;
         bool CanMove = SameBlock && IsSafeToMove(Def, Insert, AA, LIS, MRI) &&
                        !TreeWalker.IsOnStack(Reg);
-        if (CanMove && MRI.hasOneUse(Reg)) {
+        if (CanMove && HasOneUse(Reg, Def, MRI, MDT, LIS)) {
           Insert = MoveForSingleUse(Reg, Op, Def, MBB, Insert, LIS, MFI, MRI);
         } else if (ShouldRematerialize(Def, AA, TII)) {
           Insert = RematerializeCheapDef(Reg, Op, Def, MBB, Insert, LIS, MFI,
diff --git a/llvm/test/CodeGen/WebAssembly/cfg-stackify.ll b/llvm/test/CodeGen/WebAssembly/cfg-stackify.ll
index ff945706778..cbb45611cec 100644
--- a/llvm/test/CodeGen/WebAssembly/cfg-stackify.ll
+++ b/llvm/test/CodeGen/WebAssembly/cfg-stackify.ll
@@ -15,10 +15,10 @@ declare void @something()
 
 ; CHECK-LABEL: test0:
 ; CHECK: loop
-; CHECK-NOT: br
-; CHECK: i32.add
 ; CHECK-NEXT: block
-; CHECK-NEXT: i32.lt_s
+; CHECK-NEXT: i32.const
+; CHECK-NEXT: i32.add
+; CHECK:      i32.lt_s
 ; CHECK-NEXT: br_if
 ; CHECK-NEXT: return
 ; CHECK-NEXT: .LBB0_3:
@@ -29,9 +29,9 @@ declare void @something()
 ; CHECK-NEXT: end_loop
 ; OPT-LABEL: test0:
 ; OPT: loop
-; OPT-NOT: br
-; OPT: i32.add
-; OPT-NEXT: i32.ge_s
+; OPT-NEXT: i32.const
+; OPT-NEXT: i32.add
+; OPT:      i32.ge_s
 ; OPT-NEXT: br_if
 ; OPT-NOT: br
 ; OPT: call
@@ -60,10 +60,10 @@ back:
 
 ; CHECK-LABEL: test1:
 ; CHECK: loop
-; CHECK-NOT: br
-; CHECK: i32.add
 ; CHECK-NEXT: block
-; CHECK-NEXT: i32.lt_s
+; CHECK-NEXT: i32.const
+; CHECK-NEXT: i32.add
+; CHECK:      i32.lt_s
 ; CHECK-NEXT: br_if
 ; CHECK-NEXT: return
 ; CHECK-NEXT: .LBB1_3:
@@ -74,9 +74,9 @@ back:
 ; CHECK-NEXT: end_loop
 ; OPT-LABEL: test1:
 ; OPT: loop
-; OPT-NOT: br
-; OPT: i32.add
-; OPT-NEXT: i32.ge_s
+; OPT-NEXT: i32.const
+; OPT-NEXT: i32.add
+; OPT:      i32.ge_s
 ; OPT-NEXT: br_if
 ; OPT-NOT: br
 ; OPT: call
@@ -108,16 +108,22 @@ back:
 ; CHECK: block{{$}}
 ; CHECK: br_if 0, {{[^,]+}}{{$}}
 ; CHECK: .LBB2_{{[0-9]+}}:
-; CHECK: br_if 0, ${{[0-9]+}}{{$}}
+; CHECK: loop
+; CHECK: br_if 0, $pop{{[0-9]+}}{{$}}
 ; CHECK: .LBB2_{{[0-9]+}}:
+; CHECK: end_loop
+; CHECK: end_block
 ; CHECK: return{{$}}
 ; OPT-LABEL: test2:
 ; OPT-NOT: local
 ; OPT: block{{$}}
 ; OPT: br_if 0, {{[^,]+}}{{$}}
 ; OPT: .LBB2_{{[0-9]+}}:
-; OPT: br_if 0, ${{[0-9]+}}{{$}}
+; OPT: loop
+; OPT: br_if 0, $pop{{[0-9]+}}{{$}}
 ; OPT: .LBB2_{{[0-9]+}}:
+; OPT: end_loop
+; OPT: end_block
 ; OPT: return{{$}}
 define void @test2(double* nocapture %p, i32 %n) {
 entry:
diff --git a/llvm/test/CodeGen/WebAssembly/reg-stackify.ll b/llvm/test/CodeGen/WebAssembly/reg-stackify.ll
index 8518f4c1f49..bc36357c0b9 100644
--- a/llvm/test/CodeGen/WebAssembly/reg-stackify.ll
+++ b/llvm/test/CodeGen/WebAssembly/reg-stackify.ll
@@ -400,6 +400,30 @@ define i32 @no_stackify_past_epilogue() {
   ret i32 %call
 }
 
+; Stackify a loop induction variable into a loop comparison.
+
+; CHECK_LABEL: stackify_indvar:
+; CHECK:             i32.const   $push[[L5:.+]]=, 1{{$}}
+; CHECK-NEXT:        i32.add     $push[[L4:.+]]=, $[[R0:.+]], $pop[[L5]]{{$}}
+; CHECK-NEXT:        tee_local   $push[[L3:.+]]=, $[[R0]]=, $pop[[L4]]{{$}}
+; CHECK-NEXT:        i32.ne      $push[[L2:.+]]=, $0, $pop[[L3]]{{$}}
+define void @stackify_indvar(i32 %tmp, i32* %v) #0 {
+bb:
+  br label %bb3
+
+bb3:                                              ; preds = %bb3, %bb2
+  %tmp4 = phi i32 [ %tmp7, %bb3 ], [ 0, %bb ]
+  %tmp5 = load volatile i32, i32* %v, align 4
+  %tmp6 = add nsw i32 %tmp5, %tmp4
+  store volatile i32 %tmp6, i32* %v, align 4
+  %tmp7 = add nuw nsw i32 %tmp4, 1
+  %tmp8 = icmp eq i32 %tmp7, %tmp
+  br i1 %tmp8, label %bb10, label %bb3
+
+bb10:                                             ; preds = %bb9, %bb
+  ret void
+}
+
 !llvm.module.flags = !{!0}
 !llvm.dbg.cu = !{!1}