Skip to content

Commit

Permalink
[AIE2P] Refactor register bank selection for finding acc/fifo bank
Browse files Browse the repository at this point in the history
  • Loading branch information
@niwinanto
niwinanto committed Feb 13, 2025
1 parent d9fcfbc commit c5b401a
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 74 deletions.
125 changes: 62 additions & 63 deletions llvm/lib/Target/AIE/aie2p/AIE2PRegisterBankInfo.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -402,7 +402,7 @@ static bool checkFifoDstSrc(const MachineInstr &MI,
// Check if FifoRegCandidate is one of the fifo operands of the intrinsic
static bool isUsedAsFifoRegInIntrinsic(const MachineRegisterInfo &MRI,
const MachineInstr &MI,
const Register &FifoRegCandidate) {
const Register FifoRegCandidate) {
switch (cast<GIntrinsic>(MI).getIntrinsicID()) {
// TODO: To be extented with more FIFO using intrinsics
case Intrinsic::aie2p_fifo_ld_fill: {
Expand Down Expand Up @@ -485,10 +485,11 @@ static bool isUsedAsFifoRegInIntrinsic(const MachineRegisterInfo &MRI,
}
/// \returns true if the specified intrinsic has an accumulator
/// vector as one of its operands.
static bool isAccIntrinsic(const MachineRegisterInfo &MRI,
const MachineInstr &MI, const Register &AccReg) {
static bool isUsedAsAccRegInIntrinsic(const MachineRegisterInfo &MRI,
const MachineInstr &MI,
const Register AccReg) {
switch (cast<GIntrinsic>(MI).getIntrinsicID()) {
// All Intrinsics with accumlator destination operand
// All Intrinsics with accumulator destination operand
case Intrinsic::aie2p_vbroadcast_zero_acc1024:
case Intrinsic::aie2p_scd_expand_ACC1024:
case Intrinsic::aie2p_scd_ACC2048:
Expand Down Expand Up @@ -538,7 +539,7 @@ static bool isAccIntrinsic(const MachineRegisterInfo &MRI,
return true;
break;
}
// All Intrinsics with first source as accumlator operand
// All Intrinsics with first source as accumulator operand
case Intrinsic::aie2p_mcd_write_acc32: {
// Operand at Idx 1 is an accumulator operand
Register SrcReg = MI.getOperand(1).getReg();
Expand Down Expand Up @@ -690,25 +691,26 @@ static bool isAccIntrinsic(const MachineRegisterInfo &MRI,
return false;
}

bool AIE2PRegisterBankInfo::usesAccReg(const MachineInstr &MI,
const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
const Register &RegOp) const {
bool AIE2PRegisterBankInfo::isUsedAsAccRegInInstr(
const MachineInstr &MI, const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI, Register Reg) {
const MachineFunction &MF = *MI.getParent()->getParent();
const RegisterBankInfo *RBI = MF.getSubtarget().getRegBankInfo();
unsigned Op = MI.getOpcode();
switch (Op) {
default:
break;
case TargetOpcode::G_INTRINSIC:
case TargetOpcode::G_INTRINSIC_W_SIDE_EFFECTS:
return isAccIntrinsic(MRI, MI, RegOp);
return isUsedAsAccRegInIntrinsic(MRI, MI, Reg);
case TargetOpcode::COPY: {
Register DstReg = MI.getOperand(0).getReg();
if (isAccReg(DstReg))
return true;
break;
}
case TargetOpcode::G_STORE: {
auto *RB = getRegBank(MI.getOperand(0).getReg(), MRI, TRI);
auto *RB = RBI->getRegBank(MI.getOperand(0).getReg(), MRI, TRI);
if (RB == &AIE2P::AccRegBank)
return true;
break;
Expand All @@ -717,28 +719,30 @@ bool AIE2PRegisterBankInfo::usesAccReg(const MachineInstr &MI,
return false;
}

// Check if the instruction has RegOp as a fifo input.
// Check if the instruction has Reg as a fifo input.
// Similar to usesAccReg for Accumulators
bool AIE2PRegisterBankInfo::hasFifoInput(const MachineInstr &MI,
const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
const Register RegOp) const {
bool AIE2PRegisterBankInfo::isUsedAsFifoRegInInstr(
const MachineInstr &MI, const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI, Register Reg) {
const MachineFunction &MF = *MI.getParent()->getParent();
auto *RI = static_cast<const AIEBaseRegisterInfo *>(
MI.getParent()->getParent()->getSubtarget().getRegisterInfo());
MF.getSubtarget().getRegisterInfo());
const RegisterBankInfo *RBI = MF.getSubtarget().getRegBankInfo();

switch (MI.getOpcode()) {
default:
break;
case TargetOpcode::G_INTRINSIC:
case TargetOpcode::G_INTRINSIC_W_SIDE_EFFECTS:
return isUsedAsFifoRegInIntrinsic(MRI, MI, RegOp);
return isUsedAsFifoRegInIntrinsic(MRI, MI, Reg);
case TargetOpcode::COPY: {
Register DstReg = MI.getOperand(0).getReg();
if (RI->isFifoPhysReg(DstReg))
return true;
break;
}
case TargetOpcode::G_STORE: {
auto *RB = getRegBank(MI.getOperand(0).getReg(), MRI, TRI);
auto *RB = RBI->getRegBank(MI.getOperand(0).getReg(), MRI, TRI);
if (RB == &AIE2P::FifoRegBank)
return true;
break;
Expand All @@ -747,18 +751,18 @@ bool AIE2PRegisterBankInfo::hasFifoInput(const MachineInstr &MI,
return false;
}

bool AIE2PRegisterBankInfo::isUseAccInsn(const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
const Register &RegOp,
unsigned Depth) const {
bool AIE2PRegisterBankInfo::registerBankCheckFunc(
RegisterUsedAsSpecificBankFcn RegisterUsedAsSpecificBank,
const MachineRegisterInfo &MRI, const TargetRegisterInfo &TRI, Register Reg,
unsigned Depth) const {
if (Depth > MaxDepthBankSearch)
return false;

auto IsCopyToVReg = [](const MachineInstr &MI) {
return (MI.isCopy() && MI.getOperand(0).getReg().isVirtual());
};

for (auto &UseMI : MRI.use_nodbg_instructions(RegOp)) {
for (auto &UseMI : MRI.use_nodbg_instructions(Reg)) {
const unsigned UseOpcode = UseMI.getOpcode();

// skip copies, bitcasts and phis
Expand All @@ -767,43 +771,31 @@ bool AIE2PRegisterBankInfo::isUseAccInsn(const MachineRegisterInfo &MRI,
Register DefReg = UseMI.getOperand(0).getReg();
if (DefReg.isPhysical())
continue;
if (isUseAccInsn(MRI, TRI, DefReg, Depth + 1))
if (registerBankCheckFunc(RegisterUsedAsSpecificBank, MRI, TRI, DefReg,
Depth + 1))
return true;
} else if (usesAccReg(UseMI, MRI, TRI, RegOp)) {
} else if (RegisterUsedAsSpecificBank(UseMI, MRI, TRI, Reg)) {
return true;
}
}
return false;
}

// Check if RegOp is used as a fifo register.
bool AIE2PRegisterBankInfo::isUseFifoInsn(const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
const Register RegOp,
unsigned Depth) const {
if (Depth > MaxDepthBankSearch)
return false;

auto IsCopyToVReg = [](const MachineInstr &MI) {
return (MI.isCopy() && MI.getOperand(0).getReg().isVirtual());
};

for (auto &UseMI : MRI.use_nodbg_instructions(RegOp)) {
const unsigned UseOpcode = UseMI.getOpcode();
const RegisterBank *AIE2PRegisterBankInfo::getPreferredRegBankForVectorTy(
const MachineRegisterInfo &MRI, const TargetRegisterInfo &TRI,
Register Reg) const {
auto RegTy = MRI.getType(Reg);
auto RegSize = RegTy.getSizeInBits();
if (RegSize == 256)
return &getRegBank(AIE2P::VRegBankID);
if (RegSize == 2048)
return &getRegBank(AIE2P::AccRegBankID);

// skip copies, bitcasts and phis
if (UseOpcode == TargetOpcode::G_BITCAST || IsCopyToVReg(UseMI) ||
UseMI.isPHI()) {
Register DefReg = UseMI.getOperand(0).getReg();
if (DefReg.isPhysical())
continue;
if (isUseFifoInsn(MRI, TRI, DefReg, Depth + 1))
return true;
} else if (hasFifoInput(UseMI, MRI, TRI, RegOp)) {
return true;
}
}
return false;
if (registerBankCheckFunc(isUsedAsAccRegInInstr, MRI, TRI, Reg))
return &getRegBank(AIE2P::AccRegBankID);
if (registerBankCheckFunc(isUsedAsFifoRegInInstr, MRI, TRI, Reg))
return &getRegBank(AIE2P::FifoRegBankID);
return &getRegBank(AIE2P::VRegBankID);
}

const RegisterBankInfo::InstructionMapping &
Expand Down Expand Up @@ -833,7 +825,8 @@ AIE2PRegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
auto *RB = getRegBank(DstReg, MRI, TRI);
if (RB == &AIE2P::AccRegBank)
return AIEBaseRegisterBankInfo::getInstrMapping(MI);
if (isUseAccInsn(MRI, TRI, DstReg)) {
if (&AIE2P::AccRegBank ==
getPreferredRegBankForVectorTy(MRI, TRI, DstReg)) {
OpRegBankIdx[0] = getAccPartialMappingIdx(DstType);
for (unsigned Idx = 2; Idx < NumOperands; ++Idx) {
LLT Type = MRI.getType(MI.getOperand(Idx).getReg());
Expand All @@ -851,7 +844,7 @@ AIE2PRegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
LLT Ty = MRI.getType(MO.getReg());
if (!Ty.isValid())
continue;
if (isAccIntrinsic(MRI, MI, MI.getOperand(Idx).getReg())) {
if (isUsedAsAccRegInIntrinsic(MRI, MI, MI.getOperand(Idx).getReg())) {
LLT Type = MRI.getType(MI.getOperand(Idx).getReg());
OpRegBankIdx[Idx] = getAccPartialMappingIdx(Type);
continue;
Expand All @@ -874,11 +867,13 @@ AIE2PRegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
auto *RB = getRegBank(DstReg, MRI, TRI);
if (RB == &AIE2P::AccRegBank)
return AIEBaseRegisterBankInfo::getInstrMapping(MI);
if (isUseAccInsn(MRI, TRI, DstReg)) {
if (&AIE2P::AccRegBank ==
getPreferredRegBankForVectorTy(MRI, TRI, DstReg)) {
OpRegBankIdx[0] = getAccPartialMappingIdx(Type);
return AIEBaseRegisterBankInfo::getInstrMappingFinal(MI, Cost, OpSize,
OpRegBankIdx);
} else if (isUseFifoInsn(MRI, TRI, DstReg)) {
} else if (&AIE2P::FifoRegBank ==
getPreferredRegBankForVectorTy(MRI, TRI, DstReg)) {
OpRegBankIdx[0] = getFifoPartialMappingIdx(Type);
return AIEBaseRegisterBankInfo::getInstrMappingFinal(MI, Cost, OpSize,
OpRegBankIdx);
Expand All @@ -893,7 +888,7 @@ AIE2PRegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
LLT SrcType = MRI.getType(SrcReg);
// Check if we already know the register bank.
auto *RB = getRegBank(SrcReg, MRI, TRI);
if (isUseAccInsn(MRI, TRI, DstReg))
if (&AIE2P::AccRegBank == getPreferredRegBankForVectorTy(MRI, TRI, DstReg))
OpRegBankIdx[0] = getAccPartialMappingIdx(DstType);
else
OpRegBankIdx[0] = getPartialMappingIdx(DstType);
Expand Down Expand Up @@ -953,7 +948,8 @@ AIE2PRegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
OpRegBankIdx);
}

if (isUseAccInsn(MRI, TRI, DstReg)) {
if (&AIE2P::AccRegBank ==
getPreferredRegBankForVectorTy(MRI, TRI, DstReg)) {
OpRegBankIdx[0] = getAccPartialMappingIdx(DstType);
for (unsigned Idx = 1; Idx < NumOperands; ++Idx) {
LLT Type = MRI.getType(MI.getOperand(Idx).getReg());
Expand All @@ -974,9 +970,11 @@ AIE2PRegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
auto *RB = getRegBank(UseCandidate, MRI, TRI);
if (RB == &AIE2P::AccRegBank || RB == &AIE2P::FifoRegBank)
return AIEBaseRegisterBankInfo::getInstrMapping(MI);
if (isUseAccInsn(MRI, TRI, UseCandidate))
const auto *PreferredRB =
getPreferredRegBankForVectorTy(MRI, TRI, UseCandidate);
if (&AIE2P::AccRegBank == PreferredRB)
isAccRegMapping = true;
if (isUseFifoInsn(MRI, TRI, UseCandidate))
if (&AIE2P::FifoRegBank == PreferredRB)
isFifoPhysRegMapping = true;
// size of accu and fifo vector on aie2p >= 512.
MachineMemOperand *MMO = *MI.memoperands_begin();
Expand All @@ -992,9 +990,10 @@ AIE2PRegisterBankInfo::getInstrMapping(const MachineInstr &MI) const {
UseCandidate = DefMI->getOperand(0).getReg();
Type = MRI.getType(MI.getOperand(0).getReg());
}
if (isUseAccInsn(MRI, TRI, UseCandidate))
PreferredRB = getPreferredRegBankForVectorTy(MRI, TRI, UseCandidate);
if (&AIE2P::AccRegBank == PreferredRB)
isAccRegMapping = true;
if (isUseFifoInsn(MRI, TRI, UseCandidate))
if (&AIE2P::FifoRegBank == PreferredRB)
isFifoPhysRegMapping = true;
if (isAccRegMapping) {
OpRegBankIdx[0] = getAccPartialMappingIdx(Type);
Expand Down
32 changes: 21 additions & 11 deletions llvm/lib/Target/AIE/aie2p/AIE2PRegisterBankInfo.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,8 @@
#define LLVM_LIB_TARGET_AIE2P_AIE2PREGISTERBANKINFO_H

#include "AIEBaseRegisterBankInfo.h"
#include "llvm/CodeGen/RegisterBank.h"
#include <optional>

#define GET_REGBANK_DECLARATIONS
#include "AIE2PGenRegisterBank.inc"
Expand Down Expand Up @@ -77,17 +79,25 @@ class AIE2PRegisterBankInfo final : public AIE2PGenRegisterBankInfo {
getInstrMapping(const MachineInstr &MI) const override;
const RegisterBank &getRegBankFromRegClass(const TargetRegisterClass &RC,
LLT) const override;
bool usesAccReg(const MachineInstr &MI, const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI, const Register &AccReg) const;
bool isUseAccInsn(const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI, const Register &AccReg,
unsigned Depth = 0) const;
bool hasFifoInput(const MachineInstr &MI, const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
const Register FifoReg) const;
bool isUseFifoInsn(const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI, const Register FifoReg,
unsigned Depth = 0) const;
static bool isUsedAsAccRegInInstr(const MachineInstr &MI,
const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
Register AccReg);
static bool isUsedAsFifoRegInInstr(const MachineInstr &MI,
const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
Register FifoReg);
using RegisterUsedAsSpecificBankFcn =
std::function<bool(const MachineInstr &MI, const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI, Register Reg)>;
bool registerBankCheckFunc(
RegisterUsedAsSpecificBankFcn RegisterUsedAsSpecificBank,
const MachineRegisterInfo &MRI, const TargetRegisterInfo &TRI,
Register Reg, unsigned Depth = 0) const;
const RegisterBank *
getPreferredRegBankForVectorTy(const MachineRegisterInfo &MRI,
const TargetRegisterInfo &TRI,
Register Reg) const;
};
} // end namespace llvm
#endif

0 comments on commit c5b401a

Please sign in to comment.