[WIP][CodeGen] Encode liveness for COPY instructions after virtRegRewriter pass.

[vg0204]

As for the Greedy RA, the virtRegRewriter pass is the last place that holds livenes info, even at subregister level. So, now that information can be extracted and encoded on COPY instruction. This is in reference to SWDEV-498533

This information for COPY can later be used to identify partially live regsiters precisely, assuming the liveness information used is not invalidated by any kind if IR muatation later[#151124 ]. It enables us to stop using implict & implicit-def operands while expandinf the copies later in COpyPhysReg.

NOTE: Currently, as we do not support any MachineOperand or MachineOperand attribute( associated with Register operand) that could hold liveness information as the 64-bit laneBitmask, so just for sake of EXPERIMENTATION, encoded this info as an immediate 64-bit unsigned integer.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

@llvm/pr-subscribers-llvm-regalloc

Author: Vikash Gupta (vg0204)

Changes

As for the Greedy RA, the virtRegRewriter pass is the last place that holds livenes info, even at subregister level. So, now that information can be extracted and encoded on COPY instruction.

This information for COPY can later be used to identify partially live regsiters precisely, assuming the liveness information used is not invalidated by any kind if IR muatation later. It enables us to stop using implict & implicit-def operands while expandinf the copies later in COpyPhysReg.

NOTE: Currently, as we do not support any MachineOperand or MachineOperand attribute( associated with Register operand) that could hold liveness information as the 64-bit laneBitmask, so just for sake of EXPERIMENTATION, encoded this info as an immediate 64-bit unsigned integer.

---

Full diff: https://github.com/llvm/llvm-project/pull/151123.diff

2 Files Affected:

• (modified) llvm/include/llvm/Target/Target.td (+1-1)
• (modified) llvm/lib/CodeGen/VirtRegMap.cpp (+88-1)

diff --git a/llvm/include/llvm/Target/Target.td b/llvm/include/llvm/Target/Target.td
index 4c83f8a580aa0..1f125c2cf87de 100644
--- a/llvm/include/llvm/Target/Target.td
+++ b/llvm/include/llvm/Target/Target.td
@@ -1323,7 +1323,7 @@ def REG_SEQUENCE : StandardPseudoInstruction {
 }
 def COPY : StandardPseudoInstruction {
   let OutOperandList = (outs unknown:$dst);
-  let InOperandList = (ins unknown:$src);
+  let InOperandList = (ins unknown:$src, variable_ops);
   let AsmString = "";
   let hasSideEffects = false;
   let isAsCheapAsAMove = true;
diff --git a/llvm/lib/CodeGen/VirtRegMap.cpp b/llvm/lib/CodeGen/VirtRegMap.cpp
index 99ba893d6f096..227c0ae813934 100644
--- a/llvm/lib/CodeGen/VirtRegMap.cpp
+++ b/llvm/lib/CodeGen/VirtRegMap.cpp
@@ -213,6 +213,8 @@ class VirtRegRewriter {
   void rewrite();
   void addMBBLiveIns();
   bool readsUndefSubreg(const MachineOperand &MO) const;
+  uint64_t calcLiveRegUnitMask(const MachineOperand &MO,
+                               MCRegister PhysReg) const;
   void addLiveInsForSubRanges(const LiveInterval &LI, MCRegister PhysReg) const;
   void handleIdentityCopy(MachineInstr &MI);
   void expandCopyBundle(MachineInstr &MI) const;
@@ -474,6 +476,77 @@ bool VirtRegRewriter::readsUndefSubreg(const MachineOperand &MO) const {
   return true;
 }
 
+// Return LaneBitmask value as unint64_t for PhysReg assigned to MO,
+// representing its live register units at its parent MI. In case of undef or
+// fully live MO, return 0u.
+uint64_t VirtRegRewriter::calcLiveRegUnitMask(const MachineOperand &MO,
+                                              MCRegister PhysReg) const {
+  Register Reg = MO.getReg();
+  const LiveInterval &LI = LIS->getInterval(Reg);
+  const MachineInstr &MI = *MO.getParent();
+  SlotIndex MIIndex = LIS->getInstructionIndex(MI);
+  unsigned SubRegIdx = MO.getSubReg();
+  LaneBitmask UseMask = SubRegIdx
+                            ? TRI->getSubRegIndexLaneMask(SubRegIdx)
+                            : (Reg.isVirtual() ? MRI->getMaxLaneMaskForVReg(Reg)
+                                               : LaneBitmask::getNone());
+
+  LaneBitmask LiveRegUnitMask;
+  DenseSet<unsigned> LiveRegUnits;
+
+  // dbgs() << "\n********** " << printReg(Reg, TRI) << "[ " <<
+  // printReg(PhysReg, TRI) << " ]" << " **********\n";
+
+  if (MO.isUndef())
+    return 0u;
+
+  assert(LI.liveAt(MIIndex) &&
+         "Reads of completely dead register should be marked undef already");
+
+  if (LI.hasSubRanges()) {
+    for (MCRegUnitMaskIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
+      unsigned Unit = (*Units).first;
+      LaneBitmask Mask = (*Units).second;
+      for (const LiveInterval::SubRange &S : LI.subranges()) {
+        if ((S.LaneMask & UseMask & Mask).any() && S.liveAt(MIIndex)) {
+          LiveRegUnits.insert(Unit);
+        }
+      }
+    }
+  } else {
+    for (MCRegUnitMaskIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
+      unsigned Unit = (*Units).first;
+      const LiveRange &UnitRange = LIS->getRegUnit(Unit);
+      LaneBitmask Mask = (*Units).second;
+
+      if (UnitRange.liveAt(MIIndex) && (UseMask & Mask).any())
+        LiveRegUnits.insert(Unit);
+    }
+  }
+
+  // Consider the exact subregister & create new UseMask as per the RC for it.
+  if (SubRegIdx != 0) {
+    PhysReg = TRI->getSubReg(PhysReg, SubRegIdx);
+    UseMask = (TRI->getMinimalPhysRegClass(PhysReg))->getLaneMask();
+  }
+
+  for (MCRegUnitMaskIterator Units(PhysReg, TRI); Units.isValid(); ++Units) {
+    unsigned Unit = (*Units).first;
+    LaneBitmask Mask = (*Units).second;
+    if (LiveRegUnits.count(Unit)) {
+      // dbgs() << "LIVE DEF UNIT : " << printRegUnit(Unit, TRI) << '\n';
+      LiveRegUnitMask |= Mask;
+    }
+  }
+
+  // dbgs() << "UseMask : " << PrintLaneMask(UseMask) << '\n';
+  // dbgs() << "LiveRegUnitMask : " << PrintLaneMask(LiveRegUnitMask) << '\n';
+  if (UseMask == LiveRegUnitMask)
+    return 0u;
+
+  return LiveRegUnitMask.getAsInteger();
+}
+
 void VirtRegRewriter::handleIdentityCopy(MachineInstr &MI) {
   if (!MI.isIdentityCopy())
     return;
@@ -495,7 +568,11 @@ void VirtRegRewriter::handleIdentityCopy(MachineInstr &MI) {
   // give us additional liveness information: The target (super-)register
   // must not be valid before this point. Replace the COPY with a KILL
   // instruction to maintain this information.
-  if (MI.getOperand(1).isUndef() || MI.getNumOperands() > 2) {
+
+  // Avoid COPY with an exact 3 operand, wiith third operand be Mask, as
+  // it same as a COPY with no additional liveness information.
+  if (MI.getOperand(1).isUndef() || MI.getNumOperands() > 3 ||
+      (MI.getNumOperands() == 3 && !MI.getOperand(2).isImm())) {
     MI.setDesc(TII->get(TargetOpcode::KILL));
     LLVM_DEBUG(dbgs() << "  replace by: " << MI);
     return;
@@ -641,11 +718,14 @@ void VirtRegRewriter::rewrite() {
   SmallVector<Register, 8> SuperDeads;
   SmallVector<Register, 8> SuperDefs;
   SmallVector<Register, 8> SuperKills;
+  uint64_t Mask;
 
   for (MachineFunction::iterator MBBI = MF->begin(), MBBE = MF->end();
        MBBI != MBBE; ++MBBI) {
     LLVM_DEBUG(MBBI->print(dbgs(), Indexes));
     for (MachineInstr &MI : llvm::make_early_inc_range(MBBI->instrs())) {
+      // reset for each MI.
+      Mask = 0u;
       for (MachineOperand &MO : MI.operands()) {
         // Make sure MRI knows about registers clobbered by regmasks.
         if (MO.isRegMask())
@@ -663,6 +743,9 @@ void VirtRegRewriter::rewrite() {
         RewriteRegs.insert(PhysReg);
         assert(!MRI->isReserved(PhysReg) && "Reserved register assignment");
 
+        if (MO.isUse() && MI.isCopy())
+          Mask = calcLiveRegUnitMask(MO, PhysReg);
+
         // Preserve semantics of sub-register operands.
         unsigned SubReg = MO.getSubReg();
         if (SubReg != 0) {
@@ -739,6 +822,10 @@ void VirtRegRewriter::rewrite() {
         MO.setIsRenamable(true);
       }
 
+      // Add LaneBitmask as MO_Imm
+      if (MI.isCopy() && Mask)
+        MI.addOperand(*MF, MachineOperand::CreateImm(Mask));
+
       // Add any missing super-register kills after rewriting the whole
       // instruction.
       while (!SuperKills.empty())