Skip to content

Commit 50b440d

Browse files
MacDuesvkeerthy
MacDue
authored and
svkeerthy
committed
[AArch64][SME] Add common helper for expanding conditional pseudos (NFC) (#155398)
Factors out some common slightly tricky code. Hopefully makes adding new pseudos simpler.
1 parent d16bded commit 50b440d

File tree

1 file changed

+100
-77
lines changed

1 file changed

+100
-77
lines changed

‎llvm/lib/Target/AArch64/AArch64ExpandPseudoInsts.cpp‎

Lines changed: 100 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -92,9 +92,18 @@ class AArch64ExpandPseudo : public MachineFunctionPass {
9292
bool expandCALL_BTI(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI);
9393
bool expandStoreSwiftAsyncContext(MachineBasicBlock &MBB,
9494
MachineBasicBlock::iterator MBBI);
95-
MachineBasicBlock *
96-
expandCommitOrRestoreZASave(MachineBasicBlock &MBB,
97-
MachineBasicBlock::iterator MBBI);
95+
struct ConditionalBlocks {
96+
MachineBasicBlock &CondBB;
97+
MachineBasicBlock &EndBB;
98+
};
99+
ConditionalBlocks expandConditionalPseudo(MachineBasicBlock &MBB,
100+
MachineBasicBlock::iterator MBBI,
101+
DebugLoc DL,
102+
MachineInstrBuilder &Branch);
103+
MachineBasicBlock *expandRestoreZASave(MachineBasicBlock &MBB,
104+
MachineBasicBlock::iterator MBBI);
105+
MachineBasicBlock *expandCommitZASave(MachineBasicBlock &MBB,
106+
MachineBasicBlock::iterator MBBI);
98107
MachineBasicBlock *expandCondSMToggle(MachineBasicBlock &MBB,
99108
MachineBasicBlock::iterator MBBI);
100109
};
@@ -991,72 +1000,97 @@ bool AArch64ExpandPseudo::expandStoreSwiftAsyncContext(
9911000
return true;
9921001
}
9931002

994-
static constexpr unsigned ZERO_ALL_ZA_MASK = 0b11111111;
995-
996-
MachineBasicBlock *AArch64ExpandPseudo::expandCommitOrRestoreZASave(
997-
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
998-
MachineInstr &MI = *MBBI;
999-
bool IsRestoreZA = MI.getOpcode() == AArch64::RestoreZAPseudo;
1000-
assert((MI.getOpcode() == AArch64::RestoreZAPseudo ||
1001-
MI.getOpcode() == AArch64::CommitZASavePseudo) &&
1002-
"Expected ZA commit or restore");
1003+
AArch64ExpandPseudo::ConditionalBlocks
1004+
AArch64ExpandPseudo::expandConditionalPseudo(MachineBasicBlock &MBB,
1005+
MachineBasicBlock::iterator MBBI,
1006+
DebugLoc DL,
1007+
MachineInstrBuilder &Branch) {
10031008
assert((std::next(MBBI) != MBB.end() ||
1004-
MI.getParent()->successors().begin() !=
1005-
MI.getParent()->successors().end()) &&
1006-
"Unexpected unreachable in block that restores ZA");
1007-
1008-
// Compare TPIDR2_EL0 value against 0.
1009-
DebugLoc DL = MI.getDebugLoc();
1010-
MachineInstrBuilder Branch =
1011-
BuildMI(MBB, MBBI, DL,
1012-
TII->get(IsRestoreZA ? AArch64::CBZX : AArch64::CBNZX))
1013-
.add(MI.getOperand(0));
1009+
MBB.successors().begin() != MBB.successors().end()) &&
1010+
"Unexpected unreachable in block");
10141011

10151012
// Split MBB and create two new blocks:
1016-
// - MBB now contains all instructions before RestoreZAPseudo.
1017-
// - SMBB contains the [Commit|RestoreZA]Pseudo instruction only.
1018-
// - EndBB contains all instructions after [Commit|RestoreZA]Pseudo.
1013+
// - MBB now contains all instructions before the conditional pseudo.
1014+
// - CondBB contains the conditional pseudo instruction only.
1015+
// - EndBB contains all instructions after the conditional pseudo.
10191016
MachineInstr &PrevMI = *std::prev(MBBI);
1020-
MachineBasicBlock *SMBB = MBB.splitAt(PrevMI, /*UpdateLiveIns*/ true);
1021-
MachineBasicBlock *EndBB = std::next(MI.getIterator()) == SMBB->end()
1022-
? *SMBB->successors().begin()
1023-
: SMBB->splitAt(MI, /*UpdateLiveIns*/ true);
1024-
1025-
// Add the SMBB label to the CB[N]Z instruction & create a branch to EndBB.
1026-
Branch.addMBB(SMBB);
1017+
MachineBasicBlock *CondBB = MBB.splitAt(PrevMI, /*UpdateLiveIns*/ true);
1018+
MachineBasicBlock *EndBB =
1019+
std::next(MBBI) == CondBB->end()
1020+
? *CondBB->successors().begin()
1021+
: CondBB->splitAt(*MBBI, /*UpdateLiveIns*/ true);
1022+
1023+
// Add the SMBB label to the branch instruction & create a branch to EndBB.
1024+
Branch.addMBB(CondBB);
10271025
BuildMI(&MBB, DL, TII->get(AArch64::B))
10281026
.addMBB(EndBB);
10291027
MBB.addSuccessor(EndBB);
10301028

1029+
// Create branch from CondBB to EndBB. Users of this helper should insert new
1030+
// instructions at CondBB.back() -- i.e. before the branch.
1031+
BuildMI(CondBB, DL, TII->get(AArch64::B)).addMBB(EndBB);
1032+
return {*CondBB, *EndBB};
1033+
}
1034+
1035+
MachineBasicBlock *
1036+
AArch64ExpandPseudo::expandRestoreZASave(MachineBasicBlock &MBB,
1037+
MachineBasicBlock::iterator MBBI) {
1038+
MachineInstr &MI = *MBBI;
1039+
DebugLoc DL = MI.getDebugLoc();
1040+
1041+
// Compare TPIDR2_EL0 against 0. Restore ZA if TPIDR2_EL0 is zero.
1042+
MachineInstrBuilder Branch =
1043+
BuildMI(MBB, MBBI, DL, TII->get(AArch64::CBZX)).add(MI.getOperand(0));
1044+
1045+
auto [CondBB, EndBB] = expandConditionalPseudo(MBB, MBBI, DL, Branch);
10311046
// Replace the pseudo with a call (BL).
10321047
MachineInstrBuilder MIB =
1033-
BuildMI(*SMBB, SMBB->end(), DL, TII->get(AArch64::BL));
1048+
BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::BL));
10341049
// Copy operands (mainly the regmask) from the pseudo.
10351050
for (unsigned I = 2; I < MI.getNumOperands(); ++I)
10361051
MIB.add(MI.getOperand(I));
1052+
// Mark the TPIDR2 block pointer (X0) as an implicit use.
1053+
MIB.addReg(MI.getOperand(1).getReg(), RegState::Implicit);
10371054

