Prevent dead uses in register coalescer after rematerialization

[rampitec]

Summary

The coalescer does not check if register uses are available
at the point of rematerialization. If it attempts to rematerialize
an instruction with such uses it can end up with use without a def.

LiveRangeEdit does such check during rematerialization, so just
call LiveRangeEdit::allUsesAvailableAt() to avoid the problem.

[arsenm]

LGTM

llvm/include/llvm/CodeGen/LiveRangeEdit.h

private:
 
   /// Rematted - Values that were actually rematted, and so need to have their
   /// live range trimmed or entirely removed.
   SmallPtrSet<const VNInfo *, 4> Rematted;
 
   /// scanRemattable - Identify the Parent values that may rematerialize.
   void scanRemattable(AAResults *aa);
 
-  /// 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;
-
   /// foldAsLoad - If LI has a single use and a single def that can be folded as
   /// a load, eliminate the register by folding the def into the use.
   bool foldAsLoad(LiveInterval *LI, SmallVectorImpl<MachineInstr *> &Dead);
 
   using ToShrinkSet = SetVector<LiveInterval *, SmallVector<LiveInterval *, 8>,
                                 SmallPtrSet<LiveInterval *, 8>>;
 
   /// Helper for eliminateDeadDefs.
public:
   struct Remat {
     VNInfo *ParentVNI;              // parent_'s value at the remat location.
     MachineInstr *OrigMI = nullptr; // Instruction defining OrigVNI. It contains
                                     // the real expr for remat.
 
     explicit Remat(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.
   /// 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

llvm/lib/CodeGen/RegisterCoalescer.cpp

class RegisterCoalescer : public MachineFunctionPass,
 
     /// 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);
+
     /// LiveRangeEdit callback for eliminateDeadDefs().
     void LRE_WillEraseInstruction(MachineInstr *MI) override;
 
     /// Coalesce the LocalWorkList.
     void coalesceLocals();
 
     /// Join compatible live intervals
     void joinAllIntervals();

 }
 
 void RegisterCoalescer::eliminateDeadDefs() {
   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 (!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))
+    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);
   MachineInstr &NewMI = *std::prev(MII);
   NewMI.setDebugLoc(DL);
 

llvm/test/CodeGen/AMDGPU/coalescer-remat-dead-use.mir

+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py
+# RUN: llc -march=amdgcn -mcpu=gfx900 -o - -verify-coalescing -run-pass=simple-register-coalescing %s | FileCheck -check-prefix=GCN %s
+
+---
+# Do not rematerialize V_MOV_B32 at COPY because source register %1 is killed.
+
+name:            no_remat_killed_src_in_inst
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; GCN-LABEL: name: no_remat_killed_src_in_inst
+    ; GCN: liveins: $vgpr0
+    ; GCN: [[COPY:%[0-9]+]]:vgpr_32 = COPY $vgpr0
+    ; GCN: [[V_ADD_U32_e32_:%[0-9]+]]:vgpr_32 = V_ADD_U32_e32 1, [[COPY]], implicit $exec
+    ; GCN: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 [[V_ADD_U32_e32_]], implicit $exec
+    ; GCN: $vgpr0 = COPY [[V_MOV_B32_e32_]]
+    ; GCN: SI_RETURN_TO_EPILOG $vgpr0
+    %0:vgpr_32 = COPY $vgpr0
+    %1:vgpr_32 = V_ADD_U32_e32 1, %0, implicit $exec
+    %2:vgpr_32 = V_MOV_B32_e32 killed %1, implicit $exec
+    $vgpr0 = COPY killed %2
+    SI_RETURN_TO_EPILOG killed $vgpr0
+...
+---
+# Do not rematerialize V_MOV_B32 at COPY because source register %1 is killed
+# after the MOV but before the COPY.
+
+name:            no_remat_killed_src_after_inst
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; GCN-LABEL: name: no_remat_killed_src_after_inst
+    ; GCN: liveins: $vgpr0
+    ; GCN: [[COPY:%[0-9]+]]:vgpr_32 = COPY $vgpr0
+    ; GCN: [[V_ADD_U32_e32_:%[0-9]+]]:vgpr_32 = V_ADD_U32_e32 1, [[COPY]], implicit $exec
+    ; GCN: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 [[V_ADD_U32_e32_]], implicit $exec
+    ; GCN: KILL [[V_ADD_U32_e32_]]
+    ; GCN: $vgpr0 = COPY [[V_MOV_B32_e32_]]
+    ; GCN: SI_RETURN_TO_EPILOG $vgpr0
+    %0:vgpr_32 = COPY $vgpr0
+    %1:vgpr_32 = V_ADD_U32_e32 1, %0, implicit $exec
+    %2:vgpr_32 = V_MOV_B32_e32 %1, implicit $exec
+    KILL %1
+    $vgpr0 = COPY killed %2
+    SI_RETURN_TO_EPILOG killed $vgpr0
+...
+---
+# Even if %1 is not killed do not rematerialize V_MOV_B32 so that we do not
+# extend %1 liverange.
+
+name:            no_remat_alive_src_in_inst_unused
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; GCN-LABEL: name: no_remat_alive_src_in_inst_unused
+    ; GCN: liveins: $vgpr0
+    ; GCN: [[COPY:%[0-9]+]]:vgpr_32 = COPY $vgpr0
+    ; GCN: [[V_ADD_U32_e32_:%[0-9]+]]:vgpr_32 = V_ADD_U32_e32 1, [[COPY]], implicit $exec
+    ; GCN: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 [[V_ADD_U32_e32_]], implicit $exec
+    ; GCN: $vgpr0 = COPY [[V_MOV_B32_e32_]]
+    ; GCN: SI_RETURN_TO_EPILOG $vgpr0
+    %0:vgpr_32 = COPY $vgpr0
+    %1:vgpr_32 = V_ADD_U32_e32 1, %0, implicit $exec
+    %2:vgpr_32 = V_MOV_B32_e32 %1, implicit $exec
+    $vgpr0 = COPY killed %2
+    SI_RETURN_TO_EPILOG killed $vgpr0
+...
+---
+# Rematerialize V_MOV_B32 since %1 is available at COPY and still alive.
+
+name:            remat_alive_src_in_inst_used_and_available
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    liveins: $vgpr0
+
+    ; GCN-LABEL: name: remat_alive_src_in_inst_used_and_available
+    ; GCN: liveins: $vgpr0
+    ; GCN: [[COPY:%[0-9]+]]:vgpr_32 = COPY $vgpr0
+    ; GCN: [[V_ADD_U32_e32_:%[0-9]+]]:vgpr_32 = V_ADD_U32_e32 1, [[COPY]], implicit $exec
+    ; GCN: $vgpr0 = V_MOV_B32_e32 [[V_ADD_U32_e32_]], implicit $exec
+    ; GCN: SI_RETURN_TO_EPILOG $vgpr0
+    %0:vgpr_32 = COPY $vgpr0
+    %1:vgpr_32 = V_ADD_U32_e32 1, %0, implicit $exec
+    %2:vgpr_32 = V_MOV_B32_e32 %1, implicit $exec
+    $vgpr0 = COPY killed %2, implicit %1
+    SI_RETURN_TO_EPILOG killed $vgpr0
+...