1207 lines · cpp
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-exception6//7//===----------------------------------------------------------------------===//8//9// This pass implements the SME ABI requirements for ZA state. This includes10// 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 a15// specific state (e.g., "ACTIVE" or "SAVED") and inserting the necessary state16// transitions to ensure ZA is in the required state before instructions. State17// transitions represent actions such as setting up or restoring a lazy save.18// Certain points within a function may also have predefined states independent19// of any instructions, for example, a "shared_za" function is always entered20// and exited in the "ACTIVE" state.21//22// To handle ZA state across control flow, we make use of edge bundling. This23// assigns each block an "incoming" and "outgoing" edge bundle (representing24// 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 is26// joined with the incoming bundle of all successors. The result is that each27// bundle can be assigned a single ZA state, which ensures the state required by28// all a blocks' successors is the same, and that each basic block will always29// be entered with the same ZA state. This eliminates the need for splitting30// 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// BB138// BB239//40// Initial Bundles Joined Bundles41//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 the54// 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/MachineRegisterInfo.h"67#include "llvm/CodeGen/TargetRegisterInfo.h"68 69using namespace llvm;70 71#define DEBUG_TYPE "aarch64-machine-sme-abi"72 73namespace {74 75// Note: For agnostic ZA, we assume the function is always entered/exited in the76// "ACTIVE" state -- this _may_ not be the case (since OFF is also a77// possibility, but for the purpose of placing ZA saves/restores, that does not78// matter).79enum ZAState : uint8_t {80 // Any/unknown state (not valid)81 ANY = 0,82 83 // ZA is in use and active (i.e. within the accumulator)84 ACTIVE,85 86 // ZA is active, but ZT0 has been saved.87 // This handles the edge case of sharedZA && !sharesZT0.88 ACTIVE_ZT0_SAVED,89 90 // A ZA save has been set up or committed (i.e. ZA is dormant or off)91 // If the function uses ZT0 it must also be saved.92 LOCAL_SAVED,93 94 // ZA has been committed to the lazy save buffer of the current function.95 // If the function uses ZT0 it must also be saved.96 // ZA is off.97 LOCAL_COMMITTED,98 99 // The ZA/ZT0 state on entry to the function.100 ENTRY,101 102 // ZA is off.103 OFF,104 105 // The number of ZA states (not a valid state)106 NUM_ZA_STATE107};108 109/// A bitmask enum to record live physical registers that the "emit*" routines110/// may need to preserve. Note: This only tracks registers we may clobber.111enum LiveRegs : uint8_t {112 None = 0,113 NZCV = 1 << 0,114 W0 = 1 << 1,115 W0_HI = 1 << 2,116 X0 = W0 | W0_HI,117 LLVM_MARK_AS_BITMASK_ENUM(/* LargestValue = */ W0_HI)118};119 120/// Holds the virtual registers live physical registers have been saved to.121struct PhysRegSave {122 LiveRegs PhysLiveRegs;123 Register StatusFlags = AArch64::NoRegister;124 Register X0Save = AArch64::NoRegister;125};126 127/// Contains the needed ZA state (and live registers) at an instruction. That is128/// the state ZA must be in _before_ "InsertPt".129struct InstInfo {130 ZAState NeededState{ZAState::ANY};131 MachineBasicBlock::iterator InsertPt;132 LiveRegs PhysLiveRegs = LiveRegs::None;133};134 135/// Contains the needed ZA state for each instruction in a block. Instructions136/// that do not require a ZA state are not recorded.137struct BlockInfo {138 SmallVector<InstInfo> Insts;139 ZAState FixedEntryState{ZAState::ANY};140 ZAState DesiredIncomingState{ZAState::ANY};141 ZAState DesiredOutgoingState{ZAState::ANY};142 LiveRegs PhysLiveRegsAtEntry = LiveRegs::None;143 LiveRegs PhysLiveRegsAtExit = LiveRegs::None;144};145 146/// Contains the needed ZA state information for all blocks within a function.147struct FunctionInfo {148 SmallVector<BlockInfo> Blocks;149 std::optional<MachineBasicBlock::iterator> AfterSMEProloguePt;150 LiveRegs PhysLiveRegsAfterSMEPrologue = LiveRegs::None;151};152 153/// State/helpers that is only needed when emitting code to handle154/// saving/restoring ZA.155class EmitContext {156public:157 EmitContext() = default;158 159 /// Get or create a TPIDR2 block in \p MF.160 int getTPIDR2Block(MachineFunction &MF) {161 if (TPIDR2BlockFI)162 return *TPIDR2BlockFI;163 MachineFrameInfo &MFI = MF.getFrameInfo();164 TPIDR2BlockFI = MFI.CreateStackObject(16, Align(16), false);165 return *TPIDR2BlockFI;166 }167 168 /// Get or create agnostic ZA buffer pointer in \p MF.169 Register getAgnosticZABufferPtr(MachineFunction &MF) {170 if (AgnosticZABufferPtr != AArch64::NoRegister)171 return AgnosticZABufferPtr;172 Register BufferPtr =173 MF.getInfo<AArch64FunctionInfo>()->getEarlyAllocSMESaveBuffer();174 AgnosticZABufferPtr =175 BufferPtr != AArch64::NoRegister176 ? BufferPtr177 : MF.getRegInfo().createVirtualRegister(&AArch64::GPR64RegClass);178 return AgnosticZABufferPtr;179 }180 181 int getZT0SaveSlot(MachineFunction &MF) {182 if (ZT0SaveFI)183 return *ZT0SaveFI;184 MachineFrameInfo &MFI = MF.getFrameInfo();185 ZT0SaveFI = MFI.CreateSpillStackObject(64, Align(16));186 return *ZT0SaveFI;187 }188 189 /// Returns true if the function must allocate a ZA save buffer on entry. This190 /// will be the case if, at any point in the function, a ZA save was emitted.191 bool needsSaveBuffer() const {192 assert(!(TPIDR2BlockFI && AgnosticZABufferPtr) &&193 "Cannot have both a TPIDR2 block and agnostic ZA buffer");194 return TPIDR2BlockFI || AgnosticZABufferPtr != AArch64::NoRegister;195 }196 197private:198 std::optional<int> ZT0SaveFI;199 std::optional<int> TPIDR2BlockFI;200 Register AgnosticZABufferPtr = AArch64::NoRegister;201};202 203/// Checks if \p State is a legal edge bundle state. For a state to be a legal204/// bundle state, it must be possible to transition from it to any other bundle205/// state without losing any ZA state. This is the case for ACTIVE/LOCAL_SAVED,206/// as you can transition between those states by saving/restoring ZA. The OFF207/// state would not be legal, as transitioning to it drops the content of ZA.208static bool isLegalEdgeBundleZAState(ZAState State) {209 switch (State) {210 case ZAState::ACTIVE: // ZA state within the accumulator/ZT0.211 case ZAState::ACTIVE_ZT0_SAVED: // ZT0 is saved (ZA is active).212 case ZAState::LOCAL_SAVED: // ZA state may be saved on the stack.213 case ZAState::LOCAL_COMMITTED: // ZA state is saved on the stack.214 return true;215 default:216 return false;217 }218}219 220StringRef getZAStateString(ZAState State) {221#define MAKE_CASE(V) \222 case V: \223 return #V;224 switch (State) {225 MAKE_CASE(ZAState::ANY)226 MAKE_CASE(ZAState::ACTIVE)227 MAKE_CASE(ZAState::ACTIVE_ZT0_SAVED)228 MAKE_CASE(ZAState::LOCAL_SAVED)229 MAKE_CASE(ZAState::LOCAL_COMMITTED)230 MAKE_CASE(ZAState::ENTRY)231 MAKE_CASE(ZAState::OFF)232 default:233 llvm_unreachable("Unexpected ZAState");234 }235#undef MAKE_CASE236}237 238static bool isZAorZTRegOp(const TargetRegisterInfo &TRI,239 const MachineOperand &MO) {240 if (!MO.isReg() || !MO.getReg().isPhysical())241 return false;242 return any_of(TRI.subregs_inclusive(MO.getReg()), [](const MCPhysReg &SR) {243 return AArch64::MPR128RegClass.contains(SR) ||244 AArch64::ZTRRegClass.contains(SR);245 });246}247 248/// Returns the required ZA state needed before \p MI and an iterator pointing249/// to where any code required to change the ZA state should be inserted.250static std::pair<ZAState, MachineBasicBlock::iterator>251getInstNeededZAState(const TargetRegisterInfo &TRI, MachineInstr &MI,252 SMEAttrs SMEFnAttrs) {253 MachineBasicBlock::iterator InsertPt(MI);254 255 // Note: InOutZAUsePseudo, RequiresZASavePseudo, and RequiresZT0SavePseudo are256 // intended to mark the position immediately before a call. Due to257 // SelectionDAG constraints, these markers occur after the ADJCALLSTACKDOWN,258 // so we use std::prev(InsertPt) to get the position before the call.259 260 if (MI.getOpcode() == AArch64::InOutZAUsePseudo)261 return {ZAState::ACTIVE, std::prev(InsertPt)};262 263 // Note: If we need to save both ZA and ZT0 we use RequiresZASavePseudo.264 if (MI.getOpcode() == AArch64::RequiresZASavePseudo)265 return {ZAState::LOCAL_SAVED, std::prev(InsertPt)};266 267 // If we only need to save ZT0 there's two cases to consider:268 // 1. The function has ZA state (that we don't need to save).269 // - In this case we switch to the "ACTIVE_ZT0_SAVED" state.270 // This only saves ZT0.271 // 2. The function does not have ZA state272 // - In this case we switch to "LOCAL_COMMITTED" state.273 // This saves ZT0 and turns ZA off.274 if (MI.getOpcode() == AArch64::RequiresZT0SavePseudo) {275 return {SMEFnAttrs.hasZAState() ? ZAState::ACTIVE_ZT0_SAVED276 : ZAState::LOCAL_COMMITTED,277 std::prev(InsertPt)};278 }279 280 if (MI.isReturn()) {281 bool ZAOffAtReturn = SMEFnAttrs.hasPrivateZAInterface();282 return {ZAOffAtReturn ? ZAState::OFF : ZAState::ACTIVE, InsertPt};283 }284 285 for (auto &MO : MI.operands()) {286 if (isZAorZTRegOp(TRI, MO))287 return {ZAState::ACTIVE, InsertPt};288 }289 290 return {ZAState::ANY, InsertPt};291}292 293struct MachineSMEABI : public MachineFunctionPass {294 inline static char ID = 0;295 296 MachineSMEABI(CodeGenOptLevel OptLevel = CodeGenOptLevel::Default)297 : MachineFunctionPass(ID), OptLevel(OptLevel) {}298 299 bool runOnMachineFunction(MachineFunction &MF) override;300 301 StringRef getPassName() const override { return "Machine SME ABI pass"; }302 303 void getAnalysisUsage(AnalysisUsage &AU) const override {304 AU.setPreservesCFG();305 AU.addRequired<EdgeBundlesWrapperLegacy>();306 AU.addPreservedID(MachineLoopInfoID);307 AU.addPreservedID(MachineDominatorsID);308 MachineFunctionPass::getAnalysisUsage(AU);309 }310 311 /// Collects the needed ZA state (and live registers) before each instruction312 /// within the machine function.313 FunctionInfo collectNeededZAStates(SMEAttrs SMEFnAttrs);314 315 /// Assigns each edge bundle a ZA state based on the needed states of blocks316 /// that have incoming or outgoing edges in that bundle.317 SmallVector<ZAState> assignBundleZAStates(const EdgeBundles &Bundles,318 const FunctionInfo &FnInfo);319 320 /// Inserts code to handle changes between ZA states within the function.321 /// E.g., ACTIVE -> LOCAL_SAVED will insert code required to save ZA.322 void insertStateChanges(EmitContext &, const FunctionInfo &FnInfo,323 const EdgeBundles &Bundles,324 ArrayRef<ZAState> BundleStates);325 326 /// Propagates desired states forwards (from predecessors -> successors) if327 /// \p Forwards, otherwise, propagates backwards (from successors ->328 /// predecessors).329 void propagateDesiredStates(FunctionInfo &FnInfo, bool Forwards = true);330 331 void emitZT0SaveRestore(EmitContext &, MachineBasicBlock &MBB,332 MachineBasicBlock::iterator MBBI, bool IsSave);333 334 // Emission routines for private and shared ZA functions (using lazy saves).335 void emitSMEPrologue(MachineBasicBlock &MBB,336 MachineBasicBlock::iterator MBBI);337 void emitRestoreLazySave(EmitContext &, MachineBasicBlock &MBB,338 MachineBasicBlock::iterator MBBI,339 LiveRegs PhysLiveRegs);340 void emitSetupLazySave(EmitContext &, MachineBasicBlock &MBB,341 MachineBasicBlock::iterator MBBI);342 void emitAllocateLazySaveBuffer(EmitContext &, MachineBasicBlock &MBB,343 MachineBasicBlock::iterator MBBI);344 void emitZAMode(MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI,345 bool ClearTPIDR2, bool On);346 347 // Emission routines for agnostic ZA functions.348 void emitSetupFullZASave(MachineBasicBlock &MBB,349 MachineBasicBlock::iterator MBBI,350 LiveRegs PhysLiveRegs);351 // Emit a "full" ZA save or restore. It is "full" in the sense that this352 // function will emit a call to __arm_sme_save or __arm_sme_restore, which353 // handles saving and restoring both ZA and ZT0.354 void emitFullZASaveRestore(EmitContext &, MachineBasicBlock &MBB,355 MachineBasicBlock::iterator MBBI,356 LiveRegs PhysLiveRegs, bool IsSave);357 void emitAllocateFullZASaveBuffer(EmitContext &, MachineBasicBlock &MBB,358 MachineBasicBlock::iterator MBBI,359 LiveRegs PhysLiveRegs);360 361 /// Attempts to find an insertion point before \p Inst where the status flags362 /// are not live. If \p Inst is `Block.Insts.end()` a point before the end of363 /// the block is found.364 std::pair<MachineBasicBlock::iterator, LiveRegs>365 findStateChangeInsertionPoint(MachineBasicBlock &MBB, const BlockInfo &Block,366 SmallVectorImpl<InstInfo>::const_iterator Inst);367 void emitStateChange(EmitContext &, MachineBasicBlock &MBB,368 MachineBasicBlock::iterator MBBI, ZAState From,369 ZAState To, LiveRegs PhysLiveRegs);370 371 // Helpers for switching between lazy/full ZA save/restore routines.372 void emitZASave(EmitContext &Context, MachineBasicBlock &MBB,373 MachineBasicBlock::iterator MBBI, LiveRegs PhysLiveRegs) {374 if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())375 return emitFullZASaveRestore(Context, MBB, MBBI, PhysLiveRegs,376 /*IsSave=*/true);377 return emitSetupLazySave(Context, MBB, MBBI);378 }379 void emitZARestore(EmitContext &Context, MachineBasicBlock &MBB,380 MachineBasicBlock::iterator MBBI, LiveRegs PhysLiveRegs) {381 if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())382 return emitFullZASaveRestore(Context, MBB, MBBI, PhysLiveRegs,383 /*IsSave=*/false);384 return emitRestoreLazySave(Context, MBB, MBBI, PhysLiveRegs);385 }386 void emitAllocateZASaveBuffer(EmitContext &Context, MachineBasicBlock &MBB,387 MachineBasicBlock::iterator MBBI,388 LiveRegs PhysLiveRegs) {389 if (AFI->getSMEFnAttrs().hasAgnosticZAInterface())390 return emitAllocateFullZASaveBuffer(Context, MBB, MBBI, PhysLiveRegs);391 return emitAllocateLazySaveBuffer(Context, MBB, MBBI);392 }393 394 /// Save live physical registers to virtual registers.395 PhysRegSave createPhysRegSave(LiveRegs PhysLiveRegs, MachineBasicBlock &MBB,396 MachineBasicBlock::iterator MBBI, DebugLoc DL);397 /// Restore physical registers from a save of their previous values.398 void restorePhyRegSave(const PhysRegSave &RegSave, MachineBasicBlock &MBB,399 MachineBasicBlock::iterator MBBI, DebugLoc DL);400 401private:402 CodeGenOptLevel OptLevel = CodeGenOptLevel::Default;403 404 MachineFunction *MF = nullptr;405 const AArch64Subtarget *Subtarget = nullptr;406 const AArch64RegisterInfo *TRI = nullptr;407 const AArch64FunctionInfo *AFI = nullptr;408 const TargetInstrInfo *TII = nullptr;409 MachineRegisterInfo *MRI = nullptr;410 MachineLoopInfo *MLI = nullptr;411};412 413static LiveRegs getPhysLiveRegs(LiveRegUnits const &LiveUnits) {414 LiveRegs PhysLiveRegs = LiveRegs::None;415 if (!LiveUnits.available(AArch64::NZCV))416 PhysLiveRegs |= LiveRegs::NZCV;417 // We have to track W0 and X0 separately as otherwise things can get418 // confused if we attempt to preserve X0 but only W0 was defined.419 if (!LiveUnits.available(AArch64::W0))420 PhysLiveRegs |= LiveRegs::W0;421 if (!LiveUnits.available(AArch64::W0_HI))422 PhysLiveRegs |= LiveRegs::W0_HI;423 return PhysLiveRegs;424}425 426static void setPhysLiveRegs(LiveRegUnits &LiveUnits, LiveRegs PhysLiveRegs) {427 if (PhysLiveRegs & LiveRegs::NZCV)428 LiveUnits.addReg(AArch64::NZCV);429 if (PhysLiveRegs & LiveRegs::W0)430 LiveUnits.addReg(AArch64::W0);431 if (PhysLiveRegs & LiveRegs::W0_HI)432 LiveUnits.addReg(AArch64::W0_HI);433}434 435[[maybe_unused]] bool isCallStartOpcode(unsigned Opc) {436 switch (Opc) {437 case AArch64::TLSDESC_CALLSEQ:438 case AArch64::TLSDESC_AUTH_CALLSEQ:439 case AArch64::ADJCALLSTACKDOWN:440 return true;441 default:442 return false;443 }444}445 446FunctionInfo MachineSMEABI::collectNeededZAStates(SMEAttrs SMEFnAttrs) {447 assert((SMEFnAttrs.hasAgnosticZAInterface() || SMEFnAttrs.hasZT0State() ||448 SMEFnAttrs.hasZAState()) &&449 "Expected function to have ZA/ZT0 state!");450 451 SmallVector<BlockInfo> Blocks(MF->getNumBlockIDs());452 LiveRegs PhysLiveRegsAfterSMEPrologue = LiveRegs::None;453 std::optional<MachineBasicBlock::iterator> AfterSMEProloguePt;454 455 for (MachineBasicBlock &MBB : *MF) {456 BlockInfo &Block = Blocks[MBB.getNumber()];457 458 if (MBB.isEntryBlock()) {459 // Entry block:460 Block.FixedEntryState = ZAState::ENTRY;461 } else if (MBB.isEHPad()) {462 // EH entry block:463 Block.FixedEntryState = ZAState::LOCAL_COMMITTED;464 }465 466 LiveRegUnits LiveUnits(*TRI);467 LiveUnits.addLiveOuts(MBB);468 469 Block.PhysLiveRegsAtExit = getPhysLiveRegs(LiveUnits);470 auto FirstTerminatorInsertPt = MBB.getFirstTerminator();471 auto FirstNonPhiInsertPt = MBB.getFirstNonPHI();472 for (MachineInstr &MI : reverse(MBB)) {473 MachineBasicBlock::iterator MBBI(MI);474 LiveUnits.stepBackward(MI);475 LiveRegs PhysLiveRegs = getPhysLiveRegs(LiveUnits);476 // The SMEStateAllocPseudo marker is added to a function if the save477 // buffer was allocated in SelectionDAG. It marks the end of the478 // allocation -- which is a safe point for this pass to insert any TPIDR2479 // block setup.480 if (MI.getOpcode() == AArch64::SMEStateAllocPseudo) {481 AfterSMEProloguePt = MBBI;482 PhysLiveRegsAfterSMEPrologue = PhysLiveRegs;483 }484 // Note: We treat Agnostic ZA as inout_za with an alternate save/restore.485 auto [NeededState, InsertPt] = getInstNeededZAState(*TRI, MI, SMEFnAttrs);486 assert((InsertPt == MBBI || isCallStartOpcode(InsertPt->getOpcode())) &&487 "Unexpected state change insertion point!");488 // TODO: Do something to avoid state changes where NZCV is live.489 if (MBBI == FirstTerminatorInsertPt)490 Block.PhysLiveRegsAtExit = PhysLiveRegs;491 if (MBBI == FirstNonPhiInsertPt)492 Block.PhysLiveRegsAtEntry = PhysLiveRegs;493 if (NeededState != ZAState::ANY)494 Block.Insts.push_back({NeededState, InsertPt, PhysLiveRegs});495 }496 497 // Reverse vector (as we had to iterate backwards for liveness).498 std::reverse(Block.Insts.begin(), Block.Insts.end());499 500 // Record the desired states on entry/exit of this block. These are the501 // states that would not incur a state transition.502 if (!Block.Insts.empty()) {503 Block.DesiredIncomingState = Block.Insts.front().NeededState;504 Block.DesiredOutgoingState = Block.Insts.back().NeededState;505 }506 }507 508 return FunctionInfo{std::move(Blocks), AfterSMEProloguePt,509 PhysLiveRegsAfterSMEPrologue};510}511 512void MachineSMEABI::propagateDesiredStates(FunctionInfo &FnInfo,513 bool Forwards) {514 // If `Forwards`, this propagates desired states from predecessors to515 // successors, otherwise, this propagates states from successors to516 // predecessors.517 auto GetBlockState = [](BlockInfo &Block, bool Incoming) -> ZAState & {518 return Incoming ? Block.DesiredIncomingState : Block.DesiredOutgoingState;519 };520 521 SmallVector<MachineBasicBlock *> Worklist;522 for (auto [BlockID, BlockInfo] : enumerate(FnInfo.Blocks)) {523 if (!isLegalEdgeBundleZAState(GetBlockState(BlockInfo, Forwards)))524 Worklist.push_back(MF->getBlockNumbered(BlockID));525 }526 527 while (!Worklist.empty()) {528 MachineBasicBlock *MBB = Worklist.pop_back_val();529 BlockInfo &Block = FnInfo.Blocks[MBB->getNumber()];530 531 // Pick a legal edge bundle state that matches the majority of532 // predecessors/successors.533 int StateCounts[ZAState::NUM_ZA_STATE] = {0};534 for (MachineBasicBlock *PredOrSucc :535 Forwards ? predecessors(MBB) : successors(MBB)) {536 BlockInfo &PredOrSuccBlock = FnInfo.Blocks[PredOrSucc->getNumber()];537 ZAState ZAState = GetBlockState(PredOrSuccBlock, !Forwards);538 if (isLegalEdgeBundleZAState(ZAState))539 StateCounts[ZAState]++;540 }541 542 ZAState PropagatedState = ZAState(max_element(StateCounts) - StateCounts);543 ZAState &CurrentState = GetBlockState(Block, Forwards);544 if (PropagatedState != CurrentState) {545 CurrentState = PropagatedState;546 ZAState &OtherState = GetBlockState(Block, !Forwards);547 // Propagate to the incoming/outgoing state if that is also "ANY".548 if (OtherState == ZAState::ANY)549 OtherState = PropagatedState;550 // Push any successors/predecessors that may need updating to the551 // worklist.552 for (MachineBasicBlock *SuccOrPred :553 Forwards ? successors(MBB) : predecessors(MBB)) {554 BlockInfo &SuccOrPredBlock = FnInfo.Blocks[SuccOrPred->getNumber()];555 if (!isLegalEdgeBundleZAState(GetBlockState(SuccOrPredBlock, Forwards)))556 Worklist.push_back(SuccOrPred);557 }558 }559 }560}561 562/// Assigns each edge bundle a ZA state based on the needed states of blocks563/// that have incoming or outgoing edges in that bundle.564SmallVector<ZAState>565MachineSMEABI::assignBundleZAStates(const EdgeBundles &Bundles,566 const FunctionInfo &FnInfo) {567 SmallVector<ZAState> BundleStates(Bundles.getNumBundles());568 for (unsigned I = 0, E = Bundles.getNumBundles(); I != E; ++I) {569 LLVM_DEBUG(dbgs() << "Assigning ZA state for edge bundle: " << I << '\n');570 571 // Attempt to assign a ZA state for this bundle that minimizes state572 // transitions. Edges within loops are given a higher weight as we assume573 // they will be executed more than once.574 int EdgeStateCounts[ZAState::NUM_ZA_STATE] = {0};575 for (unsigned BlockID : Bundles.getBlocks(I)) {576 LLVM_DEBUG(dbgs() << "- bb." << BlockID);577 578 const BlockInfo &Block = FnInfo.Blocks[BlockID];579 bool InEdge = Bundles.getBundle(BlockID, /*Out=*/false) == I;580 bool OutEdge = Bundles.getBundle(BlockID, /*Out=*/true) == I;581 582 bool LegalInEdge =583 InEdge && isLegalEdgeBundleZAState(Block.DesiredIncomingState);584 bool LegalOutEgde =585 OutEdge && isLegalEdgeBundleZAState(Block.DesiredOutgoingState);586 if (LegalInEdge) {587 LLVM_DEBUG(dbgs() << " DesiredIncomingState: "588 << getZAStateString(Block.DesiredIncomingState));589 EdgeStateCounts[Block.DesiredIncomingState]++;590 }591 if (LegalOutEgde) {592 LLVM_DEBUG(dbgs() << " DesiredOutgoingState: "593 << getZAStateString(Block.DesiredOutgoingState));594 EdgeStateCounts[Block.DesiredOutgoingState]++;595 }596 if (!LegalInEdge && !LegalOutEgde)597 LLVM_DEBUG(dbgs() << " (no state preference)");598 LLVM_DEBUG(dbgs() << '\n');599 }600 601 ZAState BundleState =602 ZAState(max_element(EdgeStateCounts) - EdgeStateCounts);603 604 if (BundleState == ZAState::ANY)605 BundleState = ZAState::ACTIVE;606 607 LLVM_DEBUG({608 dbgs() << "Chosen ZA state: " << getZAStateString(BundleState) << '\n'609 << "Edge counts:";610 for (auto [State, Count] : enumerate(EdgeStateCounts))611 dbgs() << " " << getZAStateString(ZAState(State)) << ": " << Count;612 dbgs() << "\n\n";613 });614 615 BundleStates[I] = BundleState;616 }617 618 return BundleStates;619}620 621std::pair<MachineBasicBlock::iterator, LiveRegs>622MachineSMEABI::findStateChangeInsertionPoint(623 MachineBasicBlock &MBB, const BlockInfo &Block,624 SmallVectorImpl<InstInfo>::const_iterator Inst) {625 LiveRegs PhysLiveRegs;626 MachineBasicBlock::iterator InsertPt;627 if (Inst != Block.Insts.end()) {628 InsertPt = Inst->InsertPt;629 PhysLiveRegs = Inst->PhysLiveRegs;630 } else {631 InsertPt = MBB.getFirstTerminator();632 PhysLiveRegs = Block.PhysLiveRegsAtExit;633 }634 635 if (!(PhysLiveRegs & LiveRegs::NZCV))636 return {InsertPt, PhysLiveRegs}; // Nothing to do (no live flags).637 638 // Find the previous state change. We can not move before this point.639 MachineBasicBlock::iterator PrevStateChangeI;640 if (Inst == Block.Insts.begin()) {641 PrevStateChangeI = MBB.begin();642 } else {643 // Note: `std::prev(Inst)` is the previous InstInfo. We only create an644 // InstInfo object for instructions that require a specific ZA state, so the645 // InstInfo is the site of the previous state change in the block (which can646 // be several MIs earlier).647 PrevStateChangeI = std::prev(Inst)->InsertPt;648 }649 650 // Note: LiveUnits will only accurately track X0 and NZCV.651 LiveRegUnits LiveUnits(*TRI);652 setPhysLiveRegs(LiveUnits, PhysLiveRegs);653 for (MachineBasicBlock::iterator I = InsertPt; I != PrevStateChangeI; --I) {654 // Don't move before/into a call (which may have a state change before it).655 if (I->getOpcode() == TII->getCallFrameDestroyOpcode() || I->isCall())656 break;657 LiveUnits.stepBackward(*I);658 if (LiveUnits.available(AArch64::NZCV))659 return {I, getPhysLiveRegs(LiveUnits)};660 }661 return {InsertPt, PhysLiveRegs};662}663 664void MachineSMEABI::insertStateChanges(EmitContext &Context,665 const FunctionInfo &FnInfo,666 const EdgeBundles &Bundles,667 ArrayRef<ZAState> BundleStates) {668 for (MachineBasicBlock &MBB : *MF) {669 const BlockInfo &Block = FnInfo.Blocks[MBB.getNumber()];670 ZAState InState = BundleStates[Bundles.getBundle(MBB.getNumber(),671 /*Out=*/false)];672 673 ZAState CurrentState = Block.FixedEntryState;674 if (CurrentState == ZAState::ANY)675 CurrentState = InState;676 677 for (auto &Inst : Block.Insts) {678 if (CurrentState != Inst.NeededState) {679 auto [InsertPt, PhysLiveRegs] =680 findStateChangeInsertionPoint(MBB, Block, &Inst);681 emitStateChange(Context, MBB, InsertPt, CurrentState, Inst.NeededState,682 PhysLiveRegs);683 CurrentState = Inst.NeededState;684 }685 }686 687 if (MBB.succ_empty())688 continue;689 690 ZAState OutState =691 BundleStates[Bundles.getBundle(MBB.getNumber(), /*Out=*/true)];692 if (CurrentState != OutState) {693 auto [InsertPt, PhysLiveRegs] =694 findStateChangeInsertionPoint(MBB, Block, Block.Insts.end());695 emitStateChange(Context, MBB, InsertPt, CurrentState, OutState,696 PhysLiveRegs);697 }698 }699}700 701static DebugLoc getDebugLoc(MachineBasicBlock &MBB,702 MachineBasicBlock::iterator MBBI) {703 if (MBBI != MBB.end())704 return MBBI->getDebugLoc();705 return DebugLoc();706}707 708void MachineSMEABI::emitSetupLazySave(EmitContext &Context,709 MachineBasicBlock &MBB,710 MachineBasicBlock::iterator MBBI) {711 DebugLoc DL = getDebugLoc(MBB, MBBI);712 713 // Get pointer to TPIDR2 block.714 Register TPIDR2 = MRI->createVirtualRegister(&AArch64::GPR64spRegClass);715 Register TPIDR2Ptr = MRI->createVirtualRegister(&AArch64::GPR64RegClass);716 BuildMI(MBB, MBBI, DL, TII->get(AArch64::ADDXri), TPIDR2)717 .addFrameIndex(Context.getTPIDR2Block(*MF))718 .addImm(0)719 .addImm(0);720 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), TPIDR2Ptr)721 .addReg(TPIDR2);722 // Set TPIDR2_EL0 to point to TPIDR2 block.723 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSR))724 .addImm(AArch64SysReg::TPIDR2_EL0)725 .addReg(TPIDR2Ptr);726}727 728PhysRegSave MachineSMEABI::createPhysRegSave(LiveRegs PhysLiveRegs,729 MachineBasicBlock &MBB,730 MachineBasicBlock::iterator MBBI,731 DebugLoc DL) {732 PhysRegSave RegSave{PhysLiveRegs};733 if (PhysLiveRegs & LiveRegs::NZCV) {734 RegSave.StatusFlags = MRI->createVirtualRegister(&AArch64::GPR64RegClass);735 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MRS), RegSave.StatusFlags)736 .addImm(AArch64SysReg::NZCV)737 .addReg(AArch64::NZCV, RegState::Implicit);738 }739 // Note: Preserving X0 is "free" as this is before register allocation, so740 // the register allocator is still able to optimize these copies.741 if (PhysLiveRegs & LiveRegs::W0) {742 RegSave.X0Save = MRI->createVirtualRegister(PhysLiveRegs & LiveRegs::W0_HI743 ? &AArch64::GPR64RegClass744 : &AArch64::GPR32RegClass);745 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), RegSave.X0Save)746 .addReg(PhysLiveRegs & LiveRegs::W0_HI ? AArch64::X0 : AArch64::W0);747 }748 return RegSave;749}750 751void MachineSMEABI::restorePhyRegSave(const PhysRegSave &RegSave,752 MachineBasicBlock &MBB,753 MachineBasicBlock::iterator MBBI,754 DebugLoc DL) {755 if (RegSave.StatusFlags != AArch64::NoRegister)756 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSR))757 .addImm(AArch64SysReg::NZCV)758 .addReg(RegSave.StatusFlags)759 .addReg(AArch64::NZCV, RegState::ImplicitDefine);760 761 if (RegSave.X0Save != AArch64::NoRegister)762 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY),763 RegSave.PhysLiveRegs & LiveRegs::W0_HI ? AArch64::X0 : AArch64::W0)764 .addReg(RegSave.X0Save);765}766 767void MachineSMEABI::emitRestoreLazySave(EmitContext &Context,768 MachineBasicBlock &MBB,769 MachineBasicBlock::iterator MBBI,770 LiveRegs PhysLiveRegs) {771 auto *TLI = Subtarget->getTargetLowering();772 DebugLoc DL = getDebugLoc(MBB, MBBI);773 Register TPIDR2EL0 = MRI->createVirtualRegister(&AArch64::GPR64RegClass);774 Register TPIDR2 = AArch64::X0;775 776 // TODO: Emit these within the restore MBB to prevent unnecessary saves.777 PhysRegSave RegSave = createPhysRegSave(PhysLiveRegs, MBB, MBBI, DL);778 779 // Enable ZA.780 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSRpstatesvcrImm1))781 .addImm(AArch64SVCR::SVCRZA)782 .addImm(1);783 // Get current TPIDR2_EL0.784 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MRS), TPIDR2EL0)785 .addImm(AArch64SysReg::TPIDR2_EL0);786 // Get pointer to TPIDR2 block.787 BuildMI(MBB, MBBI, DL, TII->get(AArch64::ADDXri), TPIDR2)788 .addFrameIndex(Context.getTPIDR2Block(*MF))789 .addImm(0)790 .addImm(0);791 // (Conditionally) restore ZA state.792 BuildMI(MBB, MBBI, DL, TII->get(AArch64::RestoreZAPseudo))793 .addReg(TPIDR2EL0)794 .addReg(TPIDR2)795 .addExternalSymbol(TLI->getLibcallName(RTLIB::SMEABI_TPIDR2_RESTORE))796 .addRegMask(TRI->SMEABISupportRoutinesCallPreservedMaskFromX0());797 // Zero TPIDR2_EL0.798 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSR))799 .addImm(AArch64SysReg::TPIDR2_EL0)800 .addReg(AArch64::XZR);801 802 restorePhyRegSave(RegSave, MBB, MBBI, DL);803}804 805void MachineSMEABI::emitZAMode(MachineBasicBlock &MBB,806 MachineBasicBlock::iterator MBBI,807 bool ClearTPIDR2, bool On) {808 DebugLoc DL = getDebugLoc(MBB, MBBI);809 810 if (ClearTPIDR2)811 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSR))812 .addImm(AArch64SysReg::TPIDR2_EL0)813 .addReg(AArch64::XZR);814 815 // Disable ZA.816 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSRpstatesvcrImm1))817 .addImm(AArch64SVCR::SVCRZA)818 .addImm(On ? 1 : 0);819}820 821void MachineSMEABI::emitAllocateLazySaveBuffer(822 EmitContext &Context, MachineBasicBlock &MBB,823 MachineBasicBlock::iterator MBBI) {824 MachineFrameInfo &MFI = MF->getFrameInfo();825 DebugLoc DL = getDebugLoc(MBB, MBBI);826 Register SP = MRI->createVirtualRegister(&AArch64::GPR64RegClass);827 Register SVL = MRI->createVirtualRegister(&AArch64::GPR64RegClass);828 Register Buffer = AFI->getEarlyAllocSMESaveBuffer();829 830 // Calculate SVL.831 BuildMI(MBB, MBBI, DL, TII->get(AArch64::RDSVLI_XI), SVL).addImm(1);832 833 // 1. Allocate the lazy save buffer.834 if (Buffer == AArch64::NoRegister) {835 // TODO: On Windows, we allocate the lazy save buffer in SelectionDAG (so836 // Buffer != AArch64::NoRegister). This is done to reuse the existing837 // expansions (which can insert stack checks). This works, but it means we838 // will always allocate the lazy save buffer (even if the function contains839 // no lazy saves). If we want to handle Windows here, we'll need to840 // implement something similar to LowerWindowsDYNAMIC_STACKALLOC.841 assert(!Subtarget->isTargetWindows() &&842 "Lazy ZA save is not yet supported on Windows");843 Buffer = MRI->createVirtualRegister(&AArch64::GPR64RegClass);844 // Get original stack pointer.845 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), SP)846 .addReg(AArch64::SP);847 // Allocate a lazy-save buffer object of the size given, normally SVL * SVL848 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSUBXrrr), Buffer)849 .addReg(SVL)850 .addReg(SVL)851 .addReg(SP);852 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), AArch64::SP)853 .addReg(Buffer);854 // We have just allocated a variable sized object, tell this to PEI.855 MFI.CreateVariableSizedObject(Align(16), nullptr);856 }857 858 // 2. Setup the TPIDR2 block.859 {860 // Note: This case just needs to do `SVL << 48`. It is not implemented as we861 // generally don't support big-endian SVE/SME.862 if (!Subtarget->isLittleEndian())863 reportFatalInternalError(864 "TPIDR2 block initialization is not supported on big-endian targets");865 866 // Store buffer pointer and num_za_save_slices.867 // Bytes 10-15 are implicitly zeroed.868 BuildMI(MBB, MBBI, DL, TII->get(AArch64::STPXi))869 .addReg(Buffer)870 .addReg(SVL)871 .addFrameIndex(Context.getTPIDR2Block(*MF))872 .addImm(0);873 }874}875 876static constexpr unsigned ZERO_ALL_ZA_MASK = 0b11111111;877 878void MachineSMEABI::emitSMEPrologue(MachineBasicBlock &MBB,879 MachineBasicBlock::iterator MBBI) {880 auto *TLI = Subtarget->getTargetLowering();881 DebugLoc DL = getDebugLoc(MBB, MBBI);882 883 bool ZeroZA = AFI->getSMEFnAttrs().isNewZA();884 bool ZeroZT0 = AFI->getSMEFnAttrs().isNewZT0();885 if (AFI->getSMEFnAttrs().hasPrivateZAInterface()) {886 // Get current TPIDR2_EL0.887 Register TPIDR2EL0 = MRI->createVirtualRegister(&AArch64::GPR64RegClass);888 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MRS))889 .addReg(TPIDR2EL0, RegState::Define)890 .addImm(AArch64SysReg::TPIDR2_EL0);891 // If TPIDR2_EL0 is non-zero, commit the lazy save.892 // NOTE: Functions that only use ZT0 don't need to zero ZA.893 auto CommitZASave =894 BuildMI(MBB, MBBI, DL, TII->get(AArch64::CommitZASavePseudo))895 .addReg(TPIDR2EL0)896 .addImm(ZeroZA)897 .addImm(ZeroZT0)898 .addExternalSymbol(TLI->getLibcallName(RTLIB::SMEABI_TPIDR2_SAVE))899 .addRegMask(TRI->SMEABISupportRoutinesCallPreservedMaskFromX0());900 if (ZeroZA)901 CommitZASave.addDef(AArch64::ZAB0, RegState::ImplicitDefine);902 if (ZeroZT0)903 CommitZASave.addDef(AArch64::ZT0, RegState::ImplicitDefine);904 // Enable ZA (as ZA could have previously been in the OFF state).905 BuildMI(MBB, MBBI, DL, TII->get(AArch64::MSRpstatesvcrImm1))906 .addImm(AArch64SVCR::SVCRZA)907 .addImm(1);908 } else if (AFI->getSMEFnAttrs().hasSharedZAInterface()) {909 if (ZeroZA)910 BuildMI(MBB, MBBI, DL, TII->get(AArch64::ZERO_M))911 .addImm(ZERO_ALL_ZA_MASK)912 .addDef(AArch64::ZAB0, RegState::ImplicitDefine);913 if (ZeroZT0)914 BuildMI(MBB, MBBI, DL, TII->get(AArch64::ZERO_T)).addDef(AArch64::ZT0);915 }916}917 918void MachineSMEABI::emitFullZASaveRestore(EmitContext &Context,919 MachineBasicBlock &MBB,920 MachineBasicBlock::iterator MBBI,921 LiveRegs PhysLiveRegs, bool IsSave) {922 auto *TLI = Subtarget->getTargetLowering();923 DebugLoc DL = getDebugLoc(MBB, MBBI);924 Register BufferPtr = AArch64::X0;925 926 PhysRegSave RegSave = createPhysRegSave(PhysLiveRegs, MBB, MBBI, DL);927 928 // Copy the buffer pointer into X0.929 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), BufferPtr)930 .addReg(Context.getAgnosticZABufferPtr(*MF));931 932 // Call __arm_sme_save/__arm_sme_restore.933 BuildMI(MBB, MBBI, DL, TII->get(AArch64::BL))934 .addReg(BufferPtr, RegState::Implicit)935 .addExternalSymbol(TLI->getLibcallName(936 IsSave ? RTLIB::SMEABI_SME_SAVE : RTLIB::SMEABI_SME_RESTORE))937 .addRegMask(TRI->getCallPreservedMask(938 *MF,939 CallingConv::AArch64_SME_ABI_Support_Routines_PreserveMost_From_X1));940 941 restorePhyRegSave(RegSave, MBB, MBBI, DL);942}943 944void MachineSMEABI::emitZT0SaveRestore(EmitContext &Context,945 MachineBasicBlock &MBB,946 MachineBasicBlock::iterator MBBI,947 bool IsSave) {948 DebugLoc DL = getDebugLoc(MBB, MBBI);949 Register ZT0Save = MRI->createVirtualRegister(&AArch64::GPR64spRegClass);950 951 BuildMI(MBB, MBBI, DL, TII->get(AArch64::ADDXri), ZT0Save)952 .addFrameIndex(Context.getZT0SaveSlot(*MF))953 .addImm(0)954 .addImm(0);955 956 if (IsSave) {957 BuildMI(MBB, MBBI, DL, TII->get(AArch64::STR_TX))958 .addReg(AArch64::ZT0)959 .addReg(ZT0Save);960 } else {961 BuildMI(MBB, MBBI, DL, TII->get(AArch64::LDR_TX), AArch64::ZT0)962 .addReg(ZT0Save);963 }964}965 966void MachineSMEABI::emitAllocateFullZASaveBuffer(967 EmitContext &Context, MachineBasicBlock &MBB,968 MachineBasicBlock::iterator MBBI, LiveRegs PhysLiveRegs) {969 // Buffer already allocated in SelectionDAG.970 if (AFI->getEarlyAllocSMESaveBuffer())971 return;972 973 DebugLoc DL = getDebugLoc(MBB, MBBI);974 Register BufferPtr = Context.getAgnosticZABufferPtr(*MF);975 Register BufferSize = MRI->createVirtualRegister(&AArch64::GPR64RegClass);976 977 PhysRegSave RegSave = createPhysRegSave(PhysLiveRegs, MBB, MBBI, DL);978 979 // Calculate the SME state size.980 {981 auto *TLI = Subtarget->getTargetLowering();982 const AArch64RegisterInfo *TRI = Subtarget->getRegisterInfo();983 BuildMI(MBB, MBBI, DL, TII->get(AArch64::BL))984 .addExternalSymbol(TLI->getLibcallName(RTLIB::SMEABI_SME_STATE_SIZE))985 .addReg(AArch64::X0, RegState::ImplicitDefine)986 .addRegMask(TRI->getCallPreservedMask(987 *MF, CallingConv::988 AArch64_SME_ABI_Support_Routines_PreserveMost_From_X1));989 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), BufferSize)990 .addReg(AArch64::X0);991 }992 993 // Allocate a buffer object of the size given __arm_sme_state_size.994 {995 MachineFrameInfo &MFI = MF->getFrameInfo();996 BuildMI(MBB, MBBI, DL, TII->get(AArch64::SUBXrx64), AArch64::SP)997 .addReg(AArch64::SP)998 .addReg(BufferSize)999 .addImm(AArch64_AM::getArithExtendImm(AArch64_AM::UXTX, 0));1000 BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::COPY), BufferPtr)1001 .addReg(AArch64::SP);1002 1003 // We have just allocated a variable sized object, tell this to PEI.1004 MFI.CreateVariableSizedObject(Align(16), nullptr);1005 }1006 1007 restorePhyRegSave(RegSave, MBB, MBBI, DL);1008}1009 1010struct FromState {1011 ZAState From;1012 1013 constexpr uint8_t to(ZAState To) const {1014 static_assert(NUM_ZA_STATE < 16, "expected ZAState to fit in 4-bits");1015 return uint8_t(From) << 4 | uint8_t(To);1016 }1017};1018 1019constexpr FromState transitionFrom(ZAState From) { return FromState{From}; }1020 1021void MachineSMEABI::emitStateChange(EmitContext &Context,1022 MachineBasicBlock &MBB,1023 MachineBasicBlock::iterator InsertPt,1024 ZAState From, ZAState To,1025 LiveRegs PhysLiveRegs) {1026 // ZA not used.1027 if (From == ZAState::ANY || To == ZAState::ANY)1028 return;1029 1030 // If we're exiting from the ENTRY state that means that the function has not1031 // used ZA, so in the case of private ZA/ZT0 functions we can omit any set up.1032 if (From == ZAState::ENTRY && To == ZAState::OFF)1033 return;1034 1035 // TODO: Avoid setting up the save buffer if there's no transition to1036 // LOCAL_SAVED.1037 if (From == ZAState::ENTRY) {1038 assert(&MBB == &MBB.getParent()->front() &&1039 "ENTRY state only valid in entry block");1040 emitSMEPrologue(MBB, MBB.getFirstNonPHI());1041 if (To == ZAState::ACTIVE)1042 return; // Nothing more to do (ZA is active after the prologue).1043 1044 // Note: "emitNewZAPrologue" zeros ZA, so we may need to setup a lazy save1045 // if "To" is "ZAState::LOCAL_SAVED". It may be possible to improve this1046 // case by changing the placement of the zero instruction.1047 From = ZAState::ACTIVE;1048 }1049 1050 SMEAttrs SMEFnAttrs = AFI->getSMEFnAttrs();1051 bool IsAgnosticZA = SMEFnAttrs.hasAgnosticZAInterface();1052 bool HasZT0State = SMEFnAttrs.hasZT0State();1053 bool HasZAState = IsAgnosticZA || SMEFnAttrs.hasZAState();1054 1055 switch (transitionFrom(From).to(To)) {1056 // This section handles: ACTIVE <-> ACTIVE_ZT0_SAVED1057 case transitionFrom(ZAState::ACTIVE).to(ZAState::ACTIVE_ZT0_SAVED):1058 emitZT0SaveRestore(Context, MBB, InsertPt, /*IsSave=*/true);1059 break;1060 case transitionFrom(ZAState::ACTIVE_ZT0_SAVED).to(ZAState::ACTIVE):1061 emitZT0SaveRestore(Context, MBB, InsertPt, /*IsSave=*/false);1062 break;1063 1064 // This section handles: ACTIVE[_ZT0_SAVED] -> LOCAL_SAVED1065 case transitionFrom(ZAState::ACTIVE).to(ZAState::LOCAL_SAVED):1066 case transitionFrom(ZAState::ACTIVE_ZT0_SAVED).to(ZAState::LOCAL_SAVED):1067 if (HasZT0State && From == ZAState::ACTIVE)1068 emitZT0SaveRestore(Context, MBB, InsertPt, /*IsSave=*/true);1069 if (HasZAState)1070 emitZASave(Context, MBB, InsertPt, PhysLiveRegs);1071 break;1072 1073 // This section handles: ACTIVE -> LOCAL_COMMITTED1074 case transitionFrom(ZAState::ACTIVE).to(ZAState::LOCAL_COMMITTED):1075 // TODO: We could support ZA state here, but this transition is currently1076 // only possible when we _don't_ have ZA state.1077 assert(HasZT0State && !HasZAState && "Expect to only have ZT0 state.");1078 emitZT0SaveRestore(Context, MBB, InsertPt, /*IsSave=*/true);1079 emitZAMode(MBB, InsertPt, /*ClearTPIDR2=*/false, /*On=*/false);1080 break;1081 1082 // This section handles: LOCAL_COMMITTED -> (OFF|LOCAL_SAVED)1083 case transitionFrom(ZAState::LOCAL_COMMITTED).to(ZAState::OFF):1084 case transitionFrom(ZAState::LOCAL_COMMITTED).to(ZAState::LOCAL_SAVED):1085 // These transistions are a no-op.1086 break;1087 1088 // This section handles: LOCAL_(SAVED|COMMITTED) -> ACTIVE[_ZT0_SAVED]1089 case transitionFrom(ZAState::LOCAL_COMMITTED).to(ZAState::ACTIVE):1090 case transitionFrom(ZAState::LOCAL_COMMITTED).to(ZAState::ACTIVE_ZT0_SAVED):1091 case transitionFrom(ZAState::LOCAL_SAVED).to(ZAState::ACTIVE):1092 if (HasZAState)1093 emitZARestore(Context, MBB, InsertPt, PhysLiveRegs);1094 else1095 emitZAMode(MBB, InsertPt, /*ClearTPIDR2=*/false, /*On=*/true);1096 if (HasZT0State && To == ZAState::ACTIVE)1097 emitZT0SaveRestore(Context, MBB, InsertPt, /*IsSave=*/false);1098 break;1099 1100 // This section handles transistions to OFF (not previously covered)1101 case transitionFrom(ZAState::ACTIVE).to(ZAState::OFF):1102 case transitionFrom(ZAState::ACTIVE_ZT0_SAVED).to(ZAState::OFF):1103 case transitionFrom(ZAState::LOCAL_SAVED).to(ZAState::OFF):1104 assert(SMEFnAttrs.hasPrivateZAInterface() &&1105 "Did not expect to turn ZA off in shared/agnostic ZA function");1106 emitZAMode(MBB, InsertPt, /*ClearTPIDR2=*/From == ZAState::LOCAL_SAVED,1107 /*On=*/false);1108 break;1109 1110 default:1111 dbgs() << "Error: Transition from " << getZAStateString(From) << " to "1112 << getZAStateString(To) << '\n';1113 llvm_unreachable("Unimplemented state transition");1114 }1115}1116 1117} // end anonymous namespace1118 1119INITIALIZE_PASS(MachineSMEABI, "aarch64-machine-sme-abi", "Machine SME ABI",1120 false, false)1121 1122bool MachineSMEABI::runOnMachineFunction(MachineFunction &MF) {1123 if (!MF.getSubtarget<AArch64Subtarget>().hasSME())1124 return false;1125 1126 AFI = MF.getInfo<AArch64FunctionInfo>();1127 SMEAttrs SMEFnAttrs = AFI->getSMEFnAttrs();1128 if (!SMEFnAttrs.hasZAState() && !SMEFnAttrs.hasZT0State() &&1129 !SMEFnAttrs.hasAgnosticZAInterface())1130 return false;1131 1132 assert(MF.getRegInfo().isSSA() && "Expected to be run on SSA form!");1133 1134 this->MF = &MF;1135 Subtarget = &MF.getSubtarget<AArch64Subtarget>();1136 TII = Subtarget->getInstrInfo();1137 TRI = Subtarget->getRegisterInfo();1138 MRI = &MF.getRegInfo();1139 1140 const EdgeBundles &Bundles =1141 getAnalysis<EdgeBundlesWrapperLegacy>().getEdgeBundles();1142 1143 FunctionInfo FnInfo = collectNeededZAStates(SMEFnAttrs);1144 1145 if (OptLevel != CodeGenOptLevel::None) {1146 // Propagate desired states forward, then backwards. Most of the propagation1147 // should be done in the forward step, and backwards propagation is then1148 // used to fill in the gaps. Note: Doing both in one step can give poor1149 // results. For example, consider this subgraph:1150 //1151 // ┌─────┐1152 // ┌─┤ BB0 ◄───┐1153 // │ └─┬───┘ │1154 // │ ┌─▼───◄──┐│1155 // │ │ BB1 │ ││1156 // │ └─┬┬──┘ ││1157 // │ │└─────┘│1158 // │ ┌─▼───┐ │1159 // │ │ BB2 ├───┘1160 // │ └─┬───┘1161 // │ ┌─▼───┐1162 // └─► BB3 │1163 // └─────┘1164 //1165 // If:1166 // - "BB0" and "BB2" (outer loop) has no state preference1167 // - "BB1" (inner loop) desires the ACTIVE state on entry/exit1168 // - "BB3" desires the LOCAL_SAVED state on entry1169 //1170 // If we propagate forwards first, ACTIVE is propagated from BB1 to BB2,1171 // then from BB2 to BB0. Which results in the inner and outer loops having1172 // the "ACTIVE" state. This avoids any state changes in the loops.1173 //1174 // If we propagate backwards first, we _could_ propagate LOCAL_SAVED from1175 // BB3 to BB0, which would result in a transition from ACTIVE -> LOCAL_SAVED1176 // in the outer loop.1177 for (bool Forwards : {true, false})1178 propagateDesiredStates(FnInfo, Forwards);1179 }1180 1181 SmallVector<ZAState> BundleStates = assignBundleZAStates(Bundles, FnInfo);1182 1183 EmitContext Context;1184 insertStateChanges(Context, FnInfo, Bundles, BundleStates);1185 1186 if (Context.needsSaveBuffer()) {1187 if (FnInfo.AfterSMEProloguePt) {1188 // Note: With inline stack probes the AfterSMEProloguePt may not be in the1189 // entry block (due to the probing loop).1190 MachineBasicBlock::iterator MBBI = *FnInfo.AfterSMEProloguePt;1191 emitAllocateZASaveBuffer(Context, *MBBI->getParent(), MBBI,1192 FnInfo.PhysLiveRegsAfterSMEPrologue);1193 } else {1194 MachineBasicBlock &EntryBlock = MF.front();1195 emitAllocateZASaveBuffer(1196 Context, EntryBlock, EntryBlock.getFirstNonPHI(),1197 FnInfo.Blocks[EntryBlock.getNumber()].PhysLiveRegsAtEntry);1198 }1199 }1200 1201 return true;1202}1203 1204FunctionPass *llvm::createMachineSMEABIPass(CodeGenOptLevel OptLevel) {1205 return new MachineSMEABI(OptLevel);1206}1207