1038-
if (IsRestoreZA) {
1039-
// Mark the TPIDR2 block pointer (X0) as an implicit use.
1040-
MIB.addReg(MI.getOperand(1).getReg(), RegState::Implicit);
1041-
} else /*CommitZA*/ {
1055+
MI.eraseFromParent();
1056+
return &EndBB;
1057+
}
1058+
1059+
static constexpr unsigned ZERO_ALL_ZA_MASK = 0b11111111;
1060+
1061+
MachineBasicBlock *
1062+
AArch64ExpandPseudo::expandCommitZASave(MachineBasicBlock &MBB,
1063+
MachineBasicBlock::iterator MBBI) {
1064+
MachineInstr &MI = *MBBI;
1065+
DebugLoc DL = MI.getDebugLoc();
1066+
1067+
// Compare TPIDR2_EL0 against 0. Commit ZA if TPIDR2_EL0 is non-zero.
1068+
MachineInstrBuilder Branch =
1069+
BuildMI(MBB, MBBI, DL, TII->get(AArch64::CBNZX)).add(MI.getOperand(0));
1070+
1071+
auto [CondBB, EndBB] = expandConditionalPseudo(MBB, MBBI, DL, Branch);
1072+
// Replace the pseudo with a call (BL).
1073+
MachineInstrBuilder MIB =
1074+
BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::BL));
1075+
// Copy operands (mainly the regmask) from the pseudo.
1076+
for (unsigned I = 2; I < MI.getNumOperands(); ++I)
1077+
MIB.add(MI.getOperand(I));
1078+
// Clear TPIDR2_EL0.
1079+
BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::MSR))
1080+
.addImm(AArch64SysReg::TPIDR2_EL0)
1081+
.addReg(AArch64::XZR);
1082+
bool ZeroZA = MI.getOperand(1).getImm() != 0;
1083+
if (ZeroZA) {
10421084
[[maybe_unused]] auto *TRI =
10431085
MBB.getParent()->getSubtarget().getRegisterInfo();
1044-
// Clear TPIDR2_EL0.
1045-
BuildMI(*SMBB, SMBB->end(), DL, TII->get(AArch64::MSR))
1046-
.addImm(AArch64SysReg::TPIDR2_EL0)
1047-
.addReg(AArch64::XZR);
1048-
bool ZeroZA = MI.getOperand(1).getImm() != 0;
1049-
if (ZeroZA) {
1050-
assert(MI.definesRegister(AArch64::ZAB0, TRI) && "should define ZA!");
1051-
BuildMI(*SMBB, SMBB->end(), DL, TII->get(AArch64::ZERO_M))
1052-
.addImm(ZERO_ALL_ZA_MASK)
1053-
.addDef(AArch64::ZAB0, RegState::ImplicitDefine);
1054-
}
1086+
assert(MI.definesRegister(AArch64::ZAB0, TRI) && "should define ZA!");
1087+
BuildMI(CondBB, CondBB.back(), DL, TII->get(AArch64::ZERO_M))
1088+
.addImm(ZERO_ALL_ZA_MASK)
1089+
.addDef(AArch64::ZAB0, RegState::ImplicitDefine);
10551090
}
10561091

