[RegAlloc] Account for use availability when applying rematerializable weight discount

llvm/include/llvm/CodeGen/CalcSpillWeights.h

@@ -81,6 +81,14 @@ class VirtRegMap;
     static bool isRematerializable(const LiveInterval &LI,
                                    const LiveIntervals &LIS,
                                    const VirtRegMap &VRM,
+                                   const MachineRegisterInfo &MRI,
+                                   const TargetInstrInfo &TII);
+
+    /// \returns true if all registers used by \p MI are also available with the
+    /// same value at \p UseIdx.
+    static bool allUsesAvailableAt(const MachineInstr *MI, SlotIndex UseIdx,
+                                   const LiveIntervals &LIS,
+                                   const MachineRegisterInfo &MRI,
                                    const TargetInstrInfo &TII);
 
   protected:

llvm/include/llvm/CodeGen/LiveRangeEdit.h

@@ -189,11 +189,6 @@ class LiveRangeEdit : private MachineRegisterInfo::Delegate {
     explicit Remat(const VNInfo *ParentVNI) : ParentVNI(ParentVNI) {}
   };
 
-  /// allUsesAvailableAt - Return true if all registers used by OrigMI at
-  /// OrigIdx are also available with the same value at UseIdx.
-  bool allUsesAvailableAt(const MachineInstr *OrigMI, SlotIndex OrigIdx,
-                          SlotIndex UseIdx) const;
-
   /// canRematerializeAt - Determine if ParentVNI can be rematerialized at
   /// UseIdx. It is assumed that parent_.getVNINfoAt(UseIdx) == ParentVNI.
   bool canRematerializeAt(Remat &RM, VNInfo *OrigVNI, SlotIndex UseIdx);

llvm/lib/CodeGen/CalcSpillWeights.cpp

