[WIP][AMDGPU][CopyPhysReg] Expand the COPY instructions using the encoded liveness mask.

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -1473,8 +1473,9 @@ class MachineInstr
 
   /// Return true is the instruction is an identity copy.
   bool isIdentityCopy() const {
-    return isCopy() && getOperand(0).getReg() == getOperand(1).getReg() &&
-      getOperand(0).getSubReg() == getOperand(1).getSubReg();
+    return (isCopy() || isCopyLaneMask()) &&
+           getOperand(0).getReg() == getOperand(1).getReg() &&
+           getOperand(0).getSubReg() == getOperand(1).getSubReg();
   }
 
   /// Return true if this is a transient instruction that is either very likely

llvm/lib/CodeGen/ExpandPostRAPseudos.cpp

@@ -168,6 +168,7 @@ bool ExpandPostRA::run(MachineFunction &MF) {
         MadeChange |= LowerSubregToReg(&MI);
         break;
       case TargetOpcode::COPY:
+      case TargetOpcode::COPY_LANEMASK:
         TII->lowerCopy(&MI, TRI);
         MadeChange = true;
         break;

llvm/lib/CodeGen/TargetInstrInfo.cpp

@@ -884,12 +884,21 @@ void TargetInstrInfo::lowerCopy(
   if (IdentityCopy || SrcMO.isUndef()) {
     // No need to insert an identity copy instruction, but replace with a KILL
     // if liveness is changed.
-    if (SrcMO.isUndef() || MI->getNumOperands() > 2) {
+    if (MI->getOpcode() == TargetOpcode::COPY &&
+        (SrcMO.isUndef() || MI->getNumOperands() > 2)) {
       // We must make sure the super-register gets killed. Replace the
       // instruction with KILL.
       MI->setDesc(get(TargetOpcode::KILL));
       return;
     }
+    if (MI->getOpcode() == TargetOpcode::COPY_LANEMASK &&
+        (SrcMO.isUndef() || MI->getNumOperands() > 3)) {
+      // We must make sure the super-register gets killed. Replace the
+      // instruction with KILL.
+      MI->setDesc(get(TargetOpcode::KILL));
+      return;
+    }
+
     // Vanilla identity copy.
     MI->eraseFromParent();
     return;
@@ -900,7 +909,10 @@ void TargetInstrInfo::lowerCopy(
               DstMO.getReg().isPhysical() ? DstMO.isRenamable() : false,
               SrcMO.getReg().isPhysical() ? SrcMO.isRenamable() : false);
 
-  if (MI->getNumOperands() > 2)
+  if (MI->getOpcode() == TargetOpcode::COPY && MI->getNumOperands() > 2)
+    transferImplicitOperands(MI, &TRI);
+  if (MI->getOpcode() == TargetOpcode::COPY_LANEMASK &&
+      MI->getNumOperands() > 3)
     transferImplicitOperands(MI, &TRI);
   MI->eraseFromParent();
 }

llvm/lib/CodeGen/VirtRegMap.cpp

@@ -213,6 +213,8 @@ class VirtRegRewriter {
   void rewrite();
   void addMBBLiveIns();
   bool readsUndefSubreg(const MachineOperand &MO) const;
+  LaneBitmask 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,80 @@ bool VirtRegRewriter::readsUndefSubreg(const MachineOperand &MO) const {
   return true;
 }
 
+// Return LaneBitmask value for PhysReg assigned to MO,
+// representing its live register units at its parent MI. In case of undef or
+// fully live MO, return 0u.
+LaneBitmask 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<MCRegUnit> LiveRegUnits;
+
+  // dbgs() << "\n********** " << printReg(Reg, TRI) << "[ " <<
+  // printReg(PhysReg, TRI) << " ]" << " **********\n";
+
+  if (MO.isUndef())
+    return LaneBitmask::getNone();
+
+  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) {
+      MCRegUnit 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) {
+      MCRegUnit 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) {
+    MCRegUnit 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 all lanes are live or dead, no need to create a COPY_LANEMASK
+  // instruction.
+  if (LiveRegUnitMask.all() || LiveRegUnitMask.none() ||
+      LiveRegUnitMask == UseMask)
+    return LaneBitmask::getNone();
+
+  return LiveRegUnitMask;
+}
+
 void VirtRegRewriter::handleIdentityCopy(MachineInstr &MI) {
   if (!MI.isIdentityCopy())
     return;
@@ -495,7 +571,14 @@ 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) {
+  if (MI.getOpcode() == TargetOpcode::COPY &&
+      (MI.getOperand(1).isUndef() || MI.getNumOperands() > 2)) {
+    MI.setDesc(TII->get(TargetOpcode::KILL));
+    LLVM_DEBUG(dbgs() << "  replace by: " << MI);
+    return;
+  }
+  if (MI.getOpcode() == TargetOpcode::COPY_LANEMASK &&
+      (MI.getOperand(1).isUndef() || MI.getNumOperands() > 3)) {
     MI.setDesc(TII->get(TargetOpcode::KILL));
     LLVM_DEBUG(dbgs() << "  replace by: " << MI);
     return;
@@ -641,11 +724,14 @@ void VirtRegRewriter::rewrite() {
   SmallVector<Register, 8> SuperDeads;
   SmallVector<Register, 8> SuperDefs;
   SmallVector<Register, 8> SuperKills;
+  LaneBitmask LaneMask;
 
   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.
+      LaneMask = LaneBitmask::getNone();
       for (MachineOperand &MO : MI.operands()) {
         // Make sure MRI knows about registers clobbered by regmasks.
         if (MO.isRegMask())
@@ -663,6 +749,9 @@ void VirtRegRewriter::rewrite() {
         RewriteRegs.insert(PhysReg);
         assert(!MRI->isReserved(PhysReg) && "Reserved register assignment");
 
+        if (MO.isUse() && MI.isCopy())
+          LaneMask = calcLiveRegUnitMask(MO, PhysReg);
+
         // Preserve semantics of sub-register operands.
         unsigned SubReg = MO.getSubReg();
         if (SubReg != 0) {
@@ -739,6 +828,13 @@ void VirtRegRewriter::rewrite() {
         MO.setIsRenamable(true);
       }
 
+      // If there are any live lanes, replace a COPY instruction with a
+      // COPY_LANEMASK instruction with the lane mask.
+      if (MI.isCopy() && LaneMask.any()) {
+        MI.setDesc(TII->get(TargetOpcode::COPY_LANEMASK));
+        MI.addOperand(*MF, MachineOperand::CreateLaneMask(LaneMask));
+      }
+
       // Add any missing super-register kills after rewriting the whole
       // instruction.
       while (!SuperKills.empty())