Skip to content

Commit f9c0638

Browse files
MacDueMacDue
committed
Fixups
Change-Id: Ib920d9cd4c2baac8dc278e076806504747619d8b
1 parent 441498f commit f9c0638

File tree

3 files changed

+35
-46
lines changed

3 files changed

+35
-46
lines changed

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

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -9103,21 +9103,17 @@ AArch64TargetLowering::LowerCall(CallLoweringInfo &CLI,
91039103

91049104
// Determine whether we need any streaming mode changes.
91059105
SMECallAttrs CallAttrs = getSMECallAttrs(MF.getFunction(), *this, CLI);
9106+
9107+
std::optional<unsigned> ZAMarkerNode;
91069108
bool UseNewSMEABILowering = getTM().useNewSMEABILowering();
9107-
bool IsAgnosticZAFunction = CallAttrs.caller().hasAgnosticZAInterface();
9108-
auto ZAMarkerNode = [&]() -> std::optional<unsigned> {
9109-
if (!UseNewSMEABILowering)
9110-
return std::nullopt;
9111-
if (IsAgnosticZAFunction) {
9112-
if (CallAttrs.requiresPreservingAllZAState())
9113-
return AArch64ISD::REQUIRES_ZA_SAVE;
9114-
return std::nullopt;
9115-
}
9116-
if (!CallAttrs.caller().hasZAState() && !CallAttrs.caller().hasZT0State())
9117-
return std::nullopt;
9118-
return CallAttrs.requiresLazySave() ? AArch64ISD::REQUIRES_ZA_SAVE
9119-
: AArch64ISD::INOUT_ZA_USE;
9120-
}();
9109+
if (UseNewSMEABILowering) {
9110+
if (CallAttrs.requiresLazySave() ||
9111+
CallAttrs.requiresPreservingAllZAState())
9112+
ZAMarkerNode = AArch64ISD::REQUIRES_ZA_SAVE;
9113+
else if (CallAttrs.caller().hasZAState() ||
9114+
CallAttrs.caller().hasZT0State())
9115+
ZAMarkerNode = AArch64ISD::INOUT_ZA_USE;
9116+
}
91219117

91229118
if (IsTailCall) {
91239119
// Check if it's really possible to do a tail call.

‎llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -264,7 +264,9 @@ class AArch64FunctionInfo final : public MachineFunctionInfo {
264264
EarlyAllocSMESaveBuffer = Ptr;
265265
}
266266

267-
Register getEarlyAllocSMESaveBuffer() { return EarlyAllocSMESaveBuffer; }
267+
Register getEarlyAllocSMESaveBuffer() const {
268+
return EarlyAllocSMESaveBuffer;
269+
}
268270

269271
// Old SME ABI lowering state getters/setters:
270272
Register getSMESaveBufferAddr() const { return SMESaveBufferAddr; };

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

Lines changed: 22 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -200,7 +200,7 @@ struct MachineSMEABI : public MachineFunctionPass {
200200

201201
/// Inserts code to handle changes between ZA states within the function.
202202
/// E.g., ACTIVE -> LOCAL_SAVED will insert code required to save ZA.
203-
void insertStateChanges(bool IsAgnosticZA);
203+
void insertStateChanges();
204204

205205
// Emission routines for private and shared ZA functions (using lazy saves).
206206
void emitNewZAPrologue(MachineBasicBlock &MBB,
@@ -227,26 +227,25 @@ struct MachineSMEABI : public MachineFunctionPass {
227227
LiveRegs PhysLiveRegs);
228228

229229
void emitStateChange(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
230-
ZAState From, ZAState To, LiveRegs PhysLiveRegs,
231-
bool IsAgnosticZA);
230+
ZAState From, ZAState To, LiveRegs PhysLiveRegs);
232231

233232
// Helpers for switching between lazy/full ZA save/restore routines.
234233
void emitZASave(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
235-
LiveRegs PhysLiveRegs, bool IsAgnosticZA) {
236-
if (IsAgnosticZA)
234+
LiveRegs PhysLiveRegs) {
235+
if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())
237236
return emitFullZASaveRestore(MBB, MBBI, PhysLiveRegs, /*IsSave=*/true);
238237
return emitSetupLazySave(MBB, MBBI);
239238
}
240239
void emitZARestore(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
241-
LiveRegs PhysLiveRegs, bool IsAgnosticZA) {
242-
if (IsAgnosticZA)
240+
LiveRegs PhysLiveRegs) {
241+
if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())
243242
return emitFullZASaveRestore(MBB, MBBI, PhysLiveRegs, /*IsSave=*/false);
244243
return emitRestoreLazySave(MBB, MBBI, PhysLiveRegs);
245244
}
246245
void emitAllocateZASaveBuffer(MachineBasicBlock &MBB,
247246
MachineBasicBlock::iterator MBBI,
248-
LiveRegs PhysLiveRegs, bool IsAgnosticZA) {
249-
if (IsAgnosticZA)
247+
LiveRegs PhysLiveRegs) {
248+
if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())
250249
return emitAllocateFullZASaveBuffer(MBB, MBBI, PhysLiveRegs);
251250
return emitAllocateLazySaveBuffer(MBB, MBBI);
252251
}
@@ -288,13 +287,13 @@ struct MachineSMEABI : public MachineFunctionPass {
288287
std::optional<MachineBasicBlock::iterator> AfterSMEProloguePt;
289288
Register AgnosticZABufferPtr = AArch64::NoRegister;
290289
LiveRegs PhysLiveRegsAfterSMEPrologue = LiveRegs::None;
291-
bool HasFullZASaveRestore = false;
292290
} State;
293291