@@ -81,12 +81,15 @@ Register VirtRegAuxInfo::copyHint(const MachineInstr *MI, Register Reg,
 bool VirtRegAuxInfo::isRematerializable(const LiveInterval &LI,
                                         const LiveIntervals &LIS,
                                         const VirtRegMap &VRM,
+                                        const MachineRegisterInfo &MRI,
                                         const TargetInstrInfo &TII) {
   Register Reg = LI.reg();
   Register Original = VRM.getOriginal(Reg);
+  SmallDenseMap<unsigned, MachineInstr *> VNIDefs;
   for (LiveInterval::const_vni_iterator I = LI.vni_begin(), E = LI.vni_end();
        I != E; ++I) {
     const VNInfo *VNI = *I;
+    const VNInfo *OrigVNI = VNI;
     if (VNI->isUnused())
       continue;
     if (VNI->isPHIDef())
@@ -124,6 +127,75 @@ bool VirtRegAuxInfo::isRematerializable(const LiveInterval &LI,
 
     if (!TII.isTriviallyReMaterializable(*MI))
       return false;
+
+    VNIDefs[OrigVNI->id] = MI;
+  }
+
+  // If MI has register uses, it will only be rematerializable if its uses are
+  // also live at the indices it will be rematerialized at.
+  for (MachineOperand &MO : MRI.reg_nodbg_operands(LI.reg())) {
+    if (MO.isUndef() || !MO.readsReg())
+      continue;
+    SlotIndex UseIdx = LIS.getInstructionIndex(*MO.getParent());
+    MachineInstr *Def = VNIDefs[LI.getVNInfoAt(UseIdx)->id];
+    assert(Def && "Use with no def");
+    if (!allUsesAvailableAt(Def, UseIdx, LIS, MRI, TII))
+      return false;
+  }
+
+  return true;
+}
+
+bool VirtRegAuxInfo::allUsesAvailableAt(const MachineInstr *MI,
+                                        SlotIndex UseIdx,
+                                        const LiveIntervals &LIS,
+                                        const MachineRegisterInfo &MRI,
+                                        const TargetInstrInfo &TII) {
+  SlotIndex OrigIdx = LIS.getInstructionIndex(*MI).getRegSlot(true);
+  UseIdx = std::max(UseIdx, UseIdx.getRegSlot(true));
+  for (const MachineOperand &MO : MI->operands()) {
+    if (!MO.isReg() || !MO.getReg() || !MO.readsReg())
+      continue;
+
+    // We can't remat physreg uses, unless it is a constant or target wants
+    // to ignore this use.
+    if (MO.getReg().isPhysical()) {
+      if (MRI.isConstantPhysReg(MO.getReg()) || TII.isIgnorableUse(MO))
+        continue;
+      return false;
+    }
+
+    const LiveInterval &li = LIS.getInterval(MO.getReg());
+    const VNInfo *OVNI = li.getVNInfoAt(OrigIdx);
+    if (!OVNI)
+      continue;
+
+    // Don't allow rematerialization immediately after the original def.
+    // It would be incorrect if OrigMI redefines the register.
+    // See PR14098.
+    if (SlotIndex::isSameInstr(OrigIdx, UseIdx))
+      return false;
+
+    if (OVNI != li.getVNInfoAt(UseIdx))
+      return false;
+
+    // Check that subrange is live at UseIdx.
+    if (li.hasSubRanges()) {
+      const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
+      unsigned SubReg = MO.getSubReg();
+      LaneBitmask LM = SubReg ? TRI->getSubRegIndexLaneMask(SubReg)
+                              : MRI.getMaxLaneMaskForVReg(MO.getReg());
+      for (const LiveInterval::SubRange &SR : li.subranges()) {
+        if ((SR.LaneMask & LM).none())
+          continue;
+        if (!SR.liveAt(UseIdx))
+          return false;
+        // Early exit if all used lanes are checked. No need to continue.
+        LM &= ~SR.LaneMask;
+        if (LM.none())
+          break;
+      }
+    }
   }
   return true;
 }
@@ -339,7 +411,7 @@ float VirtRegAuxInfo::weightCalcHelper(LiveInterval &LI, SlotIndex *Start,
   // it is a preferred candidate for spilling.
   // FIXME: this gets much more complicated once we support non-trivial
   // re-materialization.
-  if (isRematerializable(LI, LIS, VRM, *MF.getSubtarget().getInstrInfo()))
+  if (isRematerializable(LI, LIS, VRM, MRI, *MF.getSubtarget().getInstrInfo()))
     TotalWeight *= 0.5F;
 
   // Finally, we scale the weight by the scale factor of register class.

llvm/lib/CodeGen/LiveRangeEdit.cpp

@@ -92,60 +92,6 @@ bool LiveRangeEdit::anyRematerializable() {
   return !Remattable.empty();
 }
 
-/// allUsesAvailableAt - Return true if all registers used by OrigMI at
-/// OrigIdx are also available with the same value at UseIdx.
-bool LiveRangeEdit::allUsesAvailableAt(const MachineInstr *OrigMI,
-                                       SlotIndex OrigIdx,
-                                       SlotIndex UseIdx) const {
-  OrigIdx = OrigIdx.getRegSlot(true);
-  UseIdx = std::max(UseIdx, UseIdx.getRegSlot(true));
-  for (const MachineOperand &MO : OrigMI->operands()) {
-    if (!MO.isReg() || !MO.getReg() || !MO.readsReg())
-      continue;
-
-    // We can't remat physreg uses, unless it is a constant or target wants
-    // to ignore this use.
-    if (MO.getReg().isPhysical()) {
-      if (MRI.isConstantPhysReg(MO.getReg()) || TII.isIgnorableUse(MO))
-        continue;
-      return false;
-    }
-
-    LiveInterval &li = LIS.getInterval(MO.getReg());
-    const VNInfo *OVNI = li.getVNInfoAt(OrigIdx);
-    if (!OVNI)
-      continue;
-
-    // Don't allow rematerialization immediately after the original def.
-    // It would be incorrect if OrigMI redefines the register.
-    // See PR14098.
-    if (SlotIndex::isSameInstr(OrigIdx, UseIdx))
-      return false;
-
-    if (OVNI != li.getVNInfoAt(UseIdx))
-      return false;
-
-    // Check that subrange is live at UseIdx.
-    if (li.hasSubRanges()) {
-      const TargetRegisterInfo *TRI = MRI.getTargetRegisterInfo();
-      unsigned SubReg = MO.getSubReg();
-      LaneBitmask LM = SubReg ? TRI->getSubRegIndexLaneMask(SubReg)
-                              : MRI.getMaxLaneMaskForVReg(MO.getReg());
-      for (LiveInterval::SubRange &SR : li.subranges()) {
-        if ((SR.LaneMask & LM).none())
-          continue;
-        if (!SR.liveAt(UseIdx))
-          return false;
-        // Early exit if all used lanes are checked. No need to continue.
-        LM &= ~SR.LaneMask;
-        if (LM.none())
-          break;
-      }
-    }
-  }
-  return true;
-}
-
 bool LiveRangeEdit::canRematerializeAt(Remat &RM, VNInfo *OrigVNI,
                                        SlotIndex UseIdx) {
   assert(ScannedRemattable && "Call anyRematerializable first");
@@ -155,12 +101,10 @@ bool LiveRangeEdit::canRematerializeAt(Remat &RM, VNInfo *OrigVNI,
     return false;
 
   // No defining instruction provided.
-  SlotIndex DefIdx;
   assert(RM.OrigMI && "No defining instruction for remattable value");
-  DefIdx = LIS.getInstructionIndex(*RM.OrigMI);
 
   // Verify that all used registers are available with the same values.
-  if (!allUsesAvailableAt(RM.OrigMI, DefIdx, UseIdx))
+  if (!VirtRegAuxInfo::allUsesAvailableAt(RM.OrigMI, UseIdx, LIS, MRI, TII))
     return false;
 
   return true;
@@ -221,8 +165,8 @@ bool LiveRangeEdit::foldAsLoad(LiveInterval *LI,
 
   // Since we're moving the DefMI load, make sure we're not extending any live
   // ranges.
-  if (!allUsesAvailableAt(DefMI, LIS.getInstructionIndex(*DefMI),
-                          LIS.getInstructionIndex(*UseMI)))
+  if (!VirtRegAuxInfo::allUsesAvailableAt(
+          DefMI, LIS.getInstructionIndex(*UseMI), LIS, MRI, TII))
     return false;
 
   // We also need to make sure it is safe to move the load.

llvm/lib/CodeGen/MLRegAllocEvictAdvisor.cpp

@@ -927,7 +927,7 @@ MLEvictAdvisor::getLIFeatureComponents(const LiveInterval &LI) const {
       Ret.HintWeights += Freq;
   }
   Ret.IsRemat = VirtRegAuxInfo::isRematerializable(
-      LI, *LIS, *VRM, *MF.getSubtarget().getInstrInfo());
+      LI, *LIS, *VRM, *MRI, *MF.getSubtarget().getInstrInfo());
   return Ret;
 }
 

llvm/lib/CodeGen/RegisterCoalescer.cpp

@@ -20,6 +20,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/CodeGen/CalcSpillWeights.h"
 #include "llvm/CodeGen/LiveInterval.h"
 #include "llvm/CodeGen/LiveIntervals.h"
 #include "llvm/CodeGen/LiveRangeEdit.h"
@@ -1393,10 +1394,7 @@ bool RegisterCoalescer::reMaterializeTrivialDef(const CoalescerPair &CP,
     }
   }
 
-  SmallVector<Register, 8> NewRegs;
-  LiveRangeEdit Edit(&SrcInt, NewRegs, *MF, *LIS, nullptr, this);
-  SlotIndex DefIdx = LIS->getInstructionIndex(*DefMI);
-  if (!Edit.allUsesAvailableAt(DefMI, DefIdx, CopyIdx))
+  if (!VirtRegAuxInfo::allUsesAvailableAt(DefMI, CopyIdx, *LIS, *MRI, *TII))
     return false;
 
   DebugLoc DL = CopyMI->getDebugLoc();
@@ -1405,6 +1403,8 @@ bool RegisterCoalescer::reMaterializeTrivialDef(const CoalescerPair &CP,
       std::next(MachineBasicBlock::iterator(CopyMI));
   LiveRangeEdit::Remat RM(ValNo);
   RM.OrigMI = DefMI;
+  SmallVector<Register, 8> NewRegs;
+  LiveRangeEdit Edit(&SrcInt, NewRegs, *MF, *LIS, nullptr, this);
   Edit.rematerializeAt(*MBB, MII, DstReg, RM, *TRI, false, SrcIdx, CopyMI);
   MachineInstr &NewMI = *std::prev(MII);
   NewMI.setDebugLoc(DL);