[RegisterCoalescer] Use LiveRangeEdit to handle rematerialization

[Carrot]

Summary

This patch uses the API provided by LiveRangeEdit to handle rematerialization. It will make future maintenance and improvement more easier.

No functional change.

[rampitec]

llvm/lib/CodeGen/RegisterCoalescer.cpp

1307

isAsCheapAsAMove is a very cheap check, better to leave it here and bail early rather then defer it until you call canRematerializeAt.

[rampitec]

LGTM, but let others review it too.

[arsenm]

llvm/lib/CodeGen/LiveRangeEdit.cpp

197

Shouldn't this return getRegSlot like the fallthrough case?

[arsenm]

LGTM

llvm/include/llvm/CodeGen/LiveRangeEdit.h

public:
   /// canRematerializeAt - Determine if ParentVNI can be rematerialized at
   /// UseIdx. It is assumed that parent_.getVNINfoAt(UseIdx) == ParentVNI.
   /// When cheapAsAMove is set, only cheap remats are allowed.
   bool canRematerializeAt(Remat &RM, VNInfo *OrigVNI, SlotIndex UseIdx,
                           bool cheapAsAMove);
 
   /// rematerializeAt - Rematerialize RM.ParentVNI into DestReg by inserting an
   /// instruction into MBB before MI. The new instruction is mapped, but
-  /// liveness is not updated.
+  /// liveness is not updated. If ReplaceIndexMI is not null it will be replaced
+  /// by new MI in the index map.
   /// Return the SlotIndex of the new instruction.
   SlotIndex rematerializeAt(MachineBasicBlock &MBB,
                             MachineBasicBlock::iterator MI, unsigned DestReg,
                             const Remat &RM, const TargetRegisterInfo &,
-                            bool Late = false);
+                            bool Late = false, unsigned SubIdx = 0,
+                            MachineInstr *ReplaceIndexMI = nullptr);
 
   /// markRematerialized - explicitly mark a value as rematerialized after doing
   /// it manually.
   void markRematerialized(const VNInfo *ParentVNI) {
     Rematted.insert(ParentVNI);
   }
 
   /// didRematerialize - Return true if ParentVNI was rematerialized anywhere.

llvm/lib/CodeGen/LiveRangeEdit.cpp

bool LiveRangeEdit::canRematerializeAt(Remat &RM, VNInfo *OrigVNI,
   return true;
 }
 
 SlotIndex LiveRangeEdit::rematerializeAt(MachineBasicBlock &MBB,
                                          MachineBasicBlock::iterator MI,
                                          unsigned DestReg,
                                          const Remat &RM,
                                          const TargetRegisterInfo &tri,
-                                         bool Late) {
+                                         bool Late,
+                                         unsigned SubIdx,
+                                         MachineInstr *ReplaceIndexMI) {
   assert(RM.OrigMI && "Invalid remat");
-  TII.reMaterialize(MBB, MI, DestReg, 0, *RM.OrigMI, tri);
+  TII.reMaterialize(MBB, MI, DestReg, SubIdx, *RM.OrigMI, tri);
   // DestReg of the cloned instruction cannot be Dead. Set isDead of DestReg
   // to false anyway in case the isDead flag of RM.OrigMI's dest register
   // is true.
   (*--MI).getOperand(0).setIsDead(false);
   Rematted.insert(RM.ParentVNI);
+  if (ReplaceIndexMI)
+    return LIS.ReplaceMachineInstrInMaps(*ReplaceIndexMI, *MI);
   return LIS.getSlotIndexes()->insertMachineInstrInMaps(*MI, Late).getRegSlot();
 }
 
 void LiveRangeEdit::eraseVirtReg(Register Reg) {
   if (TheDelegate && TheDelegate->LRE_CanEraseVirtReg(Reg))
     LIS.removeInterval(Reg);
 }
 

llvm/lib/CodeGen/RegisterCoalescer.cpp

class RegisterCoalescer : public MachineFunctionPass,
     /// lateLiveIntervalUpdate is called.
     DenseSet<Register> ToBeUpdated;
 
     /// Record how many times the large live interval with many valnos
     /// has been tried to join with other live interval.
     DenseMap<Register, unsigned long> LargeLIVisitCounter;
 
     /// Recursively eliminate dead defs in DeadDefs.
-    void eliminateDeadDefs();
-
-    /// 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);
+    void eliminateDeadDefs(LiveRangeEdit *Edit = nullptr);
 
     /// LiveRangeEdit callback for eliminateDeadDefs().
     void LRE_WillEraseInstruction(MachineInstr *MI) override;
 
     /// Coalesce the LocalWorkList.
     void coalesceLocals();
 
     /// Join compatible live intervals
