Skip to content

AMDGPU: Handle rewriting VGPR MFMA fed from AGPR copy #153022

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+257 −298
Open

AMDGPU: Handle rewriting VGPR MFMA fed from AGPR copy #153022

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
306 changes: 191 additions & 115 deletions llvm/lib/Target/AMDGPU/AMDGPURewriteAGPRCopyMFMA.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,10 @@
/// MFMA opcode.
///
/// TODO:
/// - Handle rewrites of phis. This must be more careful than normal about the
/// reassignment. We do not want to introduce an AGPR-to-AGPR copy inside of a
/// loop, so it depends on the exact assignment of the copy.
///
/// - Update LiveIntervals incrementally instead of recomputing from scratch
///
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -60,6 +64,32 @@ class AMDGPURewriteAGPRCopyMFMAImpl {
return TII.isMAI(MI) && AMDGPU::getMFMASrcCVDstAGPROp(MI.getOpcode()) != -1;
}

/// Find AV_* registers assigned to AGPRs (or virtual registers which were
/// already required to be AGPR).
///
/// \return the assigned physical register that \p VReg is assigned to if it
/// is an AGPR, otherwise MCRegister().
MCRegister getAssignedAGPR(Register VReg) const {
MCRegister PhysReg = VRM.getPhys(VReg);
if (!PhysReg)
return MCRegister();

const TargetRegisterClass *VirtRegRC = MRI.getRegClass(VReg);
if (!TRI.hasAGPRs(VirtRegRC))
return MCRegister();

if (!TRI.hasVGPRs(VirtRegRC))
return PhysReg;

// If this is an AV register, we have to check if the actual assignment is
// to an AGPR
const TargetRegisterClass *AssignedRC = TRI.getPhysRegBaseClass(PhysReg);
return TRI.isAGPRClass(AssignedRC) ? PhysReg : MCRegister();
}

bool tryReassigningMFMAChain(MachineInstr &MFMA, unsigned HintOpIdx,
MCPhysReg PhysRegHint) const;

/// Compute the register class constraints based on the uses of \p Reg,
/// excluding MFMA uses from which can be rewritten to change the register
/// class constraint. This should be nearly identical to
Expand All @@ -74,6 +104,8 @@ class AMDGPURewriteAGPRCopyMFMAImpl {
Register Reg, SmallVectorImpl<MachineInstr *> &RewriteCandidates,
SmallSetVector<Register, 4> &RewriteRegs) const;

bool tryFoldCopiesToAGPR(Register VReg, MCRegister AssignedAGPR) const;
bool tryFoldCopiesFromAGPR(Register VReg, MCRegister AssignedAGPR) const;
bool run(MachineFunction &MF) const;
};

Expand Down Expand Up @@ -154,6 +186,88 @@ bool AMDGPURewriteAGPRCopyMFMAImpl::recomputeRegClassExceptRewritable(
return true;
}

bool AMDGPURewriteAGPRCopyMFMAImpl::tryReassigningMFMAChain(
MachineInstr &MFMA, unsigned HintOpIdx, MCPhysReg PhysRegHint) const {
// src2 and dst have the same physical class constraint; try to preserve
// the original src2 subclass if one were to exist.
SmallVector<MachineInstr *, 4> RewriteCandidates = {&MFMA};
SmallSetVector<Register, 4> RewriteRegs;

Register MFMAHintReg = MFMA.getOperand(HintOpIdx).getReg();
// Make sure we reassign the MFMA we found the copy from first. We want
// to ensure dst ends up in the physreg we were originally copying to.
RewriteRegs.insert(MFMAHintReg);

// We've found av = COPY (MFMA), and need to verify that we can trivially
// rewrite src2 to use the new AGPR. If we can't trivially replace it,
// we're going to induce as many copies as we would have emitted in the
// first place, as well as need to assign another register, and need to
// figure out where to put them. The live range splitting is smarter than
// anything we're doing here, so trust it did something reasonable.
//
// Note recomputeRegClassExceptRewritable will consider the constraints of
// this MFMA's src2 as well as the src2/dst of any transitive MFMA users.
if (!recomputeRegClassExceptRewritable(MFMAHintReg, RewriteCandidates,
RewriteRegs)) {
LLVM_DEBUG(dbgs() << "Could not recompute the regclass of dst reg "
<< printReg(MFMAHintReg, &TRI) << '\n');
return false;
}

// If src2 and dst are different registers, we need to also reassign the
// input to an available AGPR if it is compatible with all other uses.
//
// If we can't reassign it, we'd need to introduce a different copy
// which is likely worse than the copy we'd be saving.
//
// It's likely that the MFMA is used in sequence with other MFMAs; if we
// cannot migrate the full use/def chain of MFMAs, we would need to
// introduce intermediate copies somewhere. So we only make the
// transform if all the interfering MFMAs can also be migrated. Collect
// the set of rewritable MFMAs and check if we can assign an AGPR at
// that point.
//
// If any of the MFMAs aren't reassignable, we give up and rollback to
// the original register assignments.

using RecoloringStack =
SmallVector<std::pair<const LiveInterval *, MCRegister>, 8>;
RecoloringStack TentativeReassignments;

for (Register RewriteReg : RewriteRegs) {
LiveInterval &LI = LIS.getInterval(RewriteReg);
TentativeReassignments.push_back({&LI, VRM.getPhys(RewriteReg)});
LRM.unassign(LI);
}

if (!attemptReassignmentsToAGPR(RewriteRegs, PhysRegHint)) {
// Roll back the register assignments to the original state.
for (auto [LI, OldAssign] : TentativeReassignments) {
if (VRM.hasPhys(LI->reg()))
LRM.unassign(*LI);
LRM.assign(*LI, OldAssign);
}

return false;
}

// Fixup the register classes of the virtual registers now that we've
// committed to the reassignments.
for (Register InterferingReg : RewriteRegs) {
const TargetRegisterClass *EquivalentAGPRRegClass =
TRI.getEquivalentAGPRClass(MRI.getRegClass(InterferingReg));
MRI.setRegClass(InterferingReg, EquivalentAGPRRegClass);
}

for (MachineInstr *RewriteCandidate : RewriteCandidates) {
int NewMFMAOp =
AMDGPU::getMFMASrcCVDstAGPROp(RewriteCandidate->getOpcode());
RewriteCandidate->setDesc(TII.get(NewMFMAOp));
}

return true;
}

