AMDGPU/GlobalISel: Avoid handling registers twice in waterfall loops

[arsenm]

Summary

When multiple instructions are moved into a waterfall loop, it's
possible some of them re-use the same operands. Avoid creating
multiple sequences of readfirstlanes for them. None of the current
uses will hit this, but will be used in a future patch.

[kerbowa]

LGTM

[arsenm]

5a8c0f552b382cbf4397aff9cc3b0d9418bdf3ff

llvm/lib/Target/AMDGPU/AMDGPURegisterBankInfo.cpp

bool AMDGPURegisterBankInfo::executeInWaterfallLoop(
   MachineIRBuilder &B,
   iterator_range<MachineBasicBlock::iterator> Range,
   SmallSet<Register, 4> &SGPROperandRegs,
   MachineRegisterInfo &MRI) const {
   SmallVector<Register, 4> ResultRegs;
   SmallVector<Register, 4> InitResultRegs;
   SmallVector<Register, 4> PhiRegs;
 
+  // Track use registers which have already been expanded with a readfirstlane
+  // sequence. This may have multiple uses if moving a sequence.
+  DenseMap<Register, Register> WaterfalledRegMap;
+
   MachineBasicBlock &MBB = B.getMBB();
   MachineFunction *MF = &B.getMF();
 
   const TargetRegisterClass *WaveRC = TRI->getWaveMaskRegClass();
   const unsigned WaveAndOpc = Subtarget.isWave32() ?
     AMDGPU::S_AND_B32 : AMDGPU::S_AND_B64;
   const unsigned MovTermOpc = Subtarget.isWave32() ?
     AMDGPU::S_MOV_B32_term : AMDGPU::S_MOV_B64_term;
#endif
 
   assert(std::distance(NewBegin, NewEnd) == OrigRangeSize);
 
   for (MachineInstr &MI : make_range(NewBegin, NewEnd)) {
     for (MachineOperand &Op : MI.uses()) {
       if (!Op.isReg() || Op.isDef())
         continue;
 
-      if (!SGPROperandRegs.count(Op.getReg()))
+      Register OldReg = Op.getReg();
+      if (!SGPROperandRegs.count(OldReg))
         continue;
 
+      // See if we already processed this register in another instruction in the
+      // sequence.
+      auto OldVal = WaterfalledRegMap.find(OldReg);
+      if (OldVal != WaterfalledRegMap.end()) {
+        Op.setReg(OldVal->second);
+        continue;
+      }
+
       LLT OpTy = MRI.getType(Op.getReg());
       unsigned OpSize = OpTy.getSizeInBits();
 
       // Can only do a readlane of 32-bit pieces.
       if (OpSize == 32) {
         // Avoid extra copies in the simple case of one 32-bit register.
         Register CurrentLaneOpReg
           = MRI.createVirtualRegister(&AMDGPU::SReg_32_XM0RegClass);
for (MachineOperand &Op : MI.uses()) {
           Op.setReg(Merge.getReg(0));
         } else {
           auto Merge = B.buildMerge(OpTy, ReadlanePieces);
           Op.setReg(Merge.getReg(0));
         }
 
         MRI.setRegBank(Op.getReg(), AMDGPU::SGPRRegBank);
       }
+
+      // Make sure we don't re-process this register again.
+      WaterfalledRegMap.insert(std::make_pair(OldReg, Op.getReg()));
     }
   }
 
   B.setInsertPt(*LoopBB, LoopBB->end());
 
   // Update EXEC, save the original EXEC value to VCC.
   B.buildInstr(AndSaveExecOpc)
     .addDef(NewExec)