CodeGen: Add TIED_COPY instruction

[arsenm]

Summary

I need a way to be able to constrain a register def to use the same
input register throughout a loop. For a VGPR, we don't model the other
live values in other lanes. In cases where we are explicitly
scalarizing operations and writing different lanes sequentially, we
need to model this as an input tied operand to the def instruction
inside the loop. The verifier doesn't allow adding a tied implicit use
to an explicit def, so move this constraint into a copy we can use to
accumulate the real result. This should be folded out after allocation
like a normal copy when possible.

[nhaehnle]

Do you have a use case for this, or can you explain a bit more what this solves, perhaps with an example?

Normal waterfall loops ought to be solved in the same way as any other loop with divergent exit: a value written inside the loop is observed from outside the loop, so there should be a COPY to a vgpr / vector register bank if necessary, but then the regular register allocation etc. does the rest. We don't need a TIED_COPY there today, or rather, I don't think we should need one.

[arsenm]

llvm/include/llvm/CodeGen/MachineInstr.h

public:
   bool isCopy() const {
     return getOpcode() == TargetOpcode::COPY;
   }
 
   bool isFullCopy() const {
     return isCopy() && !getOperand(0).getSubReg() && !getOperand(1).getSubReg();
   }
 
+  bool isTiedCopy() const {
+    return getOpcode() == TargetOpcode::TIED_COPY;
+  }
+
   bool isExtractSubreg() const {
     return getOpcode() == TargetOpcode::EXTRACT_SUBREG;
   }
 
   /// Return true if the instruction behaves like a copy.
   /// This does not include native copy instructions.
   bool isCopyLike() const {
     return isCopy() || isSubregToReg();
   }
 
   /// 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();
   }
 
+  bool isIdentityTiedCopy() const {
+    return isTiedCopy() && getOperand(0).getReg() == getOperand(1).getReg() &&
+      getOperand(0).getSubReg() == getOperand(1).getSubReg();
+  }
+
   /// Return true if this instruction doesn't produce any output in the form of
   /// executable instructions.
   bool isMetaInstruction() const {
     switch (getOpcode()) {
     default:
       return false;
     case TargetOpcode::IMPLICIT_DEF:
     case TargetOpcode::KILL:

llvm/include/llvm/Support/TargetOpcodes.def

 /// After register coalescing references of v1024 should be replace with
 /// v1027:3, v1025 with v1027:4, etc.
   HANDLE_TARGET_OPCODE(REG_SEQUENCE)
 
 /// COPY - Target-independent register copy. This instruction can also be
 /// used to copy between subregisters of virtual registers.
   HANDLE_TARGET_OPCODE(COPY)
 
+/// COPY - Target-independent register copy with tied constraint.
+HANDLE_TARGET_OPCODE(TIED_COPY)
+
 /// BUNDLE - This instruction represents an instruction bundle. Instructions
 /// which immediately follow a BUNDLE instruction which are marked with
 /// 'InsideBundle' flag are inside the bundle.
 HANDLE_TARGET_OPCODE(BUNDLE)
 
 /// Lifetime markers.
 HANDLE_TARGET_OPCODE(LIFETIME_START)
 HANDLE_TARGET_OPCODE(LIFETIME_END)

llvm/include/llvm/Target/Target.td

 def COPY : StandardPseudoInstruction {
   let OutOperandList = (outs unknown:$dst);
   let InOperandList = (ins unknown:$src);
   let AsmString = "";
   let hasSideEffects = 0;
   let isAsCheapAsAMove = 1;
   let hasNoSchedulingInfo = 0;
 }
+
+// A copy with the requirement that the source and result registers
+// must be allocated to the same register. This can be used for
+// constraining a value to a single register inside a loop in SSA
+// form.
+def TIED_COPY : StandardPseudoInstruction {
+  let OutOperandList = (outs unknown:$dst);
+  let InOperandList = (ins unknown:$src);
+  let Constraints = "$dst = $src";
+  let AsmString = "";
+  let hasSideEffects = 0;
+  let isAsCheapAsAMove = 1;
+}
+
 def BUNDLE : StandardPseudoInstruction {
   let OutOperandList = (outs);
   let InOperandList = (ins variable_ops);
   let AsmString = "BUNDLE";
   let hasSideEffects = 0;
 }
 def LIFETIME_START : StandardPseudoInstruction {
   let OutOperandList = (outs);

llvm/lib/CodeGen/ExpandPostRAPseudos.cpp

for (MachineBasicBlock::iterator mi = mbbi->begin(), me = mbbi->end();
       }
 
       // Expand standard pseudos.
       switch (MI.getOpcode()) {
       case TargetOpcode::SUBREG_TO_REG:
         MadeChange |= LowerSubregToReg(&MI);
         break;
       case TargetOpcode::COPY:
+      case TargetOpcode::TIED_COPY:
         MadeChange |= LowerCopy(&MI);
         break;
       case TargetOpcode::DBG_VALUE:
         continue;
       case TargetOpcode::INSERT_SUBREG:
       case TargetOpcode::EXTRACT_SUBREG:
         llvm_unreachable("Sub-register pseudos should have been eliminated.");
       }
     }
   }
 
   return MadeChange;
 }

