[llvm] [AMDGPU] Post-RA Peephole for Two-Address Instructions (PR #207731)

[Carl Ritson via llvm-commits]

================
@@ -763,13 +802,185 @@ MachineInstrBuilder SIPreEmitPeephole::createUnpackedMI(MachineInstr &I,
   return NewMI;
 }
 
+bool SIPreEmitPeephole::isLateThreeAddrPeepholeCandidate(
+    const MachineInstr &MI) const {
+  // This is a subset of the opcodes that are convertible to three-address
+  // instructions according to . This is intended. For some cases,
+  // convertToThreeAddress could early-clobber, which would require us to roll
+  // back the transformation. However, rollback is not side-effect free for some
+  // cases in convertToThreeAddress that fold immediates. Therefore, we limit
+  // this peephole to only the instructions for which rollback will not be
+  // necessary.
+  switch (MI.getOpcode()) {
+  case AMDGPU::V_MAC_F16_e32:
+  case AMDGPU::V_MAC_F16_e64:
+  case AMDGPU::V_MAC_F32_e32:
+  case AMDGPU::V_MAC_F32_e64:
+  case AMDGPU::V_MAC_LEGACY_F32_e32:
+  case AMDGPU::V_MAC_LEGACY_F32_e64:
+  case AMDGPU::V_FMAC_F16_e32:
+  case AMDGPU::V_FMAC_F16_e64:
+  case AMDGPU::V_FMAC_F16_t16_e64:
+  case AMDGPU::V_FMAC_F16_fake16_e64:
+  case AMDGPU::V_FMAC_F32_e32:
+  case AMDGPU::V_FMAC_F32_e64:
+  case AMDGPU::V_FMAC_LEGACY_F32_e32:
+  case AMDGPU::V_FMAC_LEGACY_F32_e64:
+  case AMDGPU::V_FMAC_F64_e32:
+  case AMDGPU::V_FMAC_F64_e64:
+    return true;
+  default:
+    return false;
+  }
+}
+
+Register SIPreEmitPeephole::checkCopy(MachineInstr &TwoAddress,
+                                      MachineInstr &MaybeCopy, Register Src,
+                                      const TargetRegisterClass *NewDstRC) {
+  if (MaybeCopy.isBundle())
+    return Register();
+
+  if (!TII->isFoldableCopy(MaybeCopy) ||
+      MaybeCopy.getOpcode() == AMDGPU::WWM_COPY)
+    return Register();
+
+  if (TII->hasAnyModifiersSet(MaybeCopy))
+    return Register();
+
+  auto CopySrc = MaybeCopy.getOperand(TII->getFoldableCopySrcIdx(MaybeCopy));
+  if (!CopySrc.isReg() || CopySrc.getReg() != Src ||
+      CopySrc.getSubReg() != AMDGPU::NoSubRegister)
+    return Register();
+
+  MachineOperand Dst = MaybeCopy.getOperand(0);
+  if (!Dst.isReg() || Dst.getSubReg() != AMDGPU::NoSubRegister)
+    return Register();
+
+  Register DstReg = Dst.getReg();
+  if (TRI->getPhysRegBaseClass(DstReg) != TRI->getPhysRegBaseClass(Src))
+    return Register();
+
+  if (NewDstRC && !NewDstRC->contains(DstReg.asMCReg()))
+    return Register();
+
+  // Check that it is safe to define the copy destination register at the
+  // current position of the two-address instruction.
+  SmallPtrSet<MachineInstr *, 1> Ignore{&TwoAddress, &MaybeCopy};
+  if (!RDI->isSafeToDefRegAt(&TwoAddress, DstReg, Ignore))
+    return Register();
+
+  // Check that the two-address instruction and the copy have the same exec
+  // mask.
+  if (!RDI->hasSameReachingDef(&TwoAddress, &MaybeCopy, AMDGPU::EXEC))
+    return Register();
+
+  return DstReg;
+}
+
+Register SIPreEmitPeephole::shouldReplaceWithThreeAddress(
+    MachineInstr &MI, SmallVectorImpl<MachineInstr *> &Copies) {
+  MachineOperand *Dst = TII->getNamedOperand(MI, AMDGPU::OpName::vdst);
+  if (!Dst || !Dst->isReg())
+    return Register();
+
+  const TargetRegisterClass *NewDstRC =
+      TII->getRegClass(MI.getDesc(), Dst->getOperandNo());
+
+  auto CheckRegisterCopies =
+      [&](Register Reg, const TargetRegisterClass *ReplacementRC) -> Register {
+    if (!Reg.isPhysical())
+      return Register();
+    if (RDI->getLocalLiveOutMIDef(MI.getParent(), Reg) == &MI)
+      return Register();
+
+    SmallPtrSet<MachineInstr *, 1> Uses;
+    RDI->getReachingLocalUses(&MI, Reg, Uses);
+    if (Uses.size() != 1)
+      return Register();
+
+    MachineInstr &Use = **Uses.begin();
+    Register ReplacementReg = checkCopy(MI, Use, Reg, ReplacementRC);
+
+    if (!ReplacementReg)
+      return Register();
+
+    Copies.push_back(&Use);
+    return ReplacementReg;
+  };
+
+  Register DstReg = Dst->getReg();
+  if (Register Replacement = CheckRegisterCopies(DstReg, NewDstRC))
+    return Replacement;
+
+  // If we didn't find a copy for the full register, we might be dealing with a
+  // 64-bit register and copies for the individual parts. Check both
+  // subregisters for matching copies.
+  Register Sub0 = TRI->getSubReg(Dst->getReg(), AMDGPU::sub0);
+  Register Sub1 = TRI->getSubReg(Dst->getReg(), AMDGPU::sub1);
+  if (!Sub0 || !Sub1)
+    return Register();
+
+  auto Sub0Replacment = CheckRegisterCopies(Sub0, nullptr);
+  auto Sub1Replacement = CheckRegisterCopies(Sub1, nullptr);
----------------
perlfu wrote:

```suggestion
  Register Sub1Replacement = CheckRegisterCopies(Sub1, nullptr);
```

https://github.com/llvm/llvm-project/pull/207731