294292
MachineFunction *MF = nullptr;
295293
EdgeBundles *Bundles = nullptr;
296294
const AArch64Subtarget *Subtarget = nullptr;
297295
const AArch64RegisterInfo *TRI = nullptr;
296+
const AArch64FunctionInfo *AFI = nullptr;
298297
const TargetInstrInfo *TII = nullptr;
299298
MachineRegisterInfo *MRI = nullptr;
300299
};
@@ -425,7 +424,7 @@ void MachineSMEABI::assignBundleZAStates() {
425424
}
426425
}
427426

428-
void MachineSMEABI::insertStateChanges(bool IsAgnosticZA) {
427+
void MachineSMEABI::insertStateChanges() {
429428
for (MachineBasicBlock &MBB : *MF) {
430429
const BlockInfo &Block = State.Blocks[MBB.getNumber()];
431430
ZAState InState = State.BundleStates[Bundles->getBundle(MBB.getNumber(),
@@ -438,7 +437,7 @@ void MachineSMEABI::insertStateChanges(bool IsAgnosticZA) {
438437
for (auto &Inst : Block.Insts) {
439438
if (CurrentState != Inst.NeededState)
440439
emitStateChange(MBB, Inst.InsertPt, CurrentState, Inst.NeededState,
441-
Inst.PhysLiveRegs, IsAgnosticZA);
440+
Inst.PhysLiveRegs);
442441
CurrentState = Inst.NeededState;
443442
}
444443

@@ -449,7 +448,7 @@ void MachineSMEABI::insertStateChanges(bool IsAgnosticZA) {
449448
State.BundleStates[Bundles->getBundle(MBB.getNumber(), /*Out=*/true)];
450449
if (CurrentState != OutState)
451450
emitStateChange(MBB, MBB.getFirstTerminator(), CurrentState, OutState,
452-
Block.PhysLiveRegsAtExit, IsAgnosticZA);
451+
Block.PhysLiveRegsAtExit);
453452
}
454453
}
455454

@@ -582,8 +581,6 @@ void MachineSMEABI::emitZAOff(MachineBasicBlock &MBB,
582581
void MachineSMEABI::emitAllocateLazySaveBuffer(
583582
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
584583
MachineFrameInfo &MFI = MF->getFrameInfo();
585-
auto *AFI = MF->getInfo<AArch64FunctionInfo>();
586-
587584
DebugLoc DL = getDebugLoc(MBB, MBBI);
588585
Register SP = MRI->createVirtualRegister(&AArch64::GPR64RegClass);
589586
Register SVL = MRI->createVirtualRegister(&AArch64::GPR64RegClass);
@@ -680,7 +677,6 @@ void MachineSMEABI::emitFullZASaveRestore(MachineBasicBlock &MBB,
680677
MachineBasicBlock::iterator MBBI,
681678
LiveRegs PhysLiveRegs, bool IsSave) {
682679
auto *TLI = Subtarget->getTargetLowering();
683-
State.HasFullZASaveRestore = true;
684680
DebugLoc DL = getDebugLoc(MBB, MBBI);
685681
Register BufferPtr = AArch64::X0;
686682

@@ -705,8 +701,6 @@ void MachineSMEABI::emitFullZASaveRestore(MachineBasicBlock &MBB,
705701
void MachineSMEABI::emitAllocateFullZASaveBuffer(
706702
MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
707703
LiveRegs PhysLiveRegs) {
708-
auto *AFI = MF->getInfo<AArch64FunctionInfo>();
709-
710704
// Buffer already allocated in SelectionDAG.
711705
if (AFI->getEarlyAllocSMESaveBuffer())
712706
return;
@@ -751,7 +745,7 @@ void MachineSMEABI::emitAllocateFullZASaveBuffer(
751745
void MachineSMEABI::emitStateChange(MachineBasicBlock &MBB,
752746
MachineBasicBlock::iterator InsertPt,
753747
ZAState From, ZAState To,
754-
LiveRegs PhysLiveRegs, bool IsAgnosticZA) {
748+
LiveRegs PhysLiveRegs) {
755749

756750
// ZA not used.
757751
if (From == ZAState::ANY || To == ZAState::ANY)
@@ -783,13 +777,14 @@ void MachineSMEABI::emitStateChange(MachineBasicBlock &MBB,
783777
}
784778

785779
if (From == ZAState::ACTIVE && To == ZAState::LOCAL_SAVED)
786-
emitZASave(MBB, InsertPt, PhysLiveRegs, IsAgnosticZA);
780+
emitZASave(MBB, InsertPt, PhysLiveRegs);
787781
else if (From == ZAState::LOCAL_SAVED && To == ZAState::ACTIVE)
788-
emitZARestore(MBB, InsertPt, PhysLiveRegs, IsAgnosticZA);
782+
emitZARestore(MBB, InsertPt, PhysLiveRegs);
789783
else if (To == ZAState::OFF) {
790784
assert(From != ZAState::CALLER_DORMANT &&
791785
"CALLER_DORMANT to OFF should have already been handled");
792-
assert(!IsAgnosticZA && "Should not turn ZA off in agnostic ZA function");
786+
assert(!AFI->getSMEFnAttrs().hasAgnosticZAInterface() &&
787+
"Should not turn ZA off in agnostic ZA function");
793788
emitZAOff(MBB, InsertPt, /*ClearTPIDR2=*/From == ZAState::LOCAL_SAVED);
794789
} else {
795790
dbgs() << "Error: Transition from " << getZAStateString(From) << " to "
@@ -807,7 +802,7 @@ bool MachineSMEABI::runOnMachineFunction(MachineFunction &MF) {
807802
if (!MF.getSubtarget<AArch64Subtarget>().hasSME())
808803
return false;
809804

810-
auto *AFI = MF.getInfo<AArch64FunctionInfo>();
805+
AFI = MF.getInfo<AArch64FunctionInfo>();
811806
SMEAttrs SMEFnAttrs = AFI->getSMEFnAttrs();
812807
if (!SMEFnAttrs.hasZAState() && !SMEFnAttrs.hasZT0State() &&
813808
!SMEFnAttrs.hasAgnosticZAInterface())
@@ -824,27 +819,23 @@ bool MachineSMEABI::runOnMachineFunction(MachineFunction &MF) {
824819
TRI = Subtarget->getRegisterInfo();
825820
MRI = &MF.getRegInfo();
826821

827-
bool IsAgnosticZA = SMEFnAttrs.hasAgnosticZAInterface();
828-
829822
collectNeededZAStates(SMEFnAttrs);
830823
assignBundleZAStates();
831-
insertStateChanges(/*IsAgnosticZA=*/IsAgnosticZA);
824+
insertStateChanges();
832825

833826
// Allocate save buffer (if needed).
834-
if (State.HasFullZASaveRestore || State.TPIDR2Block) {
827+
if (State.AgnosticZABufferPtr != AArch64::NoRegister || State.TPIDR2Block) {
835828
if (State.AfterSMEProloguePt) {
836829
// Note: With inline stack probes the AfterSMEProloguePt may not be in the
837830
// entry block (due to the probing loop).
838831
emitAllocateZASaveBuffer(*(*State.AfterSMEProloguePt)->getParent(),
839832
*State.AfterSMEProloguePt,
840-
State.PhysLiveRegsAfterSMEPrologue,
841-
/*IsAgnosticZA=*/IsAgnosticZA);
833+
State.PhysLiveRegsAfterSMEPrologue);
842834
} else {
843835
MachineBasicBlock &EntryBlock = MF.front();
844836
emitAllocateZASaveBuffer(
845837
EntryBlock, EntryBlock.getFirstNonPHI(),
846-
State.Blocks[EntryBlock.getNumber()].PhysLiveRegsAtEntry,
847-
/*IsAgnosticZA=*/IsAgnosticZA);
838+
State.Blocks[EntryBlock.getNumber()].PhysLiveRegsAtEntry);
848839
}
849840
}
850841

0 commit comments

Comments
 (0)