llvm/lib/CodeGen/VirtRegMap.cpp

bool VirtRegRewriter::readsUndefSubreg(const MachineOperand &MO) const {
   for (const LiveInterval::SubRange &SR : LI.subranges()) {
     if ((SR.LaneMask & UseMask).any() && SR.liveAt(BaseIndex))
       return false;
   }
   return true;
 }
 
 void VirtRegRewriter::handleIdentityCopy(MachineInstr &MI) const {
-  if (!MI.isIdentityCopy())
+  if (!MI.isIdentityCopy() && !MI.isIdentityTiedCopy())
     return;
   LLVM_DEBUG(dbgs() << "Identity copy: " << MI);
   ++NumIdCopies;
 
   // Copies like:
   //    %r0 = COPY undef %r0
   //    %al = COPY %al, implicit-def %eax
   // give us additional liveness information: The target (super-)register

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

if (setM0ToIndexFromSGPR(TII, MRI, MI, Offset, UseGPRIdxMode, true)) {
 
     return &MBB;
   }
 
   const DebugLoc &DL = MI.getDebugLoc();
   MachineBasicBlock::iterator I(&MI);
 
   Register PhiReg = MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
+  Register TmpDstReg = MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
   Register InitReg = MRI.createVirtualRegister(&AMDGPU::VGPR_32RegClass);
 
   BuildMI(MBB, I, DL, TII->get(TargetOpcode::IMPLICIT_DEF), InitReg);
 
   auto InsPt = loadM0FromVGPR(TII, MBB, MI, InitReg, PhiReg,
                               Offset, UseGPRIdxMode, true);
   MachineBasicBlock *LoopBB = InsPt->getParent();
 
   if (UseGPRIdxMode) {
-    BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::V_MOV_B32_e32), Dst)
+    BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::V_MOV_B32_e32), TmpDstReg)
       .addReg(SrcReg, RegState::Undef, SubReg)
       .addReg(SrcReg, RegState::Implicit)
       .addReg(AMDGPU::M0, RegState::Implicit);
     BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::S_SET_GPR_IDX_OFF));
   } else {
-    BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::V_MOVRELS_B32_e32), Dst)
+    BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::V_MOVRELS_B32_e32), TmpDstReg)
       .addReg(SrcReg, RegState::Undef, SubReg)
       .addReg(SrcReg, RegState::Implicit);
   }
 
+  BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::TIED_COPY), Dst)
+    .addReg(TmpDstReg);
+
   MI.eraseFromParent();
 
   return LoopBB;
 }
 
 static MachineBasicBlock *emitIndirectDst(MachineInstr &MI,
                                           MachineBasicBlock &MBB,
                                           const GCNSubtarget &ST) {
static MachineBasicBlock *emitIndirectDst(MachineInstr &MI,
   }
 
   if (Val->isReg())
     MRI.clearKillFlags(Val->getReg());
 
   const DebugLoc &DL = MI.getDebugLoc();
 
   Register PhiReg = MRI.createVirtualRegister(VecRC);
+  Register TmpDstReg = MRI.createVirtualRegister(VecRC);
 
   auto InsPt = loadM0FromVGPR(TII, MBB, MI, SrcVec->getReg(), PhiReg,
                               Offset, UseGPRIdxMode, false);
   MachineBasicBlock *LoopBB = InsPt->getParent();
 
-  BuildMI(*LoopBB, InsPt, DL, MovRelDesc, Dst)
+  BuildMI(*LoopBB, InsPt, DL, MovRelDesc, TmpDstReg)
     .addReg(PhiReg)
     .add(*Val)
     .addImm(AMDGPU::sub0);
   if (UseGPRIdxMode)
     BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::S_SET_GPR_IDX_OFF));
 
+  BuildMI(*LoopBB, InsPt, DL, TII->get(AMDGPU::TIED_COPY), Dst)
+    .addReg(TmpDstReg);
+
   MI.eraseFromParent();
   return LoopBB;
 }
 
 MachineBasicBlock *SITargetLowering::EmitInstrWithCustomInserter(
   MachineInstr &MI, MachineBasicBlock *BB) const {
 
   const SIInstrInfo *TII = getSubtarget()->getInstrInfo();