/// Attempt to reassign the registers in \p InterferingRegs to be AGPRs, with a
/// preference to use \p PhysReg first. Returns false if the reassignments
/// cannot be trivially performed.
Expand Down Expand Up @@ -206,6 +320,78 @@ bool AMDGPURewriteAGPRCopyMFMAImpl::attemptReassignmentsToAGPR(
return true;
}

/// Identify copies that look like:
/// %vdst:vgpr = V_MFMA_.. %src0:av, %src1:av, %src2:vgpr
/// %agpr = COPY %vgpr
///
/// Then try to replace the transitive uses of %src2 and %vdst with the AGPR
/// versions of the MFMA. This should cover the common case.
bool AMDGPURewriteAGPRCopyMFMAImpl::tryFoldCopiesToAGPR(
Register VReg, MCRegister AssignedAGPR) const {
bool MadeChange = false;
for (MachineInstr &UseMI : MRI.def_instructions(VReg)) {
if (!UseMI.isCopy())
continue;

Register CopySrcReg = UseMI.getOperand(1).getReg();
if (!CopySrcReg.isVirtual())
continue;

// TODO: Handle loop phis copied to AGPR. e.g.
//
// loop:
// %phi:vgpr = COPY %mfma:vgpr
// %mfma:vgpr = V_MFMA_xxx_vgprcd_e64 %a, %b, %phi
// s_cbranch_vccnz loop
//
// endloop:
// %agpr = mfma
//
// We need to be sure that %phi is assigned to the same physical register as
// %mfma, or else we will just be moving copies into the loop.

for (MachineInstr &CopySrcDefMI : MRI.def_instructions(CopySrcReg)) {
if (isRewriteCandidate(CopySrcDefMI) &&
tryReassigningMFMAChain(CopySrcDefMI, 0, AssignedAGPR))
MadeChange = true;
}
}

return MadeChange;
}

/// Identify copies that look like:
/// %src:vgpr = COPY %src:agpr
/// %vdst:vgpr = V_MFMA_... %src0:av, %src1:av, %src:vgpr
///
/// Then try to replace the transitive uses of %src2 and %vdst with the AGPR
/// versions of the MFMA. This should cover rarer cases, and will generally be
/// redundant with tryFoldCopiesToAGPR.
bool AMDGPURewriteAGPRCopyMFMAImpl::tryFoldCopiesFromAGPR(
Register VReg, MCRegister AssignedAGPR) const {
bool MadeChange = false;
for (MachineInstr &UseMI : MRI.use_instructions(VReg)) {
if (!UseMI.isCopy())
continue;

Register CopyDstReg = UseMI.getOperand(0).getReg();
if (!CopyDstReg.isVirtual())
continue;

for (MachineInstr &CopyUseMI : MRI.use_instructions(CopyDstReg)) {
if (isRewriteCandidate(CopyUseMI)) {
const MachineOperand *Op =
CopyUseMI.findRegisterUseOperand(CopyDstReg, /*TRI=*/nullptr);
if (tryReassigningMFMAChain(CopyUseMI, Op->getOperandNo(),
VRM.getPhys(Op->getReg())))
MadeChange = true;
}
}
}

return MadeChange;
}

bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const {
// This only applies on subtargets that have a configurable AGPR vs. VGPR
// allocation.
Expand All @@ -222,124 +408,14 @@ bool AMDGPURewriteAGPRCopyMFMAImpl::run(MachineFunction &MF) const {

for (unsigned I = 0, E = MRI.getNumVirtRegs(); I != E; ++I) {
Register VReg = Register::index2VirtReg(I);
Register PhysReg = VRM.getPhys(VReg);
if (!PhysReg)
continue;

// Find AV_* registers assigned to AGPRs.
const TargetRegisterClass *VirtRegRC = MRI.getRegClass(VReg);
if (!TRI.hasAGPRs(VirtRegRC))
MCRegister AssignedAGPR = getAssignedAGPR(VReg);
if (!AssignedAGPR)
continue;

const TargetRegisterClass *AssignedRC = VirtRegRC;
if (TRI.hasVGPRs(VirtRegRC)) {
// If this is an AV register, we have to check if the actual assignment is
// to an AGPR
AssignedRC = TRI.getPhysRegBaseClass(PhysReg);
if (!TRI.isAGPRClass(AssignedRC))
continue;
}

LiveInterval &LI = LIS.getInterval(VReg);

for (VNInfo *VNI : LI.vnis()) {
if (VNI->isPHIDef() || VNI->isUnused())
continue;

MachineInstr *DefMI = LIS.getInstructionFromIndex(VNI->def);
if (!DefMI || !DefMI->isCopy())
continue;

Register MFMADstReg = DefMI->getOperand(1).getReg();
if (!MFMADstReg.isVirtual())
continue;

LiveInterval &CopySrcLI = LIS.getInterval(MFMADstReg);
LiveQueryResult LRQ = CopySrcLI.Query(VNI->def.getRegSlot());
MachineInstr *MFMA = LIS.getInstructionFromIndex(LRQ.valueIn()->def);
if (!MFMA || !isRewriteCandidate(*MFMA))
continue;

// src2 and dst have the same physical class constraint; try to preserve
// the original src2 subclass if one were to exist.
SmallVector<MachineInstr *, 4> RewriteCandidates = {MFMA};
SmallSetVector<Register, 4> RewriteRegs;

// Make sure we reassign the MFMA we found the copy from first. We want
// to ensure dst ends up in the physreg we were originally copying to.
RewriteRegs.insert(MFMADstReg);

// We've found av = COPY (MFMA), and need to verify that we can trivially
// rewrite src2 to use the new AGPR. If we can't trivially replace it,
// we're going to induce as many copies as we would have emitted in the
// first place, as well as need to assign another register, and need to
// figure out where to put them. The live range splitting is smarter than
// anything we're doing here, so trust it did something reasonable.
//
// Note recomputeRegClassExceptRewritable will consider the constraints of
// this MFMA's src2 as well as the src2/dst of any transitive MFMA users.
if (!recomputeRegClassExceptRewritable(MFMADstReg, RewriteCandidates,
RewriteRegs)) {
LLVM_DEBUG(dbgs() << "Could not recompute the regclass of dst reg "
<< printReg(MFMADstReg, &TRI) << '\n');
continue;
}

// If src2 and dst are different registers, we need to also reassign the
// input to an available AGPR if it is compatible with all other uses.
//
// If we can't reassign it, we'd need to introduce a different copy
// which is likely worse than the copy we'd be saving.
//
// It's likely that the MFMA is used in sequence with other MFMAs; if we
// cannot migrate the full use/def chain of MFMAs, we would need to
// introduce intermediate copies somewhere. So we only make the
// transform if all the interfering MFMAs can also be migrated. Collect
// the set of rewritable MFMAs and check if we can assign an AGPR at
// that point.
//
// If any of the MFMAs aren't reassignable, we give up and rollback to
// the original register assignments.

using RecoloringStack =
SmallVector<std::pair<const LiveInterval *, MCRegister>, 8>;
RecoloringStack TentativeReassignments;

for (Register RewriteReg : RewriteRegs) {
LiveInterval &LI = LIS.getInterval(RewriteReg);
TentativeReassignments.push_back({&LI, VRM.getPhys(RewriteReg)});
LRM.unassign(LI);
}

if (!attemptReassignmentsToAGPR(RewriteRegs, PhysReg)) {
// Roll back the register assignments to the original state.
for (auto [LI, OldAssign] : TentativeReassignments) {
if (VRM.hasPhys(LI->reg()))
LRM.unassign(*LI);
LRM.assign(*LI, OldAssign);
}

continue;
}

// Fixup the register classes of the virtual registers now that we've
// committed to the reassignments.
for (Register InterferingReg : RewriteRegs) {
const TargetRegisterClass *EquivalentAGPRRegClass =
TRI.getEquivalentAGPRClass(MRI.getRegClass(InterferingReg));
MRI.setRegClass(InterferingReg, EquivalentAGPRRegClass);
}

for (MachineInstr *RewriteCandidate : RewriteCandidates) {
int NewMFMAOp =
AMDGPU::getMFMASrcCVDstAGPROp(RewriteCandidate->getOpcode());
RewriteCandidate->setDesc(TII.get(NewMFMAOp));
}

// We likely left an identity copy behind after assignment; let
// VirtRegRewriter deal with it later.
if (tryFoldCopiesToAGPR(VReg, AssignedAGPR))
MadeChange = true;
if (tryFoldCopiesFromAGPR(VReg, AssignedAGPR))
MadeChange = true;
}
}

return MadeChange;
Expand Down
Loading
Loading