void RegisterCoalescer::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addPreserved<LiveIntervals>();
   AU.addPreserved<SlotIndexes>();
   AU.addRequired<MachineLoopInfo>();
   AU.addPreserved<MachineLoopInfo>();
   AU.addPreservedID(MachineDominatorsID);
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
-void RegisterCoalescer::eliminateDeadDefs() {
+void RegisterCoalescer::eliminateDeadDefs(LiveRangeEdit *Edit) {
+  if (Edit) {
+    Edit->eliminateDeadDefs(DeadDefs);
+    return;
+  }
   SmallVector<Register, 8> NewRegs;
   LiveRangeEdit(nullptr, NewRegs, *MF, *LIS,
                 nullptr, this).eliminateDeadDefs(DeadDefs);
 }
 
-bool RegisterCoalescer::allUsesAvailableAt(const MachineInstr *OrigMI,
-                                           SlotIndex OrigIdx,
-                                           SlotIndex UseIdx) {
-  SmallVector<Register, 8> NewRegs;
-  return LiveRangeEdit(nullptr, NewRegs, *MF, *LIS, nullptr, this)
-      .allUsesAvailableAt(OrigMI, OrigIdx, UseIdx);
-}
-
 void RegisterCoalescer::LRE_WillEraseInstruction(MachineInstr *MI) {
   // MI may be in WorkList. Make sure we don't visit it.
   ErasedInstrs.insert(MI);
 }
 
 bool RegisterCoalescer::adjustCopiesBackFrom(const CoalescerPair &CP,
                                              MachineInstr *CopyMI) {
   assert(!CP.isPartial() && "This doesn't work for partial copies.");
if (ValNo->isPHIDef() || ValNo->isUnused())
     return false;
   MachineInstr *DefMI = LIS->getInstructionFromIndex(ValNo->def);
   if (!DefMI)
     return false;
   if (DefMI->isCopyLike()) {
     IsDefCopy = true;
     return false;
   }
   if (!TII->isAsCheapAsAMove(*DefMI))
     return false;
-  if (!TII->isTriviallyReMaterializable(*DefMI))
+
+  SmallVector<Register, 8> NewRegs;
+  LiveRangeEdit Edit(&SrcInt, NewRegs, *MF, *LIS, nullptr, this);
+  if (!Edit.checkRematerializable(ValNo, DefMI))
     return false;
+
   if (!definesFullReg(*DefMI, SrcReg))
     return false;
   bool SawStore = false;
   if (!DefMI->isSafeToMove(AA, SawStore))
     return false;
   const MCInstrDesc &MCID = DefMI->getDesc();
   if (MCID.getNumDefs() != 1)
     return false;
if (!DefMI->isImplicitDef()) {
     } else {
       // Theoretically, some stack frame reference could exist. Just make sure
       // it hasn't actually happened.
       assert(Register::isVirtualRegister(DstReg) &&
              "Only expect to deal with virtual or physical registers");
     }
   }
 
-  if (!allUsesAvailableAt(DefMI, ValNo->def, CopyIdx))
+  LiveRangeEdit::Remat RM(ValNo);
+  RM.OrigMI = DefMI;
+  if (!Edit.canRematerializeAt(RM, ValNo, CopyIdx, true))
     return false;
 
   DebugLoc DL = CopyMI->getDebugLoc();
   MachineBasicBlock *MBB = CopyMI->getParent();
   MachineBasicBlock::iterator MII =
     std::next(MachineBasicBlock::iterator(CopyMI));
-  TII->reMaterialize(*MBB, MII, DstReg, SrcIdx, *DefMI, *TRI);
+  Edit.rematerializeAt(*MBB, MII, DstReg, RM, *TRI, false, SrcIdx, CopyMI);
   MachineInstr &NewMI = *std::prev(MII);
   NewMI.setDebugLoc(DL);
 
   // In a situation like the following:
   //     %0:subreg = instr              ; DefMI, subreg = DstIdx
   //     %1        = copy %0:subreg ; CopyMI, SrcIdx = 0
   // instead of widening %1 to the register class of %0 simply do:
   //     %1 = instr
for (unsigned I = CopyMI->getDesc().getNumOperands(),
     if (MO.isReg()) {
       assert(MO.isImplicit() && "No explicit operands after implicit operands.");
       // Discard VReg implicit defs.
       if (Register::isPhysicalRegister(MO.getReg()))
         ImplicitOps.push_back(MO);
     }
   }
 
-  LIS->ReplaceMachineInstrInMaps(*CopyMI, NewMI);
   CopyMI->eraseFromParent();
   ErasedInstrs.insert(CopyMI);
 
   // NewMI may have dead implicit defs (E.g. EFLAGS for MOV<bits>r0 on X86).
   // We need to remember these so we can add intervals once we insert
   // NewMI into SlotIndexes.
   SmallVector<MCRegister, 4> NewMIImplDefs;
   for (unsigned i = NewMI.getDesc().getNumOperands(),
bool RegisterCoalescer::reMaterializeTrivialDef(const CoalescerPair &CP,
   for (MachineOperand &UseMO : MRI->use_nodbg_operands(SrcReg)) {
     if (UseMO.getParent()->isCopyLike())
       NumCopyUses++;
   }
   if (NumCopyUses < LateRematUpdateThreshold) {
     // The source interval can become smaller because we removed a use.
     shrinkToUses(&SrcInt, &DeadDefs);
     if (!DeadDefs.empty())
-      eliminateDeadDefs();
+      eliminateDeadDefs(&Edit);
   } else {
     ToBeUpdated.insert(SrcReg);
   }
   return true;
 }
 
 MachineInstr *RegisterCoalescer::eliminateUndefCopy(MachineInstr *CopyMI) {
   // ProcessImplicitDefs may leave some copies of <undef> values, it only