1//===- MachineSMEABIPass.cpp ----------------------------------------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// This pass implements the SME ABI requirements for ZA state. This includes
10// implementing the lazy (and agnostic) ZA state save schemes around calls.
11//
12//===----------------------------------------------------------------------===//
13//
14// This pass works by collecting instructions that require ZA to be in a
15// specific state (e.g., "ACTIVE" or "SAVED") and inserting the necessary state
16// transitions to ensure ZA is in the required state before instructions. State
17// transitions represent actions such as setting up or restoring a lazy save.
18// Certain points within a function may also have predefined states independent
19// of any instructions, for example, a "shared_za" function is always entered
20// and exited in the "ACTIVE" state.
21//
22// To handle ZA state across control flow, we make use of edge bundling. This
23// assigns each block an "incoming" and "outgoing" edge bundle (representing
24// incoming and outgoing edges). Initially, these are unique to each block;
25// then, in the process of forming bundles, the outgoing bundle of a block is
26// joined with the incoming bundle of all successors. The result is that each
27// bundle can be assigned a single ZA state, which ensures the state required by
28// all a blocks' successors is the same, and that each basic block will always
29// be entered with the same ZA state. This eliminates the need for splitting
30// edges to insert state transitions or "phi" nodes for ZA states.
31//
32// See below for a simple example of edge bundling.
33//
34// The following shows a conditionally executed basic block (BB1):
35//
36// if (cond)
37// BB1
38// BB2
39//
40// Initial Bundles Joined Bundles
41//
42// ┌──0──┐ ┌──0──┐
43// │ BB0 │ │ BB0 │
44// └──1──┘ └──1──┘
45// ├───────┐ ├───────┐
46// ▼ │ ▼ │
47// ┌──2──┐ │ ─────► ┌──1──┐ │
48// │ BB1 │ ▼ │ BB1 │ ▼
49// └──3──┘ ┌──4──┐ └──1──┘ ┌──1──┐
50// └───►4 BB2 │ └───►1 BB2 │
51// └──5──┘ └──2──┘
52//
53// On the left are the initial per-block bundles, and on the right are the
54// joined bundles (which are the result of the EdgeBundles analysis).
55
56#include "AArch64InstrInfo.h"
57#include "AArch64MachineFunctionInfo.h"
58#include "AArch64Subtarget.h"
59#include "MCTargetDesc/AArch64AddressingModes.h"
60#include "llvm/ADT/BitmaskEnum.h"
61#include "llvm/ADT/SmallVector.h"
62#include "llvm/CodeGen/EdgeBundles.h"
63#include "llvm/CodeGen/LivePhysRegs.h"
64#include "llvm/CodeGen/MachineBasicBlock.h"
65#include "llvm/CodeGen/MachineFunctionPass.h"
66#include "llvm/CodeGen/MachineOptimizationRemarkEmitter.h"
67#include "llvm/CodeGen/MachineRegisterInfo.h"
68#include "llvm/CodeGen/TargetRegisterInfo.h"
69
70using namespace llvm;
71
72#define DEBUG_TYPE "aarch64-machine-sme-abi"
73
74namespace {
75
76// Note: For agnostic ZA, we assume the function is always entered/exited in the
77// "ACTIVE" state -- this _may_ not be the case (since OFF is also a
78// possibility, but for the purpose of placing ZA saves/restores, that does not
79// matter).
80enum ZAState : uint8_t {
81 // Any/unknown state (not valid)
82 ANY = 0,
83
84 // ZA is in use and active (i.e. within the accumulator)
85 ACTIVE,
86
87 // ZA is active, but ZT0 has been saved.
88 // This handles the edge case of sharedZA && !sharesZT0.
89 ACTIVE_ZT0_SAVED,
90
91 // A ZA save has been set up or committed (i.e. ZA is dormant or off)
92 // If the function uses ZT0 it must also be saved.
93 LOCAL_SAVED,
94
95 // ZA has been committed to the lazy save buffer of the current function.
96 // If the function uses ZT0 it must also be saved.
97 // ZA is off.
98 LOCAL_COMMITTED,
99
100 // The ZA/ZT0 state on entry to the function.
101 ENTRY,
102
103 // ZA is off.
104 OFF,
105
106 // The number of ZA states (not a valid state)
107 NUM_ZA_STATE
108};
109
110/// A bitmask enum to record live physical registers that the "emit*" routines
111/// may need to preserve. Note: This only tracks registers we may clobber.
112enum LiveRegs : uint8_t {
113 None = 0,
114 NZCV = 1 << 0,
115 W0 = 1 << 1,
116 W0_HI = 1 << 2,
117 X0 = W0 | W0_HI,
118 LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ W0_HI)
119};
120
121/// Holds the virtual registers live physical registers have been saved to.
122struct PhysRegSave {
123 LiveRegs PhysLiveRegs;
124 Register StatusFlags = AArch64::NoRegister;
125 Register X0Save = AArch64::NoRegister;
126};
127
128/// Contains the needed ZA state (and live registers) at an instruction. That is
129/// the state ZA must be in _before_ "InsertPt".
130struct InstInfo {
131 ZAState NeededState{ZAState::ANY};
132 MachineBasicBlock::iterator InsertPt;
133 LiveRegs PhysLiveRegs = LiveRegs::None;
134};
135
136/// Contains the needed ZA state for each instruction in a block. Instructions
137/// that do not require a ZA state are not recorded.
138struct BlockInfo {
139 SmallVector<InstInfo> Insts;
140 ZAState FixedEntryState{ZAState::ANY};
141 ZAState DesiredIncomingState{ZAState::ANY};
142 ZAState DesiredOutgoingState{ZAState::ANY};
143 LiveRegs PhysLiveRegsAtEntry = LiveRegs::None;
144 LiveRegs PhysLiveRegsAtExit = LiveRegs::None;
145};
146
147/// Contains the needed ZA state information for all blocks within a function.
148struct FunctionInfo {
149 SmallVector<BlockInfo> Blocks;
150 std::optional<MachineBasicBlock::iterator> AfterSMEProloguePt;
151 LiveRegs PhysLiveRegsAfterSMEPrologue = LiveRegs::None;
152};
153
154/// State/helpers that is only needed when emitting code to handle
155/// saving/restoring ZA.
156class EmitContext {
157public:
158 EmitContext() = default;
159
160 /// Get or create a TPIDR2 block in \p MF.
161 int getTPIDR2Block(MachineFunction &MF) {
162 if (TPIDR2BlockFI)
163 return *TPIDR2BlockFI;
164 MachineFrameInfo &MFI = MF.getFrameInfo();
165 TPIDR2BlockFI = MFI.CreateStackObject(Size: 16, Alignment: Align(16), isSpillSlot: false);
166 return *TPIDR2BlockFI;
167 }
168
169 /// Get or create agnostic ZA buffer pointer in \p MF.
170 Register getAgnosticZABufferPtr(MachineFunction &MF) {
171 if (AgnosticZABufferPtr != AArch64::NoRegister)
172 return AgnosticZABufferPtr;
173 Register BufferPtr =
174 MF.getInfo<AArch64FunctionInfo>()->getEarlyAllocSMESaveBuffer();
175 AgnosticZABufferPtr =
176 BufferPtr != AArch64::NoRegister
177 ? BufferPtr
178 : MF.getRegInfo().createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
179 return AgnosticZABufferPtr;
180 }
181
182 int getZT0SaveSlot(MachineFunction &MF) {
183 if (ZT0SaveFI)
184 return *ZT0SaveFI;
185 MachineFrameInfo &MFI = MF.getFrameInfo();
186 ZT0SaveFI = MFI.CreateSpillStackObject(Size: 64, Alignment: Align(16));
187 return *ZT0SaveFI;
188 }
189
190 /// Returns true if the function must allocate a ZA save buffer on entry. This
191 /// will be the case if, at any point in the function, a ZA save was emitted.
192 bool needsSaveBuffer() const {
193 assert(!(TPIDR2BlockFI && AgnosticZABufferPtr) &&
194 "Cannot have both a TPIDR2 block and agnostic ZA buffer");
195 return TPIDR2BlockFI || AgnosticZABufferPtr != AArch64::NoRegister;
196 }
197
198private:
199 std::optional<int> ZT0SaveFI;
200 std::optional<int> TPIDR2BlockFI;
201 Register AgnosticZABufferPtr = AArch64::NoRegister;
202};
203
204StringRef getZAStateString(ZAState State) {
205#define MAKE_CASE(V) \
206 case V: \
207 return #V;
208 switch (State) {
209 MAKE_CASE(ZAState::ANY)
210 MAKE_CASE(ZAState::ACTIVE)
211 MAKE_CASE(ZAState::ACTIVE_ZT0_SAVED)
212 MAKE_CASE(ZAState::LOCAL_SAVED)
213 MAKE_CASE(ZAState::LOCAL_COMMITTED)
214 MAKE_CASE(ZAState::ENTRY)
215 MAKE_CASE(ZAState::OFF)
216 default:
217 llvm_unreachable("Unexpected ZAState");
218 }
219#undef MAKE_CASE
220}
221
222static bool isZAorZTRegOp(const TargetRegisterInfo &TRI,
223 const MachineOperand &MO) {
224 if (!MO.isReg() || !MO.getReg().isPhysical())
225 return false;
226 return any_of(Range: TRI.subregs_inclusive(Reg: MO.getReg()), P: [](const MCPhysReg &SR) {
227 return AArch64::MPR128RegClass.contains(Reg: SR) ||
228 AArch64::ZTRRegClass.contains(Reg: SR);
229 });
230}
231
232/// Returns the required ZA state needed before \p MI and an iterator pointing
233/// to where any code required to change the ZA state should be inserted.
234static std::pair<ZAState, MachineBasicBlock::iterator>
235getInstNeededZAState(const TargetRegisterInfo &TRI, MachineInstr &MI,
236 SMEAttrs SMEFnAttrs) {
237 MachineBasicBlock::iterator InsertPt(MI);
238
239 // Note: InOutZAUsePseudo, RequiresZASavePseudo, and RequiresZT0SavePseudo are
240 // intended to mark the position immediately before a call. Due to
241 // SelectionDAG constraints, these markers occur after the ADJCALLSTACKDOWN,
242 // so we use std::prev(InsertPt) to get the position before the call.
243
244 if (MI.getOpcode() == AArch64::InOutZAUsePseudo)
245 return {ZAState::ACTIVE, std::prev(x: InsertPt)};
246
247 // Note: If we need to save both ZA and ZT0 we use RequiresZASavePseudo.
248 if (MI.getOpcode() == AArch64::RequiresZASavePseudo)
249 return {ZAState::LOCAL_SAVED, std::prev(x: InsertPt)};
250
251 // If we only need to save ZT0 there's two cases to consider:
252 // 1. The function has ZA state (that we don't need to save).
253 // - In this case we switch to the "ACTIVE_ZT0_SAVED" state.
254 // This only saves ZT0.
255 // 2. The function does not have ZA state
256 // - In this case we switch to "LOCAL_COMMITTED" state.
257 // This saves ZT0 and turns ZA off.
258 if (MI.getOpcode() == AArch64::RequiresZT0SavePseudo) {
259 return {SMEFnAttrs.hasZAState() ? ZAState::ACTIVE_ZT0_SAVED
260 : ZAState::LOCAL_COMMITTED,
261 std::prev(x: InsertPt)};
262 }
263
264 if (MI.isReturn()) {
265 bool ZAOffAtReturn = SMEFnAttrs.hasPrivateZAInterface();
266 return {ZAOffAtReturn ? ZAState::OFF : ZAState::ACTIVE, InsertPt};
267 }
268
269 for (auto &MO : MI.operands()) {
270 if (isZAorZTRegOp(TRI, MO))
271 return {ZAState::ACTIVE, InsertPt};
272 }
273
274 return {ZAState::ANY, InsertPt};
275}
276
277struct MachineSMEABI : public MachineFunctionPass {
278 inline static char ID = 0;
279
280 MachineSMEABI(CodeGenOptLevel OptLevel = CodeGenOptLevel::Default)
281 : MachineFunctionPass(ID), OptLevel(OptLevel) {}
282
283 bool runOnMachineFunction(MachineFunction &MF) override;
284
285 StringRef getPassName() const override { return "Machine SME ABI pass"; }
286
287 void getAnalysisUsage(AnalysisUsage &AU) const override {
288 AU.setPreservesCFG();
289 AU.addRequired<EdgeBundlesWrapperLegacy>();
290 AU.addRequired<MachineOptimizationRemarkEmitterPass>();
291 AU.addRequired<LibcallLoweringInfoWrapper>();
292 AU.addPreservedID(ID&: MachineLoopInfoID);
293 AU.addPreservedID(ID&: MachineDominatorsID);
294 MachineFunctionPass::getAnalysisUsage(AU);
295 }
296
297 /// Collects the needed ZA state (and live registers) before each instruction
298 /// within the machine function.
299 FunctionInfo collectNeededZAStates(SMEAttrs SMEFnAttrs);
300
301 /// Assigns each edge bundle a ZA state based on the desired states of
302 /// incoming and outgoing blocks in the bundle.
303 SmallVector<ZAState> assignBundleZAStates(const EdgeBundles &Bundles,
304 const FunctionInfo &FnInfo);
305
306 /// Inserts code to handle changes between ZA states within the function.
307 /// E.g., ACTIVE -> LOCAL_SAVED will insert code required to save ZA.
308 void insertStateChanges(EmitContext &, const FunctionInfo &FnInfo,
309 const EdgeBundles &Bundles,
310 ArrayRef<ZAState> BundleStates);
311
312 void addSMELibCall(MachineInstrBuilder &MIB, RTLIB::Libcall LC,
313 CallingConv::ID ExpectedCC);
314
315 void emitZT0SaveRestore(EmitContext &, MachineBasicBlock &MBB,
316 MachineBasicBlock::iterator MBBI, bool IsSave);
317
318 // Emission routines for private and shared ZA functions (using lazy saves).
319 void emitSMEPrologue(MachineBasicBlock &MBB,
320 MachineBasicBlock::iterator MBBI);
321 void emitRestoreLazySave(EmitContext &, MachineBasicBlock &MBB,
322 MachineBasicBlock::iterator MBBI,
323 LiveRegs PhysLiveRegs);
324 void emitSetupLazySave(EmitContext &, MachineBasicBlock &MBB,
325 MachineBasicBlock::iterator MBBI);
326 void emitAllocateLazySaveBuffer(EmitContext &, MachineBasicBlock &MBB,
327 MachineBasicBlock::iterator MBBI);
328 void emitZAMode(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,
329 bool ClearTPIDR2, bool On);
330
331 // Emission routines for agnostic ZA functions.
332 void emitSetupFullZASave(MachineBasicBlock &MBB,
333 MachineBasicBlock::iterator MBBI,
334 LiveRegs PhysLiveRegs);
335 // Emit a "full" ZA save or restore. It is "full" in the sense that this
336 // function will emit a call to __arm_sme_save or __arm_sme_restore, which
337 // handles saving and restoring both ZA and ZT0.
338 void emitFullZASaveRestore(EmitContext &, MachineBasicBlock &MBB,
339 MachineBasicBlock::iterator MBBI,
340 LiveRegs PhysLiveRegs, bool IsSave);
341 void emitAllocateFullZASaveBuffer(EmitContext &, MachineBasicBlock &MBB,
342 MachineBasicBlock::iterator MBBI,
343 LiveRegs PhysLiveRegs);
344
345 /// Attempts to find an insertion point before \p Inst where the status flags
346 /// are not live. If \p Inst is `Block.Insts.end()` a point before the end of
347 /// the block is found.
348 std::pair<MachineBasicBlock::iterator, LiveRegs>
349 findStateChangeInsertionPoint(MachineBasicBlock &MBB, const BlockInfo &Block,
350 SmallVectorImpl<InstInfo>::const_iterator Inst);
351 void emitStateChange(EmitContext &, MachineBasicBlock &MBB,
352 MachineBasicBlock::iterator MBBI, ZAState From,
353 ZAState To, LiveRegs PhysLiveRegs);
354
355 // Helpers for switching between lazy/full ZA save/restore routines.
356 void emitZASave(EmitContext &Context, MachineBasicBlock &MBB,
357 MachineBasicBlock::iterator MBBI, LiveRegs PhysLiveRegs) {
358 if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())
359 return emitFullZASaveRestore(Context, MBB, MBBI, PhysLiveRegs,
360 /*IsSave=*/true);
361 return emitSetupLazySave(Context, MBB, MBBI);
362 }
363 void emitZARestore(EmitContext &Context, MachineBasicBlock &MBB,
364 MachineBasicBlock::iterator MBBI, LiveRegs PhysLiveRegs) {
365 if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())
366 return emitFullZASaveRestore(Context, MBB, MBBI, PhysLiveRegs,
367 /*IsSave=*/false);
368 return emitRestoreLazySave(Context, MBB, MBBI, PhysLiveRegs);
369 }
370 void emitAllocateZASaveBuffer(EmitContext &Context, MachineBasicBlock &MBB,
371 MachineBasicBlock::iterator MBBI,
372 LiveRegs PhysLiveRegs) {
373 if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())
374 return emitAllocateFullZASaveBuffer(Context, MBB, MBBI, PhysLiveRegs);
375 return emitAllocateLazySaveBuffer(Context, MBB, MBBI);
376 }
377
378 /// Collects the reachable calls from \p MBBI marked with \p Marker. This is
379 /// intended to be used to emit lazy save remarks. Note: This stops at the
380 /// first marked call along any path.
381 void collectReachableMarkedCalls(const MachineBasicBlock &MBB,
382 MachineBasicBlock::const_iterator MBBI,
383 SmallVectorImpl<const MachineInstr *> &Calls,
384 unsigned Marker) const;
385
386 void emitCallSaveRemarks(const MachineBasicBlock &MBB,
387 MachineBasicBlock::const_iterator MBBI, DebugLoc DL,
388 unsigned Marker, StringRef RemarkName,
389 StringRef SaveName) const;
390
391 void emitError(const Twine &Message) {
392 LLVMContext &Context = MF->getFunction().getContext();
393 Context.emitError(ErrorStr: MF->getName() + ": " + Message);
394 }
395
396 /// Save live physical registers to virtual registers.
397 PhysRegSave createPhysRegSave(LiveRegs PhysLiveRegs, MachineBasicBlock &MBB,
398 MachineBasicBlock::iterator MBBI, DebugLoc DL);
399 /// Restore physical registers from a save of their previous values.
400 void restorePhyRegSave(const PhysRegSave &RegSave, MachineBasicBlock &MBB,
401 MachineBasicBlock::iterator MBBI, DebugLoc DL);
402
403private:
404 CodeGenOptLevel OptLevel = CodeGenOptLevel::Default;
405
406 MachineFunction *MF = nullptr;
407 const AArch64Subtarget *Subtarget = nullptr;
408 const AArch64RegisterInfo *TRI = nullptr;
409 const AArch64FunctionInfo *AFI = nullptr;
410 const AArch64InstrInfo *TII = nullptr;
411 const LibcallLoweringInfo *LLI = nullptr;
412
413 MachineOptimizationRemarkEmitter *ORE = nullptr;
414 MachineRegisterInfo *MRI = nullptr;
415 MachineLoopInfo *MLI = nullptr;
416};
417
418static LiveRegs getPhysLiveRegs(LiveRegUnits const &LiveUnits) {
419 LiveRegs PhysLiveRegs = LiveRegs::None;
420 if (!LiveUnits.available(Reg: AArch64::NZCV))
421 PhysLiveRegs |= LiveRegs::NZCV;
422 // We have to track W0 and X0 separately as otherwise things can get
423 // confused if we attempt to preserve X0 but only W0 was defined.
424 if (!LiveUnits.available(Reg: AArch64::W0))
425 PhysLiveRegs |= LiveRegs::W0;
426 if (!LiveUnits.available(Reg: AArch64::W0_HI))
427 PhysLiveRegs |= LiveRegs::W0_HI;
428 return PhysLiveRegs;
429}
430
431static void setPhysLiveRegs(LiveRegUnits &LiveUnits, LiveRegs PhysLiveRegs) {
432 if (PhysLiveRegs & LiveRegs::NZCV)
433 LiveUnits.addReg(Reg: AArch64::NZCV);
434 if (PhysLiveRegs & LiveRegs::W0)
435 LiveUnits.addReg(Reg: AArch64::W0);
436 if (PhysLiveRegs & LiveRegs::W0_HI)
437 LiveUnits.addReg(Reg: AArch64::W0_HI);
438}
439
440[[maybe_unused]] bool isCallStartOpcode(unsigned Opc) {
441 switch (Opc) {
442 case AArch64::BLR:
443 case AArch64::BLRA:
444 case AArch64::TLSDESC_CALLSEQ:
445 case AArch64::TLSDESC_AUTH_CALLSEQ:
446 case AArch64::ADJCALLSTACKDOWN:
447 return true;
448 default:
449 return false;
450 }
451}
452
453FunctionInfo MachineSMEABI::collectNeededZAStates(SMEAttrs SMEFnAttrs) {
454 assert((SMEFnAttrs.hasAgnosticZAInterface() || SMEFnAttrs.hasZT0State() ||
455 SMEFnAttrs.hasZAState()) &&
456 "Expected function to have ZA/ZT0 state!");
457
458 SmallVector<BlockInfo> Blocks(MF->getNumBlockIDs());
459 LiveRegs PhysLiveRegsAfterSMEPrologue = LiveRegs::None;
460 std::optional<MachineBasicBlock::iterator> AfterSMEProloguePt;
461
462 for (MachineBasicBlock &MBB : *MF) {
463 BlockInfo &Block = Blocks[MBB.getNumber()];
464
465 if (MBB.isEntryBlock()) {
466 // Entry block:
467 Block.FixedEntryState = ZAState::ENTRY;
468 } else if (MBB.isEHPad()) {
469 // EH entry block:
470 Block.FixedEntryState = ZAState::LOCAL_COMMITTED;
471 }
472
473 LiveRegUnits LiveUnits(*TRI);
474 LiveUnits.addLiveOuts(MBB);
475
476 Block.PhysLiveRegsAtExit = getPhysLiveRegs(LiveUnits);
477 auto FirstTerminatorInsertPt = MBB.getFirstTerminator();
478 auto FirstNonPhiInsertPt = MBB.getFirstNonPHI();
479 for (MachineInstr &MI : reverse(C&: MBB)) {
480 if (MI.isDebugInstr())
481 continue;
482
483 MachineBasicBlock::iterator MBBI(MI);
484 LiveUnits.stepBackward(MI);
485 LiveRegs PhysLiveRegs = getPhysLiveRegs(LiveUnits);
486 // The SMEStateAllocPseudo marker is added to a function if the save
487 // buffer was allocated in SelectionDAG. It marks the end of the
488 // allocation -- which is a safe point for this pass to insert any TPIDR2
489 // block setup.
490 if (MI.getOpcode() == AArch64::SMEStateAllocPseudo) {
491 AfterSMEProloguePt = MBBI;
492 PhysLiveRegsAfterSMEPrologue = PhysLiveRegs;
493 }
494 // Note: We treat Agnostic ZA as inout_za with an alternate save/restore.
495 auto [NeededState, InsertPt] = getInstNeededZAState(TRI: *TRI, MI, SMEFnAttrs);
496 assert((InsertPt == MBBI || isCallStartOpcode(InsertPt->getOpcode())) &&
497 "Unexpected state change insertion point!");
498 if (MBBI == FirstTerminatorInsertPt)
499 Block.PhysLiveRegsAtExit = PhysLiveRegs;
500 if (MBBI == FirstNonPhiInsertPt)
501 Block.PhysLiveRegsAtEntry = PhysLiveRegs;
502 if (NeededState != ZAState::ANY)
503 Block.Insts.push_back(Elt: {.NeededState: NeededState, .InsertPt: InsertPt, .PhysLiveRegs: PhysLiveRegs});
504 }
505
506 // Reverse vector (as we had to iterate backwards for liveness).
507 std::reverse(first: Block.Insts.begin(), last: Block.Insts.end());
508
509 // Record the desired states on entry/exit of this block. These are the
510 // states that would not incur a state transition.
511 if (!Block.Insts.empty()) {
512 Block.DesiredIncomingState = Block.Insts.front().NeededState;
513 Block.DesiredOutgoingState = Block.Insts.back().NeededState;
514 }
515 }
516
517 return FunctionInfo{.Blocks: std::move(Blocks), .AfterSMEProloguePt: AfterSMEProloguePt,
518 .PhysLiveRegsAfterSMEPrologue: PhysLiveRegsAfterSMEPrologue};
519}
520
521/// Assigns each edge bundle a ZA state based on the desired states of incoming
522/// and outgoing blocks in the bundle.
523SmallVector<ZAState>
524MachineSMEABI::assignBundleZAStates(const EdgeBundles &Bundles,
525 const FunctionInfo &FnInfo) {
526 SmallVector<ZAState> BundleStates(Bundles.getNumBundles());
527 for (unsigned I = 0, E = Bundles.getNumBundles(); I != E; ++I) {
528 std::optional<ZAState> BundleState;
529 for (unsigned BlockID : Bundles.getBlocks(Bundle: I)) {
530 const BlockInfo &Block = FnInfo.Blocks[BlockID];
531 // Check if the block is an incoming block in the bundle. Note: We skip
532 // Block.FixedEntryState != ANY to ignore EH pads (which are only
533 // reachable via exceptions).
534 if (Block.FixedEntryState != ZAState::ANY ||
535 Bundles.getBundle(N: BlockID, /*Out=*/false) != I)
536 continue;
537
538 // Pick a state that matches all incoming blocks. Fall back to "ACTIVE" if
539 // any incoming state doesn't match. This will hoist the state from
540 // incoming blocks to outgoing blocks.
541 if (!BundleState)
542 BundleState = Block.DesiredIncomingState;
543 else if (BundleState != Block.DesiredIncomingState)
544 BundleState = ZAState::ACTIVE;
545 }
546
547 if (!BundleState || BundleState == ZAState::ANY)
548 BundleState = ZAState::ACTIVE;
549
550 BundleStates[I] = *BundleState;
551 }
552
553 return BundleStates;
554}
555
556std::pair<MachineBasicBlock::iterator, LiveRegs>
557MachineSMEABI::findStateChangeInsertionPoint(
558 MachineBasicBlock &MBB, const BlockInfo &Block,
559 SmallVectorImpl<InstInfo>::const_iterator Inst) {
560 LiveRegs PhysLiveRegs;
561 MachineBasicBlock::iterator InsertPt;
562 if (Inst != Block.Insts.end()) {
563 InsertPt = Inst->InsertPt;
564 PhysLiveRegs = Inst->PhysLiveRegs;
565 } else {
566 InsertPt = MBB.getFirstTerminator();
567 PhysLiveRegs = Block.PhysLiveRegsAtExit;
568 }
569
570 if (PhysLiveRegs == LiveRegs::None)
571 return {InsertPt, PhysLiveRegs}; // Nothing to do (no live regs).
572
573 // Find the previous state change. We can not move before this point.
574 MachineBasicBlock::iterator PrevStateChangeI;
575 if (Inst == Block.Insts.begin()) {
576 PrevStateChangeI = MBB.begin();
577 } else {
578 // Note: `std::prev(Inst)` is the previous InstInfo. We only create an
579 // InstInfo object for instructions that require a specific ZA state, so the
580 // InstInfo is the site of the previous state change in the block (which can
581 // be several MIs earlier).
582 PrevStateChangeI = std::prev(x: Inst)->InsertPt;
583 }
584
585 // Note: LiveUnits will only accurately track X0 and NZCV.
586 LiveRegUnits LiveUnits(*TRI);
587 setPhysLiveRegs(LiveUnits, PhysLiveRegs);
588 auto BestCandidate = std::make_pair(x&: InsertPt, y&: PhysLiveRegs);
589 for (MachineBasicBlock::iterator I = InsertPt; I != PrevStateChangeI; --I) {
590 if (I->isDebugInstr())
591 continue;
592
593 // Don't move before/into a call (which may have a state change before it).
594 if (I->getOpcode() == TII->getCallFrameDestroyOpcode() || I->isCall())
595 break;
596 LiveUnits.stepBackward(MI: *I);
597 LiveRegs CurrentPhysLiveRegs = getPhysLiveRegs(LiveUnits);
598 // Find places where NZCV is available, but keep looking for locations where
599 // both NZCV and X0 are available, which can avoid some copies.
600 if (!(CurrentPhysLiveRegs & LiveRegs::NZCV))
601 BestCandidate = {I, CurrentPhysLiveRegs};
602 if (CurrentPhysLiveRegs == LiveRegs::None)
603 break;
604 }
605 return BestCandidate;
606}
607
608void MachineSMEABI::insertStateChanges(EmitContext &Context,
609 const FunctionInfo &FnInfo,
610 const EdgeBundles &Bundles,
611 ArrayRef<ZAState> BundleStates) {
612 for (MachineBasicBlock &MBB : *MF) {
613 const BlockInfo &Block = FnInfo.Blocks[MBB.getNumber()];
614 ZAState InState = BundleStates[Bundles.getBundle(N: MBB.getNumber(),
615 /*Out=*/false)];
616
617 ZAState CurrentState = Block.FixedEntryState;
618 if (CurrentState == ZAState::ANY)
619 CurrentState = InState;
620
621 for (auto &Inst : Block.Insts) {
622 if (CurrentState != Inst.NeededState) {
623 auto [InsertPt, PhysLiveRegs] =
624 findStateChangeInsertionPoint(MBB, Block, Inst: &Inst);
625 emitStateChange(Context, MBB, MBBI: InsertPt, From: CurrentState, To: Inst.NeededState,
626 PhysLiveRegs);
627 CurrentState = Inst.NeededState;
628 }
629 }
630
631 if (MBB.succ_empty())
632 continue;
633
634 ZAState OutState =
635 BundleStates[Bundles.getBundle(N: MBB.getNumber(), /*Out=*/true)];
636 if (CurrentState != OutState) {
637 auto [InsertPt, PhysLiveRegs] =
638 findStateChangeInsertionPoint(MBB, Block, Inst: Block.Insts.end());
639 emitStateChange(Context, MBB, MBBI: InsertPt, From: CurrentState, To: OutState,
640 PhysLiveRegs);
641 }
642 }
643}
644
645static DebugLoc getDebugLoc(MachineBasicBlock &MBB,
646 MachineBasicBlock::iterator MBBI) {
647 if (MBB.empty())
648 return DebugLoc();
649 return MBBI != MBB.end() ? MBBI->getDebugLoc() : MBB.back().getDebugLoc();
650}
651
652/// Finds the first call (as determined by MachineInstr::isCall()) starting from
653/// \p MBBI in \p MBB marked with \p Marker (which is a marker opcode such as
654/// RequiresZASavePseudo). If a marked call is found, it is pushed to \p Calls
655/// and the function returns true.
656static bool findMarkedCall(const MachineBasicBlock &MBB,
657 MachineBasicBlock::const_iterator MBBI,
658 SmallVectorImpl<const MachineInstr *> &Calls,
659 unsigned Marker, unsigned CallDestroyOpcode) {
660 auto IsMarker = [&](auto &MI) { return MI.getOpcode() == Marker; };
661 auto MarkerInst = std::find_if(first: MBBI, last: MBB.end(), pred: IsMarker);
662 if (MarkerInst == MBB.end())
663 return false;
664 MachineBasicBlock::const_iterator I = MarkerInst;
665 while (++I != MBB.end()) {
666 if (I->isCall() || I->getOpcode() == CallDestroyOpcode)
667 break;
668 }
669 if (I != MBB.end() && I->isCall())
670 Calls.push_back(Elt: &*I);
671 // Note: This function always returns true if a "Marker" was found.
672 return true;
673}
674
675void MachineSMEABI::collectReachableMarkedCalls(
676 const MachineBasicBlock &StartMBB,
677 MachineBasicBlock::const_iterator StartInst,
678 SmallVectorImpl<const MachineInstr *> &Calls, unsigned Marker) const {
679 assert(Marker == AArch64::InOutZAUsePseudo ||
680 Marker == AArch64::RequiresZASavePseudo ||
681 Marker == AArch64::RequiresZT0SavePseudo);
682 unsigned CallDestroyOpcode = TII->getCallFrameDestroyOpcode();
683 if (findMarkedCall(MBB: StartMBB, MBBI: StartInst, Calls, Marker, CallDestroyOpcode))
684 return;
685
686 SmallPtrSet<const MachineBasicBlock *, 4> Visited;
687 SmallVector<const MachineBasicBlock *> Worklist(StartMBB.succ_rbegin(),
688 StartMBB.succ_rend());
689 while (!Worklist.empty()) {
690 const MachineBasicBlock *MBB = Worklist.pop_back_val();
691 auto [_, Inserted] = Visited.insert(Ptr: MBB);
692 if (!Inserted)
693 continue;
694
695 if (!findMarkedCall(MBB: *MBB, MBBI: MBB->begin(), Calls, Marker, CallDestroyOpcode))
696 Worklist.append(in_start: MBB->succ_rbegin(), in_end: MBB->succ_rend());
697 }
698}
699
700static StringRef getCalleeName(const MachineInstr &CallInst) {
701 assert(CallInst.isCall() && "expected a call");
702 for (const MachineOperand &MO : CallInst.operands()) {
703 if (MO.isSymbol())
704 return MO.getSymbolName();
705 if (MO.isGlobal())
706 return MO.getGlobal()->getName();
707 }
708 return {};
709}
710
711void MachineSMEABI::emitCallSaveRemarks(const MachineBasicBlock &MBB,
712 MachineBasicBlock::const_iterator MBBI,
713 DebugLoc DL, unsigned Marker,
714 StringRef RemarkName,
715 StringRef SaveName) const {
716 auto SaveRemark = [&](DebugLoc DL, const MachineBasicBlock &MBB) {
717 return MachineOptimizationRemarkAnalysis("sme", RemarkName, DL, &MBB);
718 };
719 StringRef StateName = Marker == AArch64::RequiresZT0SavePseudo ? "ZT0" : "ZA";
720 ORE->emit(RemarkBuilder: [&] {
721 return SaveRemark(DL, MBB) << SaveName << " of " << StateName
722 << " emitted in '" << MF->getName() << "'";
723 });
724 if (!ORE->allowExtraAnalysis(PassName: "sme"))
725 return;
726 SmallVector<const MachineInstr *> CallsRequiringSaves;
727 collectReachableMarkedCalls(StartMBB: MBB, StartInst: MBBI, Calls&: CallsRequiringSaves, Marker);
728 for (const MachineInstr *CallInst : CallsRequiringSaves) {
729 auto R = SaveRemark(CallInst->getDebugLoc(), *CallInst->getParent());
730 R << "call";
731 if (StringRef CalleeName = getCalleeName(CallInst: *CallInst); !CalleeName.empty())
732 R << " to '" << CalleeName << "'";
733 R << " requires " << StateName << " save";
734 ORE->emit(OptDiag&: R);
735 }
736}
737
738void MachineSMEABI::emitSetupLazySave(EmitContext &Context,
739 MachineBasicBlock &MBB,
740 MachineBasicBlock::iterator MBBI) {
741 DebugLoc DL = getDebugLoc(MBB, MBBI);
742
743 emitCallSaveRemarks(MBB, MBBI, DL, Marker: AArch64::RequiresZASavePseudo,
744 RemarkName: "SMELazySaveZA", SaveName: "lazy save");
745
746 // Get pointer to TPIDR2 block.
747 Register TPIDR2 = MRI->createVirtualRegister(RegClass: &AArch64::GPR64spRegClass);
748 Register TPIDR2Ptr = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
749 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::ADDXri), DestReg: TPIDR2)
750 .addFrameIndex(Idx: Context.getTPIDR2Block(MF&: *MF))
751 .addImm(Val: 0)
752 .addImm(Val: 0);
753 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: TPIDR2Ptr)
754 .addReg(RegNo: TPIDR2);
755 // Set TPIDR2_EL0 to point to TPIDR2 block.
756 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSR))
757 .addImm(Val: AArch64SysReg::TPIDR2_EL0)
758 .addReg(RegNo: TPIDR2Ptr);
759}
760
761PhysRegSave MachineSMEABI::createPhysRegSave(LiveRegs PhysLiveRegs,
762 MachineBasicBlock &MBB,
763 MachineBasicBlock::iterator MBBI,
764 DebugLoc DL) {
765 PhysRegSave RegSave{.PhysLiveRegs: PhysLiveRegs};
766 if (PhysLiveRegs & LiveRegs::NZCV) {
767 RegSave.StatusFlags = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
768 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MRS), DestReg: RegSave.StatusFlags)
769 .addImm(Val: AArch64SysReg::NZCV)
770 .addReg(RegNo: AArch64::NZCV, Flags: RegState::Implicit);
771 }
772 // Note: Preserving X0 is "free" as this is before register allocation, so
773 // the register allocator is still able to optimize these copies.
774 if (PhysLiveRegs & LiveRegs::W0) {
775 RegSave.X0Save = MRI->createVirtualRegister(RegClass: PhysLiveRegs & LiveRegs::W0_HI
776 ? &AArch64::GPR64RegClass
777 : &AArch64::GPR32RegClass);
778 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: RegSave.X0Save)
779 .addReg(RegNo: PhysLiveRegs & LiveRegs::W0_HI ? AArch64::X0 : AArch64::W0);
780 }
781 return RegSave;
782}
783
784void MachineSMEABI::restorePhyRegSave(const PhysRegSave &RegSave,
785 MachineBasicBlock &MBB,
786 MachineBasicBlock::iterator MBBI,
787 DebugLoc DL) {
788 if (RegSave.StatusFlags != AArch64::NoRegister)
789 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSR))
790 .addImm(Val: AArch64SysReg::NZCV)
791 .addReg(RegNo: RegSave.StatusFlags)
792 .addReg(RegNo: AArch64::NZCV, Flags: RegState::ImplicitDefine);
793
794 if (RegSave.X0Save != AArch64::NoRegister)
795 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY),
796 DestReg: RegSave.PhysLiveRegs & LiveRegs::W0_HI ? AArch64::X0 : AArch64::W0)
797 .addReg(RegNo: RegSave.X0Save);
798}
799
800void MachineSMEABI::addSMELibCall(MachineInstrBuilder &MIB, RTLIB::Libcall LC,
801 CallingConv::ID ExpectedCC) {
802 RTLIB::LibcallImpl LCImpl = LLI->getLibcallImpl(Call: LC);
803 if (LCImpl == RTLIB::Unsupported)
804 emitError(Message: "cannot lower SME ABI (SME routines unsupported)");
805 CallingConv::ID CC = LLI->getLibcallImplCallingConv(Call: LCImpl);
806 StringRef ImplName = RTLIB::RuntimeLibcallsInfo::getLibcallImplName(CallImpl: LCImpl);
807 if (CC != ExpectedCC)
808 emitError(Message: "invalid calling convention for SME routine: '" + ImplName + "'");
809 // FIXME: This assumes the ImplName StringRef is null-terminated.
810 MIB.addExternalSymbol(FnName: ImplName.data());
811 MIB.addRegMask(Mask: TRI->getCallPreservedMask(MF: *MF, CC));
812}
813
814void MachineSMEABI::emitRestoreLazySave(EmitContext &Context,
815 MachineBasicBlock &MBB,
816 MachineBasicBlock::iterator MBBI,
817 LiveRegs PhysLiveRegs) {
818 DebugLoc DL = getDebugLoc(MBB, MBBI);
819 Register TPIDR2EL0 = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
820 Register TPIDR2 = AArch64::X0;
821
822 // TODO: Emit these within the restore MBB to prevent unnecessary saves.
823 PhysRegSave RegSave = createPhysRegSave(PhysLiveRegs, MBB, MBBI, DL);
824
825 // Enable ZA.
826 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSRpstatesvcrImm1))
827 .addImm(Val: AArch64SVCR::SVCRZA)
828 .addImm(Val: 1);
829 // Get current TPIDR2_EL0.
830 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MRS), DestReg: TPIDR2EL0)
831 .addImm(Val: AArch64SysReg::TPIDR2_EL0);
832 // Get pointer to TPIDR2 block.
833 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::ADDXri), DestReg: TPIDR2)
834 .addFrameIndex(Idx: Context.getTPIDR2Block(MF&: *MF))
835 .addImm(Val: 0)
836 .addImm(Val: 0);
837 // (Conditionally) restore ZA state.
838 auto RestoreZA = BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::RestoreZAPseudo))
839 .addReg(RegNo: TPIDR2EL0)
840 .addReg(RegNo: TPIDR2);
841 addSMELibCall(
842 MIB&: RestoreZA, LC: RTLIB::SMEABI_TPIDR2_RESTORE,
843 ExpectedCC: CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0);
844 // Zero TPIDR2_EL0.
845 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSR))
846 .addImm(Val: AArch64SysReg::TPIDR2_EL0)
847 .addReg(RegNo: AArch64::XZR);
848
849 restorePhyRegSave(RegSave, MBB, MBBI, DL);
850}
851
852void MachineSMEABI::emitZAMode(MachineBasicBlock &MBB,
853 MachineBasicBlock::iterator MBBI,
854 bool ClearTPIDR2, bool On) {
855 DebugLoc DL = getDebugLoc(MBB, MBBI);
856
857 if (ClearTPIDR2)
858 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSR))
859 .addImm(Val: AArch64SysReg::TPIDR2_EL0)
860 .addReg(RegNo: AArch64::XZR);
861
862 // Disable ZA.
863 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSRpstatesvcrImm1))
864 .addImm(Val: AArch64SVCR::SVCRZA)
865 .addImm(Val: On ? 1 : 0);
866}
867
868void MachineSMEABI::emitAllocateLazySaveBuffer(
869 EmitContext &Context, MachineBasicBlock &MBB,
870 MachineBasicBlock::iterator MBBI) {
871 MachineFrameInfo &MFI = MF->getFrameInfo();
872 DebugLoc DL = getDebugLoc(MBB, MBBI);
873 Register SP = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
874 Register SVL = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
875 Register Buffer = AFI->getEarlyAllocSMESaveBuffer();
876
877 // Calculate SVL.
878 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::RDSVLI_XI), DestReg: SVL).addImm(Val: 1);
879
880 // 1. Allocate the lazy save buffer.
881 if (Buffer == AArch64::NoRegister) {
882 // TODO: On Windows, we allocate the lazy save buffer in SelectionDAG (so
883 // Buffer != AArch64::NoRegister). This is done to reuse the existing
884 // expansions (which can insert stack checks). This works, but it means we
885 // will always allocate the lazy save buffer (even if the function contains
886 // no lazy saves). If we want to handle Windows here, we'll need to
887 // implement something similar to LowerWindowsDYNAMIC_STACKALLOC.
888 assert(!Subtarget->isTargetWindows() &&
889 "Lazy ZA save is not yet supported on Windows");
890 Buffer = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
891 // Get original stack pointer.
892 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: SP)
893 .addReg(RegNo: AArch64::SP);
894 // Allocate a lazy-save buffer object of the size given, normally SVL * SVL
895 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSUBXrrr), DestReg: Buffer)
896 .addReg(RegNo: SVL)
897 .addReg(RegNo: SVL)
898 .addReg(RegNo: SP);
899 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: AArch64::SP)
900 .addReg(RegNo: Buffer);
901 // We have just allocated a variable sized object, tell this to PEI.
902 MFI.CreateVariableSizedObject(Alignment: Align(16), Alloca: nullptr);
903 }
904
905 // 2. Setup the TPIDR2 block.
906 {
907 // Note: This case just needs to do `SVL << 48`. It is not implemented as we
908 // generally don't support big-endian SVE/SME.
909 if (!Subtarget->isLittleEndian())
910 reportFatalInternalError(
911 reason: "TPIDR2 block initialization is not supported on big-endian targets");
912
913 // Store buffer pointer and num_za_save_slices.
914 // Bytes 10-15 are implicitly zeroed.
915 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::STPXi))
916 .addReg(RegNo: Buffer)
917 .addReg(RegNo: SVL)
918 .addFrameIndex(Idx: Context.getTPIDR2Block(MF&: *MF))
919 .addImm(Val: 0);
920 }
921}
922
923static constexpr unsigned ZERO_ALL_ZA_MASK = 0b11111111;
924
925void MachineSMEABI::emitSMEPrologue(MachineBasicBlock &MBB,
926 MachineBasicBlock::iterator MBBI) {
927 DebugLoc DL = getDebugLoc(MBB, MBBI);
928
929 bool ZeroZA = AFI->getSMEFnAttrs().isNewZA();
930 bool ZeroZT0 = AFI->getSMEFnAttrs().isNewZT0();
931 if (AFI->getSMEFnAttrs().hasPrivateZAInterface()) {
932 // Get current TPIDR2_EL0.
933 Register TPIDR2EL0 = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
934 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MRS))
935 .addReg(RegNo: TPIDR2EL0, Flags: RegState::Define)
936 .addImm(Val: AArch64SysReg::TPIDR2_EL0);
937 // If TPIDR2_EL0 is non-zero, commit the lazy save.
938 // NOTE: Functions that only use ZT0 don't need to zero ZA.
939 auto CommitZASave =
940 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::CommitZASavePseudo))
941 .addReg(RegNo: TPIDR2EL0)
942 .addImm(Val: ZeroZA)
943 .addImm(Val: ZeroZT0);
944 addSMELibCall(
945 MIB&: CommitZASave, LC: RTLIB::SMEABI_TPIDR2_SAVE,
946 ExpectedCC: CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X0);
947 if (ZeroZA)
948 CommitZASave.addDef(RegNo: AArch64::ZAB0, Flags: RegState::ImplicitDefine);
949 if (ZeroZT0)
950 CommitZASave.addDef(RegNo: AArch64::ZT0, Flags: RegState::ImplicitDefine);
951 // Enable ZA (as ZA could have previously been in the OFF state).
952 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::MSRpstatesvcrImm1))
953 .addImm(Val: AArch64SVCR::SVCRZA)
954 .addImm(Val: 1);
955 } else if (AFI->getSMEFnAttrs().hasSharedZAInterface()) {
956 if (ZeroZA)
957 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::ZERO_M))
958 .addImm(Val: ZERO_ALL_ZA_MASK)
959 .addDef(RegNo: AArch64::ZAB0, Flags: RegState::ImplicitDefine);
960 if (ZeroZT0)
961 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::ZERO_T)).addDef(RegNo: AArch64::ZT0);
962 }
963}
964
965void MachineSMEABI::emitFullZASaveRestore(EmitContext &Context,
966 MachineBasicBlock &MBB,
967 MachineBasicBlock::iterator MBBI,
968 LiveRegs PhysLiveRegs, bool IsSave) {
969 DebugLoc DL = getDebugLoc(MBB, MBBI);
970
971 if (IsSave)
972 emitCallSaveRemarks(MBB, MBBI, DL, Marker: AArch64::RequiresZASavePseudo,
973 RemarkName: "SMEFullZASave", SaveName: "full save");
974
975 PhysRegSave RegSave = createPhysRegSave(PhysLiveRegs, MBB, MBBI, DL);
976
977 // Copy the buffer pointer into X0.
978 Register BufferPtr = AArch64::X0;
979 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: BufferPtr)
980 .addReg(RegNo: Context.getAgnosticZABufferPtr(MF&: *MF));
981
982 // Call __arm_sme_save/__arm_sme_restore.
983 auto SaveRestoreZA = BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::BL))
984 .addReg(RegNo: BufferPtr, Flags: RegState::Implicit);
985 addSMELibCall(
986 MIB&: SaveRestoreZA,
987 LC: IsSave ? RTLIB::SMEABI_SME_SAVE : RTLIB::SMEABI_SME_RESTORE,
988 ExpectedCC: CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X1);
989
990 restorePhyRegSave(RegSave, MBB, MBBI, DL);
991}
992
993void MachineSMEABI::emitZT0SaveRestore(EmitContext &Context,
994 MachineBasicBlock &MBB,
995 MachineBasicBlock::iterator MBBI,
996 bool IsSave) {
997 DebugLoc DL = getDebugLoc(MBB, MBBI);
998
999 // Note: This will report calls that _only_ need ZT0 saved. Call that save
1000 // both ZA and ZT0 will be under the SMELazySaveZA remark. This prevents
1001 // reporting the same calls twice.
1002 if (IsSave)
1003 emitCallSaveRemarks(MBB, MBBI, DL, Marker: AArch64::RequiresZT0SavePseudo,
1004 RemarkName: "SMEZT0Save", SaveName: "spill");
1005
1006 Register ZT0Save = MRI->createVirtualRegister(RegClass: &AArch64::GPR64spRegClass);
1007
1008 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::ADDXri), DestReg: ZT0Save)
1009 .addFrameIndex(Idx: Context.getZT0SaveSlot(MF&: *MF))
1010 .addImm(Val: 0)
1011 .addImm(Val: 0);
1012
1013 if (IsSave) {
1014 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::STR_TX))
1015 .addReg(RegNo: AArch64::ZT0)
1016 .addReg(RegNo: ZT0Save);
1017 } else {
1018 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::LDR_TX), DestReg: AArch64::ZT0)
1019 .addReg(RegNo: ZT0Save);
1020 }
1021}
1022
1023void MachineSMEABI::emitAllocateFullZASaveBuffer(
1024 EmitContext &Context, MachineBasicBlock &MBB,
1025 MachineBasicBlock::iterator MBBI, LiveRegs PhysLiveRegs) {
1026 // Buffer already allocated in SelectionDAG.
1027 if (AFI->getEarlyAllocSMESaveBuffer())
1028 return;
1029
1030 DebugLoc DL = getDebugLoc(MBB, MBBI);
1031 Register BufferPtr = Context.getAgnosticZABufferPtr(MF&: *MF);
1032 Register BufferSize = MRI->createVirtualRegister(RegClass: &AArch64::GPR64RegClass);
1033
1034 PhysRegSave RegSave = createPhysRegSave(PhysLiveRegs, MBB, MBBI, DL);
1035
1036 // Calculate the SME state size.
1037 {
1038 auto SMEStateSize = BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::BL))
1039 .addReg(RegNo: AArch64::X0, Flags: RegState::ImplicitDefine);
1040 addSMELibCall(
1041 MIB&: SMEStateSize, LC: RTLIB::SMEABI_SME_STATE_SIZE,
1042 ExpectedCC: CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X1);
1043 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: BufferSize)
1044 .addReg(RegNo: AArch64::X0);
1045 }
1046
1047 // Allocate a buffer object of the size given __arm_sme_state_size.
1048 {
1049 MachineFrameInfo &MFI = MF->getFrameInfo();
1050 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: AArch64::SUBXrx64), DestReg: AArch64::SP)
1051 .addReg(RegNo: AArch64::SP)
1052 .addReg(RegNo: BufferSize)
1053 .addImm(Val: AArch64_AM::getArithExtendImm(ET: AArch64_AM::UXTX, Imm: 0));
1054 BuildMI(BB&: MBB, I: MBBI, MIMD: DL, MCID: TII->get(Opcode: TargetOpcode::COPY), DestReg: BufferPtr)
1055 .addReg(RegNo: AArch64::SP);
1056
1057 // We have just allocated a variable sized object, tell this to PEI.
1058 MFI.CreateVariableSizedObject(Alignment: Align(16), Alloca: nullptr);
1059 }
1060
1061 restorePhyRegSave(RegSave, MBB, MBBI, DL);
1062}
1063
1064struct FromState {
1065 ZAState From;
1066
1067 constexpr uint8_t to(ZAState To) const {
1068 static_assert(NUM_ZA_STATE < 16, "expected ZAState to fit in 4-bits");
1069 return uint8_t(From) << 4 | uint8_t(To);
1070 }
1071};
1072
1073constexpr FromState transitionFrom(ZAState From) { return FromState{.From: From}; }
1074
1075void MachineSMEABI::emitStateChange(EmitContext &Context,
1076 MachineBasicBlock &MBB,
1077 MachineBasicBlock::iterator InsertPt,
1078 ZAState From, ZAState To,
1079 LiveRegs PhysLiveRegs) {
1080 // ZA not used.
1081 if (From == ZAState::ANY || To == ZAState::ANY)
1082 return;
1083
1084 // If we're exiting from the ENTRY state that means that the function has not
1085 // used ZA, so in the case of private ZA/ZT0 functions we can omit any set up.
1086 if (From == ZAState::ENTRY && To == ZAState::OFF)
1087 return;
1088
1089 // TODO: Avoid setting up the save buffer if there's no transition to
1090 // LOCAL_SAVED.
1091 if (From == ZAState::ENTRY) {
1092 assert(&MBB == &MBB.getParent()->front() &&
1093 "ENTRY state only valid in entry block");
1094 emitSMEPrologue(MBB, MBBI: MBB.getFirstNonPHI());
1095 if (To == ZAState::ACTIVE)
1096 return; // Nothing more to do (ZA is active after the prologue).
1097
1098 // Note: "emitNewZAPrologue" zeros ZA, so we may need to setup a lazy save
1099 // if "To" is "ZAState::LOCAL_SAVED". It may be possible to improve this
1100 // case by changing the placement of the zero instruction.
1101 From = ZAState::ACTIVE;
1102 }
1103
1104 SMEAttrs SMEFnAttrs = AFI->getSMEFnAttrs();
1105 bool IsAgnosticZA = SMEFnAttrs.hasAgnosticZAInterface();
1106 bool HasZT0State = SMEFnAttrs.hasZT0State();
1107 bool HasZAState = IsAgnosticZA || SMEFnAttrs.hasZAState();
1108
1109 switch (transitionFrom(From).to(To)) {
1110 // This section handles: ACTIVE <-> ACTIVE_ZT0_SAVED
1111 case transitionFrom(From: ZAState::ACTIVE).to(To: ZAState::ACTIVE_ZT0_SAVED):
1112 emitZT0SaveRestore(Context, MBB, MBBI: InsertPt, /*IsSave=*/true);
1113 break;
1114 case transitionFrom(From: ZAState::ACTIVE_ZT0_SAVED).to(To: ZAState::ACTIVE):
1115 emitZT0SaveRestore(Context, MBB, MBBI: InsertPt, /*IsSave=*/false);
1116 break;
1117
1118 // This section handles: ACTIVE[_ZT0_SAVED] -> LOCAL_SAVED
1119 case transitionFrom(From: ZAState::ACTIVE).to(To: ZAState::LOCAL_SAVED):
1120 case transitionFrom(From: ZAState::ACTIVE_ZT0_SAVED).to(To: ZAState::LOCAL_SAVED):
1121 if (HasZT0State && From == ZAState::ACTIVE)
1122 emitZT0SaveRestore(Context, MBB, MBBI: InsertPt, /*IsSave=*/true);
1123 if (HasZAState)
1124 emitZASave(Context, MBB, MBBI: InsertPt, PhysLiveRegs);
1125 break;
1126
1127 // This section handles: ACTIVE -> LOCAL_COMMITTED
1128 case transitionFrom(From: ZAState::ACTIVE).to(To: ZAState::LOCAL_COMMITTED):
1129 // TODO: We could support ZA state here, but this transition is currently
1130 // only possible when we _don't_ have ZA state.
1131 assert(HasZT0State && !HasZAState && "Expect to only have ZT0 state.");
1132 emitZT0SaveRestore(Context, MBB, MBBI: InsertPt, /*IsSave=*/true);
1133 emitZAMode(MBB, MBBI: InsertPt, /*ClearTPIDR2=*/false, /*On=*/false);
1134 break;
1135
1136 // This section handles: LOCAL_COMMITTED -> (OFF|LOCAL_SAVED)
1137 case transitionFrom(From: ZAState::LOCAL_COMMITTED).to(To: ZAState::OFF):
1138 case transitionFrom(From: ZAState::LOCAL_COMMITTED).to(To: ZAState::LOCAL_SAVED):
1139 // These transitions are a no-op.
1140 break;
1141
1142 // This section handles: LOCAL_(SAVED|COMMITTED) -> ACTIVE[_ZT0_SAVED]
1143 case transitionFrom(From: ZAState::LOCAL_COMMITTED).to(To: ZAState::ACTIVE):
1144 case transitionFrom(From: ZAState::LOCAL_COMMITTED).to(To: ZAState::ACTIVE_ZT0_SAVED):
1145 case transitionFrom(From: ZAState::LOCAL_SAVED).to(To: ZAState::ACTIVE):
1146 case transitionFrom(From: ZAState::LOCAL_SAVED).to(To: ZAState::ACTIVE_ZT0_SAVED):
1147 if (HasZAState)
1148 emitZARestore(Context, MBB, MBBI: InsertPt, PhysLiveRegs);
1149 else
1150 emitZAMode(MBB, MBBI: InsertPt, /*ClearTPIDR2=*/false, /*On=*/true);
1151 if (HasZT0State && To == ZAState::ACTIVE)
1152 emitZT0SaveRestore(Context, MBB, MBBI: InsertPt, /*IsSave=*/false);
1153 break;
1154
1155 // This section handles transitions to OFF (not previously covered)
1156 case transitionFrom(From: ZAState::ACTIVE).to(To: ZAState::OFF):
1157 case transitionFrom(From: ZAState::ACTIVE_ZT0_SAVED).to(To: ZAState::OFF):
1158 case transitionFrom(From: ZAState::LOCAL_SAVED).to(To: ZAState::OFF):
1159 assert(SMEFnAttrs.hasPrivateZAInterface() &&
1160 "Did not expect to turn ZA off in shared/agnostic ZA function");
1161 emitZAMode(MBB, MBBI: InsertPt, /*ClearTPIDR2=*/From == ZAState::LOCAL_SAVED,
1162 /*On=*/false);
1163 break;
1164
1165 default:
1166 dbgs() << "Error: Transition from " << getZAStateString(State: From) << " to "
1167 << getZAStateString(State: To) << '\n';
1168 llvm_unreachable("Unimplemented state transition");
1169 }
1170}
1171
1172/// Returns true if private ZA setup can be elided. This occurs when there is
1173/// no instruction within the function that requires ZA to be active.
1174static bool canElidePrivateZASetup(const FunctionInfo &FnInfo) {
1175 for (const BlockInfo &BlockInfo : FnInfo.Blocks) {
1176 for (const InstInfo &InstInfo : BlockInfo.Insts) {
1177 if (InstInfo.NeededState == ZAState::ACTIVE ||
1178 InstInfo.NeededState == ZAState::ACTIVE_ZT0_SAVED)
1179 return false;
1180 }
1181 }
1182 return true;
1183}
1184
1185} // end anonymous namespace
1186
1187INITIALIZE_PASS(MachineSMEABI, "aarch64-machine-sme-abi", "Machine SME ABI",
1188 false, false)
1189
1190bool MachineSMEABI::runOnMachineFunction(MachineFunction &MF) {
1191 AFI = MF.getInfo<AArch64FunctionInfo>();
1192 SMEAttrs SMEFnAttrs = AFI->getSMEFnAttrs();
1193 if (!SMEFnAttrs.hasZAState() && !SMEFnAttrs.hasZT0State() &&
1194 !SMEFnAttrs.hasAgnosticZAInterface())
1195 return false;
1196
1197 Subtarget = &MF.getSubtarget<AArch64Subtarget>();
1198 if (!Subtarget->hasSME() && !SMEFnAttrs.hasAgnosticZAInterface())
1199 return false;
1200
1201 assert(MF.getRegInfo().isSSA() && "Expected to be run on SSA form!");
1202
1203 this->MF = &MF;
1204 ORE = &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE();
1205 LLI = &getAnalysis<LibcallLoweringInfoWrapper>().getLibcallLowering(
1206 M: *MF.getFunction().getParent(), Subtarget: *Subtarget);
1207 TII = Subtarget->getInstrInfo();
1208 TRI = Subtarget->getRegisterInfo();
1209 MRI = &MF.getRegInfo();
1210
1211 const EdgeBundles &Bundles =
1212 getAnalysis<EdgeBundlesWrapperLegacy>().getEdgeBundles();
1213
1214 FunctionInfo FnInfo = collectNeededZAStates(SMEFnAttrs);
1215
1216 if (SMEFnAttrs.hasPrivateZAInterface() && canElidePrivateZASetup(FnInfo))
1217 return false;
1218
1219 SmallVector<ZAState> BundleStates = assignBundleZAStates(Bundles, FnInfo);
1220
1221 EmitContext Context;
1222 insertStateChanges(Context, FnInfo, Bundles, BundleStates);
1223
1224 if (Context.needsSaveBuffer()) {
1225 if (FnInfo.AfterSMEProloguePt) {
1226 // Note: With inline stack probes the AfterSMEProloguePt may not be in the
1227 // entry block (due to the probing loop).
1228 MachineBasicBlock::iterator MBBI = *FnInfo.AfterSMEProloguePt;
1229 emitAllocateZASaveBuffer(Context, MBB&: *MBBI->getParent(), MBBI,
1230 PhysLiveRegs: FnInfo.PhysLiveRegsAfterSMEPrologue);
1231 } else {
1232 MachineBasicBlock &EntryBlock = MF.front();
1233 emitAllocateZASaveBuffer(
1234 Context, MBB&: EntryBlock, MBBI: EntryBlock.getFirstNonPHI(),
1235 PhysLiveRegs: FnInfo.Blocks[EntryBlock.getNumber()].PhysLiveRegsAtEntry);
1236 }
1237 }
1238
1239 return true;
1240}
1241
1242FunctionPass *llvm::createMachineSMEABIPass(CodeGenOptLevel OptLevel) {
1243 return new MachineSMEABI(OptLevel);
1244}
1245