1057-
BuildMI(SMBB, DL, TII->get(AArch64::B)).addMBB(EndBB);
10581092
MI.eraseFromParent();
1059-
return EndBB;
1093+
return &EndBB;
10601094
}
10611095

10621096
MachineBasicBlock *
@@ -1130,24 +1164,9 @@ AArch64ExpandPseudo::expandCondSMToggle(MachineBasicBlock &MBB,
11301164
MachineInstrBuilder Tbx =
11311165
BuildMI(MBB, MBBI, DL, TII->get(Opc)).addReg(SMReg32).addImm(0);
11321166

1133-
// Split MBB and create two new blocks:
1134-
// - MBB now contains all instructions before MSRcond_pstatesvcrImm1.
1135-
// - SMBB contains the MSRcond_pstatesvcrImm1 instruction only.
1136-
// - EndBB contains all instructions after MSRcond_pstatesvcrImm1.
1137-
MachineInstr &PrevMI = *std::prev(MBBI);
1138-
MachineBasicBlock *SMBB = MBB.splitAt(PrevMI, /*UpdateLiveIns*/ true);
1139-
MachineBasicBlock *EndBB = std::next(MI.getIterator()) == SMBB->end()
1140-
? *SMBB->successors().begin()
1141-
: SMBB->splitAt(MI, /*UpdateLiveIns*/ true);
1142-
1143-
// Add the SMBB label to the TB[N]Z instruction & create a branch to EndBB.
1144-
Tbx.addMBB(SMBB);
1145-
BuildMI(&MBB, DL, TII->get(AArch64::B))
1146-
.addMBB(EndBB);
1147-
MBB.addSuccessor(EndBB);
1148-
1167+
auto [CondBB, EndBB] = expandConditionalPseudo(MBB, MBBI, DL, Tbx);
11491168
// Create the SMSTART/SMSTOP (MSRpstatesvcrImm1) instruction in SMBB.
1150-
MachineInstrBuilder MIB = BuildMI(*SMBB, SMBB->begin(), MI.getDebugLoc(),
1169+
MachineInstrBuilder MIB = BuildMI(CondBB, CondBB.back(), MI.getDebugLoc(),
11511170
TII->get(AArch64::MSRpstatesvcrImm1));
11521171
// Copy all but the second and third operands of MSRcond_pstatesvcrImm1 (as
11531172
// these contain the CopyFromReg for the first argument and the flag to
@@ -1157,10 +1176,8 @@ AArch64ExpandPseudo::expandCondSMToggle(MachineBasicBlock &MBB,
11571176
for (unsigned i = 4; i < MI.getNumOperands(); ++i)
11581177
MIB.add(MI.getOperand(i));
11591178

1160-
BuildMI(SMBB, DL, TII->get(AArch64::B)).addMBB(EndBB);
1161-
11621179
MI.eraseFromParent();
1163-
return EndBB;
1180+
return &EndBB;
11641181
}
11651182

11661183
bool AArch64ExpandPseudo::expandMultiVecPseudo(
@@ -1674,15 +1691,21 @@ bool AArch64ExpandPseudo::expandMI(MachineBasicBlock &MBB,
16741691
return expandCALL_BTI(MBB, MBBI);
16751692
case AArch64::StoreSwiftAsyncContext:
16761693
return expandStoreSwiftAsyncContext(MBB, MBBI);
1694+
case AArch64::RestoreZAPseudo:
16771695
case AArch64::CommitZASavePseudo:
1678-
case AArch64::RestoreZAPseudo: {
1679-
auto *NewMBB = expandCommitOrRestoreZASave(MBB, MBBI);
1680-
if (NewMBB != &MBB)
1681-
NextMBBI = MBB.end(); // The NextMBBI iterator is invalidated.
1682-
return true;
1683-
}
16841696
case AArch64::MSRpstatePseudo: {
1685-
auto *NewMBB = expandCondSMToggle(MBB, MBBI);
1697+
auto *NewMBB = [&] {
1698+
switch (Opcode) {
1699+
case AArch64::RestoreZAPseudo:
1700+
return expandRestoreZASave(MBB, MBBI);
1701+
case AArch64::CommitZASavePseudo:
1702+
return expandCommitZASave(MBB, MBBI);
1703+
case AArch64::MSRpstatePseudo:
1704+
return expandCondSMToggle(MBB, MBBI);
1705+
default:
1706+
llvm_unreachable("Unexpected conditional pseudo!");
1707+
}
1708+
}();
16861709
if (NewMBB != &MBB)
16871710
NextMBBI = MBB.end(); // The NextMBBI iterator is invalidated.
16881711
return true;

0 commit comments

Comments
 (0)