Skip to content

[AArch64][SME] Make getRegAllocationHints more specific for multi-vector loads #123081

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

Merged
kmclaughlin-arm merged 7 commits into from
Jan 30, 2025
Merged

[AArch64][SME] Make getRegAllocationHints more specific for multi-vector loads #123081

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
f9d4ac8
Add uses of %scalable_arg in SME2 dot intrinsics tests
kmclaughlin-arm Jan 15, 2025
68e65db
[AArch64][SME] Make getRegAllocationHints stricter for multi-vector l…
kmclaughlin-arm Jan 15, 2025
7635e0e
- Instead of finding a hint based on the operand index from the FORM_…
kmclaughlin-arm Jan 17, 2025
81c3d47
- Rewrote search for operand which already has a register assigned
kmclaughlin-arm Jan 27, 2025
607e3df
- Allow mix of x2 & x4 multivector loads and intrinsics
kmclaughlin-arm Jan 29, 2025
65e937d
- Renamed LdOps to UseOps
kmclaughlin-arm Jan 29, 2025
a651b43
- Moved MRI def after if condition for hasSME()
kmclaughlin-arm Jan 30, 2025
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
99 changes: 81 additions & 18 deletions llvm/lib/Target/AArch64/AArch64RegisterInfo.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
#include "MCTargetDesc/AArch64InstPrinter.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/BinaryFormat/Dwarf.h"
#include "llvm/CodeGen/LiveRegMatrix.h"
#include "llvm/CodeGen/MachineFrameInfo.h"
#include "llvm/CodeGen/MachineInstrBuilder.h"
#include "llvm/CodeGen/MachineRegisterInfo.h"
Expand Down Expand Up @@ -1099,6 +1100,11 @@ bool AArch64RegisterInfo::getRegAllocationHints(
const VirtRegMap *VRM, const LiveRegMatrix *Matrix) const {
const MachineRegisterInfo &MRI = MF.getRegInfo();
Copy link
Collaborator

@sdesmalen-arm sdesmalen-arm Jan 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this can be moved down (after the if(..) condition)


auto &ST = MF.getSubtarget<AArch64Subtarget>();
if (!ST.hasSME() || !ST.isStreaming())
return TargetRegisterInfo::getRegAllocationHints(VirtReg, Order, Hints, MF,
VRM);

// The SVE calling convention preserves registers Z8-Z23. As a result, there
// are no ZPR2Strided or ZPR4Strided registers that do not overlap with the
// callee-saved registers and so by default these will be pushed to the back
Expand All @@ -1108,25 +1114,82 @@ bool AArch64RegisterInfo::getRegAllocationHints(
// instructions over reducing the number of clobbered callee-save registers,
// so we add the strided registers as a hint.
unsigned RegID = MRI.getRegClass(VirtReg)->getID();
// Look through uses of the register for FORM_TRANSPOSED_REG_TUPLE.
if ((RegID == AArch64::ZPR2StridedOrContiguousRegClassID ||
RegID == AArch64::ZPR4StridedOrContiguousRegClassID) &&
any_of(MRI.use_nodbg_instructions(VirtReg), [](const MachineInstr &Use) {
return Use.getOpcode() ==
AArch64::FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO ||
Use.getOpcode() == AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO;
})) {
const TargetRegisterClass *StridedRC =
RegID == AArch64::ZPR2StridedOrContiguousRegClassID
? &AArch64::ZPR2StridedRegClass
: &AArch64::ZPR4StridedRegClass;

for (MCPhysReg Reg : Order)
if (StridedRC->contains(Reg))
Hints.push_back(Reg);
if (RegID == AArch64::ZPR2StridedOrContiguousRegClassID ||
RegID == AArch64::ZPR4StridedOrContiguousRegClassID) {

// Look through uses of the register for FORM_TRANSPOSED_REG_TUPLE.
for (const MachineInstr &Use : MRI.use_nodbg_instructions(VirtReg)) {
if (Use.getOpcode() != AArch64::FORM_TRANSPOSED_REG_TUPLE_X2_PSEUDO &&
Use.getOpcode() != AArch64::FORM_TRANSPOSED_REG_TUPLE_X4_PSEUDO)
continue;

unsigned LdOps = Use.getNumOperands() - 1;
Copy link
Collaborator

@sdesmalen-arm sdesmalen-arm Jan 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I just realised that LdOps is a misnomer, because it is assigning the number of operands in the use (form_reg_tuple). If the use here has 4 operands, it could still be that the load has 2, e.g.

ld1 { z0, z8 }, p0/z, [...]
ld1 { z1, z9 }, p0/z, [...]
ld1 { z2, z10 }, p0/z, [...]
ld1 { z3, z11 }, p0/z, [...]
{z0, z1, z2, z3} = form_reg_tuple {z0, z8}:0, {z1, z9}:0, {z2, z10}:0, {z3, z11}:0

The uses below assume this is about the number of operands of the Use, so it seems like it's just the name that's wrong.

The same is not true for StridedRC, which uses the wrong register class (it should decided the strided RC based on RegID instead)

Copy link
Contributor Author

@kmclaughlin-arm kmclaughlin-arm Jan 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the last commit I did change StridedRC to be based on RegID; this was to support the case where the sizes of the multi-vectors used by the loads and the pseudo are different as you've described here. I've now also renamed LdOps to UseOps.

const TargetRegisterClass *StridedRC =
LdOps == 2 ? &AArch64::ZPR2StridedRegClass
: &AArch64::ZPR4StridedRegClass;

SmallVector<MCPhysReg, 4> StridedOrder;
for (MCPhysReg Reg : Order)
if (StridedRC->contains(Reg))
StridedOrder.push_back(Reg);

int OpIdx = Use.findRegisterUseOperandIdx(VirtReg, this);
assert(OpIdx != -1 && "Expected operand index from register use.");

unsigned TupleID = MRI.getRegClass(Use.getOperand(0).getReg())->getID();
bool IsMulZPR = TupleID == AArch64::ZPR2Mul2RegClassID ||
TupleID == AArch64::ZPR4Mul4RegClassID;

const MachineOperand *AssignedRegOp = llvm::find_if(
make_range(Use.operands_begin() + 1, Use.operands_end()),
[&VRM](const MachineOperand &Op) {
return VRM->hasPhys(Op.getReg());
});

if (AssignedRegOp == Use.operands_end()) {
Copy link
Collaborator

@sdesmalen-arm sdesmalen-arm Jan 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think an example would be useful here.

// There are no registers already assigned to any of the pseudo
// operands. Look for a valid starting register for the group.
for (unsigned I = 0; I < StridedOrder.size(); ++I) {
MCPhysReg Reg = StridedOrder[I];
SmallVector<MCPhysReg> Regs;
unsigned FirstStridedReg = Reg - OpIdx + 1;
Copy link
Collaborator

@sdesmalen-arm sdesmalen-arm Jan 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would avoid doing this, because Reg - OpIdx + 1 may not be an SVE tuple register, which means that getSubReg(FirstStridedReg, AArch64::zsub0) might fail.

Example, if the first tuple register in the list would be Z0_Z1 and we're looking at the second operand in the tuple form_*tuple pseudo, i.e. OpIdx = 2, then FirstStridedReg would be X26_X27.

You can instead write this as:

unsigned SubRegIdx = Use.getOperand(OpIdx).getSubReg();
if (IsMulZPR && (getSubReg(Reg, SubRegIdx) - AArch64::Z0) % UseOps !=
                    ((unsigned)OpIdx - 1))
  continue;

Copy link
Contributor Author

@kmclaughlin-arm kmclaughlin-arm Jan 30, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, I added the is_contained later to make sure getSubReg would not fail but missed that it could also fail here.


// If the FORM_TRANSPOSE nodes use the ZPRMul classes, the starting
// register of the first load should be a multiple of 2 or 4.
unsigned FirstSubReg = getSubReg(FirstStridedReg, AArch64::zsub0);
if (IsMulZPR && (FirstSubReg - AArch64::Z0) % LdOps != 0)
continue;

for (unsigned Op = 0; Op < LdOps; ++Op) {
if (!is_contained(StridedOrder, FirstStridedReg + Op) ||
getSubReg(FirstStridedReg + Op, AArch64::zsub0) !=
FirstSubReg + Op)
break;
Regs.push_back(FirstStridedReg + Op);
}

if (Regs.size() == LdOps && all_of(Regs, [&](MCPhysReg R) {
return !Matrix->isPhysRegUsed(R);
}))
Copy link
Collaborator

@sdesmalen-arm sdesmalen-arm Jan 29, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: This could be rewritten in such a way that it doesn't need SmallVector<MCPhysReg> Regs as an intermediate step, e.g.

auto IsFreeConsecutiveRegs = [&](unsigned I) {
  // conditions
};
if (all_of(iota_range<unsigned>(0U, UseOps, /*Inclusive=*/false),
           IsFreeConsecutiveReg))
  Hints.push_back(Reg);

Hints.push_back(FirstStridedReg + OpIdx - 1);
}
} else {
// At least one operand already has a physical register assigned.
// Find the starting sub-register of this and use it to work out the
// correct strided register to suggest based on the current op index.
MCPhysReg TargetStartReg =
getSubReg(VRM->getPhys(AssignedRegOp->getReg()), AArch64::zsub0) +
(OpIdx - AssignedRegOp->getOperandNo());

for (unsigned I = 0; I < StridedOrder.size(); ++I)
if (getSubReg(StridedOrder[I], AArch64::zsub0) == TargetStartReg)
Hints.push_back(StridedOrder[I]);
}

return TargetRegisterInfo::getRegAllocationHints(VirtReg, Order, Hints, MF,
VRM);
if (!Hints.empty())
return TargetRegisterInfo::getRegAllocationHints(VirtReg, Order, Hints,
MF, VRM);
}
}

for (MachineInstr &MI : MRI.def_instructions(VirtReg)) {
Expand Down
Loading