[WebAssembly] Rematerialize constants rather than hold them live in registers.

Teach the register stackifier to rematerialize constants that have multiple
uses instead of leaving them in registers. In the WebAssembly encoding, it's
the same code size to materialize most constants as it is to read a value
from a register.

llvm-svn: 258142

4 files changed, 77 insertions, 40 deletions

diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.cpp
index 028e9af0834..37dd0c4026a 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.cpp
@@ -32,6 +32,21 @@ WebAssemblyInstrInfo::WebAssemblyInstrInfo(const WebAssemblySubtarget &STI)
                               WebAssembly::ADJCALLSTACKUP),
       RI(STI.getTargetTriple()) {}
 
+bool WebAssemblyInstrInfo::isReallyTriviallyReMaterializable(
+    const MachineInstr *MI, AliasAnalysis *AA) const {
+  switch (MI->getOpcode()) {
+  case WebAssembly::CONST_I32:
+  case WebAssembly::CONST_I64:
+  case WebAssembly::CONST_F32:
+  case WebAssembly::CONST_F64:
+    // isReallyTriviallyReMaterializableGeneric misses these because of the
+    // ARGUMENTS implicit def, so we manualy override it here.
+    return true;
+  default:
+    return false;
+  }
+}
+
 void WebAssemblyInstrInfo::copyPhysReg(MachineBasicBlock &MBB,
                                        MachineBasicBlock::iterator I,
                                        DebugLoc DL, unsigned DestReg,
@@ -39,9 +54,10 @@ void WebAssemblyInstrInfo::copyPhysReg(MachineBasicBlock &MBB,
   // This method is called by post-RA expansion, which expects only pregs to
   // exist. However we need to handle both here.
   auto &MRI = MBB.getParent()->getRegInfo();
-  const TargetRegisterClass *RC = TargetRegisterInfo::isVirtualRegister(DestReg) ?
-      MRI.getRegClass(DestReg) :
-      MRI.getTargetRegisterInfo()->getMinimalPhysRegClass(SrcReg);
+  const TargetRegisterClass *RC =
+      TargetRegisterInfo::isVirtualRegister(DestReg)
+          ? MRI.getRegClass(DestReg)
+          : MRI.getTargetRegisterInfo()->getMinimalPhysRegClass(SrcReg);
 
   unsigned CopyLocalOpcode;
   if (RC == &WebAssembly::I32RegClass)
@@ -145,9 +161,7 @@ unsigned WebAssemblyInstrInfo::InsertBranch(MachineBasicBlock &MBB,
   assert(Cond.size() == 2 && "Expected a flag and a successor block");
 
   if (Cond[0].getImm()) {
-    BuildMI(&MBB, DL, get(WebAssembly::BR_IF))
-        .addOperand(Cond[1])
-        .addMBB(TBB);
+    BuildMI(&MBB, DL, get(WebAssembly::BR_IF)).addOperand(Cond[1]).addMBB(TBB);
   } else {
     BuildMI(&MBB, DL, get(WebAssembly::BR_UNLESS))
         .addOperand(Cond[1])
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.h b/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.h
index 5ddd9b36f24..bf0c1dc6c95 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.h
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.h
@@ -34,6 +34,9 @@ public:
 
   const WebAssemblyRegisterInfo &getRegisterInfo() const { return RI; }
 
+  bool isReallyTriviallyReMaterializable(const MachineInstr *MI,
+                                         AliasAnalysis *AA) const override;
+
   void copyPhysReg(MachineBasicBlock &MBB, MachineBasicBlock::iterator MI,
                    DebugLoc DL, unsigned DestReg, unsigned SrcReg,
                    bool KillSrc) const override;
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.td b/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.td
index 2e682a47547..75b76565c3f 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.td
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyInstrInfo.td
@@ -123,7 +123,7 @@ defm : LOCAL<I64>;
 defm : LOCAL<F32>;
 defm : LOCAL<F64>;
 
-let isMoveImm = 1 in {
+let isMoveImm = 1, isAsCheapAsAMove = 1, isReMaterializable = 1 in {
 def CONST_I32 : I<(outs I32:$res), (ins i32imm:$imm),
                   [(set I32:$res, imm:$imm)],
                   "i32.const\t$res, $imm">;
@@ -136,7 +136,7 @@ def CONST_F32 : I<(outs F32:$res), (ins f32imm_op:$imm),
 def CONST_F64 : I<(outs F64:$res), (ins f64imm_op:$imm),
                   [(set F64:$res, fpimm:$imm)],
                   "f64.const\t$res, $imm">;
-} // isMoveImm = 1
+} // isMoveImm = 1, isAsCheapAsAMove = 1, isReMaterializable = 1
 
 } // Defs = [ARGUMENTS]
 
diff --git a/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp b/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp
index 537c147e614..5923edd656d 100644
--- a/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp
+++ b/llvm/lib/Target/WebAssembly/WebAssemblyRegStackify.cpp
@@ -23,6 +23,7 @@
 #include "WebAssembly.h"
 #include "MCTargetDesc/WebAssemblyMCTargetDesc.h" // for WebAssembly::ARGUMENT_*
 #include "WebAssemblyMachineFunctionInfo.h"
+#include "WebAssemblySubtarget.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/CodeGen/LiveIntervalAnalysis.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
@@ -116,9 +117,9 @@ static bool IsSafeToMove(const MachineInstr *Def, const MachineInstr *Insert,
     // Ask LiveIntervals whether moving this virtual register use or def to
     // Insert will change value numbers are seen.
     const LiveInterval &LI = LIS.getInterval(Reg);
-    VNInfo *DefVNI = MO.isDef() ?
-        LI.getVNInfoAt(LIS.getInstructionIndex(Def).getRegSlot()) :
-        LI.getVNInfoBefore(LIS.getInstructionIndex(Def));
+    VNInfo *DefVNI =
+        MO.isDef() ? 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)
@@ -140,6 +141,8 @@ bool WebAssemblyRegStackify::runOnMachineFunction(MachineFunction &MF) {
   bool Changed = false;
   MachineRegisterInfo &MRI = MF.getRegInfo();
   WebAssemblyFunctionInfo &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
+  const auto *TII = MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
+  const auto *TRI = MF.getSubtarget<WebAssemblySubtarget>().getRegisterInfo();
   AliasAnalysis &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
   LiveIntervals &LIS = getAnalysis<LiveIntervals>();
 
@@ -176,10 +179,6 @@ bool WebAssemblyRegStackify::runOnMachineFunction(MachineFunction &MF) {
         if (!Def)
           continue;
 
-        // There's no use in nesting implicit defs inside anything.
-        if (Def->getOpcode() == TargetOpcode::IMPLICIT_DEF)
-          continue;
-
         // Don't nest an INLINE_ASM def into anything, because we don't have
         // constraints for $pop outputs.
         if (Def->getOpcode() == TargetOpcode::INLINEASM)
@@ -197,30 +196,51 @@ bool WebAssemblyRegStackify::runOnMachineFunction(MachineFunction &MF) {
             Def->getOpcode() == WebAssembly::ARGUMENT_F64)
           continue;
 
-        // Single-use expression trees require defs that have one use.
-        // TODO: Eventually we'll relax this, to take advantage of set_local
-        // returning its result.
-        if (!MRI.hasOneUse(Reg))
-          continue;
-
-        // For now, be conservative and don't look across block boundaries.
-        // TODO: Be more aggressive?
-        if (Def->getParent() != &MBB)
-          continue;
-
-        // Don't move instructions that have side effects or memory dependencies
-        // or other complications.
-        if (!IsSafeToMove(Def, Insert, AA, LIS, MRI))
-          continue;
-
-        Changed = true;
-        AnyStackified = true;
-        // Move the def down and nest it in the current instruction.
-        MBB.splice(Insert, &MBB, Def);
-        LIS.handleMove(Def);
-        MFI.stackifyVReg(Reg);
-        ImposeStackOrdering(Def);
-        Insert = Def;
+        if (MRI.hasOneUse(Reg) && Def->getParent() == &MBB &&
+            IsSafeToMove(Def, Insert, AA, LIS, MRI)) {
+          // A single-use def in the same block with no intervening memory or
+          // register dependencies; move the def down and nest it with the
+          // current instruction.
+          // TODO: Stackify multiple-use values, taking advantage of set_local
+          // returning its result.
+          Changed = true;
+          AnyStackified = true;
+          MBB.splice(Insert, &MBB, Def);
+          LIS.handleMove(Def);
+          MFI.stackifyVReg(Reg);
+          ImposeStackOrdering(Def);
+          Insert = Def;
+        } else if (Def->isAsCheapAsAMove() &&
+                   TII->isTriviallyReMaterializable(Def, &AA)) {
+          // A trivially cloneable instruction; clone it and nest the new copy
+          // with the current instruction.
+          Changed = true;
+          AnyStackified = true;
+          unsigned OldReg = Def->getOperand(0).getReg();
+          unsigned NewReg = MRI.createVirtualRegister(MRI.getRegClass(OldReg));
+          TII->reMaterialize(MBB, Insert, NewReg, 0, Def, *TRI);
+          Op.setReg(NewReg);
+          MachineInstr *Clone =
+              &*std::prev(MachineBasicBlock::instr_iterator(Insert));
+          LIS.InsertMachineInstrInMaps(Clone);
+          LIS.createAndComputeVirtRegInterval(NewReg);
+          MFI.stackifyVReg(NewReg);
+          ImposeStackOrdering(Clone);
+          Insert = Clone;
+
+          // If that was the last use of the original, delete the original.
+          // Otherwise shrink the LiveInterval.
+          if (MRI.use_empty(OldReg)) {
+            SlotIndex Idx = LIS.getInstructionIndex(Def).getRegSlot();
+            LIS.removePhysRegDefAt(WebAssembly::ARGUMENTS, Idx);
+            LIS.removeVRegDefAt(LIS.getInterval(OldReg), Idx);
+            LIS.removeInterval(OldReg);
+            LIS.RemoveMachineInstrFromMaps(Def);
+            Def->eraseFromParent();
+          } else {
+            LIS.shrinkToUses(&LIS.getInterval(OldReg));
+          }
+        }
       }
       if (AnyStackified)
         ImposeStackOrdering(&*MII);
@@ -236,7 +256,7 @@ bool WebAssemblyRegStackify::runOnMachineFunction(MachineFunction &MF) {
   }
 
 #ifndef NDEBUG
-  // Verify that pushes and pops are performed in FIFO order.
+  // Verify that pushes and pops are performed in LIFO order.
   SmallVector<unsigned, 0> Stack;
   for (MachineBasicBlock &MBB : MF) {
     for (MachineInstr &MI : MBB) {