diff --git a/AGENTS.md b/AGENTS.md index 22f404e..bf7597d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,12 +66,14 @@ Limits that are easy to fix the wrong way: (`-mfloat-abi=softfp`) is not a supported platform. It still builds as `arm32` and runs, and that assembly does not preserve `FPSCR`. Do not add an `FPSCR` spill to `switch_arm_gcc.S` to cover it. -- `aarch64`, `darwin_arm64`, and `win_arm64` do not preserve SME ZA - across `stackman_switch`. Do not call `__arm_za_disable`. That helper - commits the outgoing caller's ZA and turns ZA off. Nothing in the - switch installs the incoming stack's ZA. `stackman_call` returns to - the same caller and does not need it. Windows ARM64 is a different - runtime. Do not call the AAPCS64 helper there. +- `aarch64` and `darwin_arm64` keep dormant or off SME ZA across + `stackman_switch`. The scheme is `docs/scalable-matrix-extension.md`. + `stackman_sme_suspend` runs before the save callback and + `stackman_sme_resume` runs after the new stack pointer is installed. + Do not restore ZA with `__arm_sme_save` / `__arm_sme_restore`. Do not + read `TPIDR2_EL0` with `mrs`. `stackman_call` does not touch ZA. + Windows ARM64 is a different runtime. Do not call the AAPCS64 helpers + there, and do not pretend `win_arm64` preserves ZA. - On Windows ARM64, `x18` is the TEB for the whole time user code runs. Spill and restore the saved value. Do not write the live register. Do not invent `.seh_*` names for this file. diff --git a/README.md b/README.md index 999c87e..57ca7cd 100644 --- a/README.md +++ b/README.md @@ -109,8 +109,9 @@ calling convention, plus archive format: Which registers each ABI spills is listed in [`docs/platform-abi.md`](docs/platform-abi.md). -The 64-bit ARM ABIs (`aarch64`, `darwin_arm64`, and `win_arm64`) do not -preserve SME ZA context across `stackman_switch`. +`aarch64` and `darwin_arm64` preserve a dormant or off SME ZA state +across `stackman_switch`. `win_arm64` does not. See +[`docs/scalable-matrix-extension.md`](docs/scalable-matrix-extension.md). Softfp (`-mfloat-abi=softfp`) is not a supported platform, and it is not clear how widely it is used. In practice that build does not define diff --git a/docs/platform-abi.md b/docs/platform-abi.md index 3d3ddb3..120dce8 100644 --- a/docs/platform-abi.md +++ b/docs/platform-abi.md @@ -104,15 +104,20 @@ for the same reason as `arm32`. of `v8`–`v15` (`d8`–`d15`). The high 64 bits of `v8`–`v15` are caller-saved and are not spilled. `x18` is left alone. -`FPCR` and `FPSR` are stored in a 16-byte slot at the pointer passed to -the callback and written back after `STACKMAN_OP_RESTORE`. `x29` is +`FPCR` and `FPSR` are stored in a 32-byte slot at the pointer passed to +the callback and written back after `STACKMAN_OP_RESTORE`. The slot is +`FPCR`, `FPSR`, the saved TPIDR2 block pointer, and 8 bytes of padding. +A null pointer means ZA was off on entry. A non-null pointer is a +dormant ZA state. `stackman_sme_suspend` commits it with +`__arm_za_disable` before the save callback, and `stackman_sme_resume` +restores it after the new stack pointer is installed, before the +restore callback. The helpers are weak. See +[`scalable-matrix-extension.md`](scalable-matrix-extension.md). `x29` is adjusted by the stack delta. The prologue sets `x29` from `sp`, so the delta is computed from the lowered stack pointer. Do not list `x29` as an asm clobber. -`stackman_call` does not spill `FPCR`, `FPSR`, or `d8`–`d15`. - -SME ZA is not preserved. Do not call `__arm_za_disable`. +`stackman_call` does not spill `FPCR`, `FPSR`, `d8`–`d15`, or ZA. ### `win_arm64` diff --git a/docs/scalable-matrix-extension.md b/docs/scalable-matrix-extension.md new file mode 100644 index 0000000..8098a3a --- /dev/null +++ b/docs/scalable-matrix-extension.md @@ -0,0 +1,143 @@ +# Scalable Matrix Extension + +SME is the Scalable Matrix Extension. It is an optional AArch64 feature, +`FEAT_SME`, published by Arm in 2021 as part of the Armv9-A architecture. +This note is how `stackman_switch` keeps its ZA state on `aarch64` and +`darwin_arm64`. +The register spills for the other ABIs stay in `platform-abi.md`. + +## Where it comes from + +SME is a CPU feature, not an operating-system feature and not a baseline +part of AArch64. A core either implements `FEAT_SME` or it does not. Armv8-A +machines do not have it. Plenty of Armv9-A machines do not have it either. +Cores that do include Apple M4 and later, and some Neoverse and Cortex +parts from that generation. The streaming vector length is chosen by the +implementation. ZA is one square matrix of that length: from 256 bytes at +the shortest length up to 64 KB at the longest. + +Software sees it through the Arm Procedure Call Standard for the 64-bit +Architecture (AAPCS64). That document defines the ZA lazy-save scheme and +the runtime helpers (`__arm_sme_state`, `__arm_za_disable`, +`__arm_tpidr2_restore`, and the others). The helpers live in the compiler +runtime, libgcc or compiler-rt, on toolchains that implement the scheme. +They ask the platform whether this thread may use SME, and they return +without touching ZA when it may not. A library built for plain `armv8-a` +can call them. + +Linux and macOS on AArch64 use that AAPCS64 runtime. In this tree those +are the `aarch64` and `darwin_arm64` ABIs, both built from +`switch_aarch64_gcc.S`. Windows on ARM64 does not provide these helpers +and does not use this lazy-save scheme. `win_arm64` is out of scope here. +`stackman_call` returns to the same caller, so the caller's buffer stays +on the live stack and this scheme is not required there. + +SVE, the Scalable Vector Extension, widens the vector registers. Advanced +SIMD is the 128-bit `v0`–`v31` file. ZA is neither of those. It is a +separate matrix accumulator with its own enable bit. + +## The three modes + +ZA is in exactly one of these modes. The names are the AAPCS64 ones. +`PSTATE.ZA` is the enable bit. `TPIDR2_EL0` is a system register that +either holds null or points at a 16-byte TPIDR2 block on some function's +stack. Byte 0 of that block is `za_save_buffer`, the address of the rows. +Bytes 8–9 are `num_za_save_slices`. The matrix is not stored at the +pointer itself. + +**Off.** `PSTATE.ZA` is 0 and `TPIDR2_EL0` is null. The matrix is disabled +and holds nothing that matters. `__arm_za_disable` leaves this mode +behind. This is the disabled state. + +**Active.** `PSTATE.ZA` is 1 and `TPIDR2_EL0` is null. The matrix is on +and the bytes sit in the hardware register. The function that is running +is using ZA directly. There is no buffer waiting to receive a copy. + +**Dormant.** `PSTATE.ZA` is 1 and `TPIDR2_EL0` points at a TPIDR2 block. +The matrix is still on, and the bytes are still in the hardware register. +The block says where to copy them if some other function needs the +register. The buffer is stale until that copy happens. Dormant exists +because the matrix is large. A normal call does not copy it. The owner +sets the pointer, makes the call, and on return either clears the pointer +(the callee never needed ZA, so the register still holds the bytes) or +finds the pointer already cleared (the callee committed a copy, and the +owner loads the buffer back). + +Active and dormant both have the matrix turned on. The difference is the +pointer. Off is the only mode in which the matrix is disabled. + +`stackman_switch` is an ordinary C function. The procedure call standard +classifies every such function as private-ZA. On entry, ZA is off or +dormant. It is not active. A caller that was using ZA switches from +active to dormant before the call, by pointing `TPIDR2_EL0` at a block in +its own frame. At the first instruction of `stackman_switch` the bytes are +still in the hardware register. + +## How a switch keeps that state + +The goal is to put back the mode `stackman_switch` found on entry, for +whichever stack is about to run. Off stays off. Dormant becomes dormant +again, with the same block pointer and the same bytes in the register. + +Do this with the AAPCS64 helpers, not with `__arm_sme_save` and +`__arm_sme_restore`. That second pair snapshots active ZA into a private +buffer and does nothing when `TPIDR2_EL0` is already set. The mode on +entry is dormant or off, so the pair would not see it. + +Saving the outgoing stack happens before the `STACKMAN_OP_SAVE` callback. +That callback is allowed to copy or abandon the stack, and a dormant +`TPIDR2_EL0` points into it. + +1. Call `__arm_sme_state`. It is safe on a machine without SME. The + TPIDR2 value comes back in `x1`. Do not read `TPIDR2_EL0` with `mrs`. + That register traps when the thread has no access to it. +2. If the pointer is null, the mode is off. Store a null word in this + stack's spill image. +3. If the pointer is non-null, the mode is dormant. Keep the pointer, + then call `__arm_za_disable`. The helper copies the live rows into + `za_save_buffer`, clears `TPIDR2_EL0`, and turns ZA off. Store the + saved pointer in the spill image. + +The spill slot has to sit at or above the stack pointer that the callback +sees, or a copy of the stack will not carry it. The published pointer +stays the `FPCR`/`FPSR` slot. The block pointer is the next 8 bytes, with +8 bytes of padding so the slot stays 16-byte aligned: + +| Offset from the published stack pointer | Contents | +| --- | --- | +| 0 | `FPCR` | +| 8 | `FPSR` | +| 16 | saved TPIDR2 block pointer, or 0 | +| 24 | 0 | + +Restoring the incoming stack happens after its stack pointer is installed +and before the `STACKMAN_OP_RESTORE` callback. Code on that stack then +sees the mode its own `stackman_switch` entry had. + +1. Load the saved pointer from that stack's slot. +2. A null pointer means off. Leave ZA disabled and leave `TPIDR2_EL0` + null. +3. A non-null pointer means dormant. `smstart za` turns the matrix on + without changing streaming mode (`PSTATE.SM` is already 0 on entry to + this non-streaming function). Call `__arm_tpidr2_restore` with the + saved block pointer. That helper requires `TPIDR2_EL0` to be null and + ZA to be on, and it only copies the rows. Then write the pointer back + to `TPIDR2_EL0`. The mode is dormant again. + +The same steps run when the callback does not switch. Entry was dormant, +the save committed it, and the restore puts dormant back before +`stackman_switch` returns. The caller's compiler then sees its own block +pointer still in `TPIDR2_EL0` and keeps the register contents. + +Link the three helpers as weak symbols. A toolchain that does not ship +them resolves the addresses to zero. In that case store a null pointer +and skip both calls, which is the behaviour this tree has today. If +`__arm_tpidr2_restore` is missing, do not call `__arm_za_disable` either: +committing without a way to load the buffer back would drop the rows. +Windows ARM64 does not get these calls. + +Two limits belong with the mechanism. The saved pointer and +`za_save_buffer` are absolute addresses. A callback that copies a stack +to a new base has to relocate them, the same way it already has to +relocate saved frame pointers. `ZT0`, the SME2 lookup table, is +caller-saved and is not part of this state. diff --git a/stackman/platforms/switch_aarch64_gcc.S b/stackman/platforms/switch_aarch64_gcc.S index 208610d..b3e894a 100644 --- a/stackman/platforms/switch_aarch64_gcc.S +++ b/stackman/platforms/switch_aarch64_gcc.S @@ -34,6 +34,14 @@ #define CFI_DEF_CFA_OFFSET(x) .cfi_def_cfa_offset x #define CFI_OFFSET(r,o) .cfi_offset r, o #define CFI_RESTORE(r) .cfi_restore r +#endif + +#ifdef __APPLE__ +#define SME_SUSPEND _stackman_sme_suspend +#define SME_RESUME _stackman_sme_resume +#else +#define SME_SUSPEND stackman_sme_suspend +#define SME_RESUME stackman_sme_resume #endif .arch armv8-a @@ -78,22 +86,32 @@ LABEL(stackman_switch) CFI_OFFSET(79,-8) mov x3, x0 mov x0, x1 + /* x19 and x20 are already spilled. Hold the callback and the + * context across stackman_sme_suspend. It is a normal call. + */ + mov x19, x3 + mov x20, x0 + bl SME_SUSPEND + mov x4, x0 + mov x0, x20 + mov x3, x19 #APP // 70 "../platforms/switch_aarch64_gcc.h" 1 mrs x1, fpcr mrs x2, fpsr - stp x1, x2, [sp, #-16]! + stp x1, x2, [sp, #-32]! + stp x4, xzr, [sp, #16] // 0 "" 2 // 85 "../platforms/switch_aarch64_gcc.h" 1 mov x2, sp // 0 "" 2 #NO_APP - /* The 16-byte FPCR/FPSR slot is outside the compiler frame. - * Keep it out of the #APP block: Darwin assembles that block, - * and these macros are empty there. A gcc -S regen does not - * emit this, because the compiler does not see the sp write. + /* FPCR/FPSR and the saved TPIDR2 pointer sit below the compiler + * frame. The published sp is the FPCR word. Darwin assembles the + * #APP block with no .cfi_startproc, so the CFA for this slot is + * set here. A gcc -S regen does not emit it. */ - CFI_DEF_CFA_OFFSET(176) + CFI_DEF_CFA_OFFSET(192) mov w1, 0 mov x20, x0 mov x19, x3 @@ -106,12 +124,25 @@ LABEL(stackman_switch) sub x29, x29, x1 // 0 "" 2 #NO_APP + /* Null means ZA was off on entry to this stack's switch. + * A pointer means dormant: turn ZA on, copy the rows back, and + * point TPIDR2_EL0 at the same block. x19 and x20 stay live. + */ + ldr x0, [sp, #16] + cbz x0, .Lsme_resume_skip + bl SME_RESUME +.Lsme_resume_skip: + /* stackman_sme_resume is a normal call, so x2 does not + * survive it. The restore callback's third argument is the + * stack just installed. Reload that on both paths. + */ + mov x2, sp mov w1, 1 mov x0, x20 blr x19 #APP // 102 "../platforms/switch_aarch64_gcc.h" 1 - ldp x1, x2, [sp], #16 + ldp x1, x2, [sp], #32 msr fpcr, x1 msr fpsr, x2 // 0 "" 2 diff --git a/stackman/platforms/switch_aarch64_gcc.h b/stackman/platforms/switch_aarch64_gcc.h index 9058c28..957a91c 100644 --- a/stackman/platforms/switch_aarch64_gcc.h +++ b/stackman/platforms/switch_aarch64_gcc.h @@ -8,6 +8,11 @@ * including the caller-saved status bits. stackman_call does not * change them and does not spill them. * + * The same slot carries the caller's TPIDR2 block pointer so a dormant + * SME ZA state can be committed on the way out and restored on the way + * back in. See docs/scalable-matrix-extension.md. stackman_call does + * not touch it: that call returns to the same caller. + * */ /* clang cannot perform inline assembly using specific __attr__ * instructions, and so it may use a base pointer and other @@ -25,6 +30,73 @@ #endif #ifdef STACKMAN_SWITCH_IMPL + +#if !__ASSEMBLER__ +/* + * Called from switch_aarch64_gcc.S and from the inline copy below. + * Weak on purpose: a toolchain without the AAPCS64 SME runtime resolves + * these to zero, and the switch stores a null pointer. + * + * target("arch=armv8-a") keeps an SME-aware compiler from wrapping the + * helpers in a lazy save of its own. smstart za is encoded with .inst so + * this file still assembles for plain armv8-a. 0xd503457f is "smstart za". + * S3_3_C13_C0_5 is TPIDR2_EL0. + */ +extern void __arm_za_disable(void) __attribute__((weak)); +extern void __arm_tpidr2_restore(void *) __attribute__((weak)); + +typedef struct { + unsigned long long status; + void *tpidr2; +} stackman_sme_state_result; + +extern stackman_sme_state_result __arm_sme_state(void) __attribute__((weak)); + +__attribute__((visibility("hidden"), noinline, target("arch=armv8-a"))) +void *stackman_sme_suspend(void); + +__attribute__((visibility("hidden"), noinline, target("arch=armv8-a"))) +void stackman_sme_resume(void *tpidr2); + +void *stackman_sme_suspend(void) +{ + stackman_sme_state_result state; + + if (!__arm_sme_state || !__arm_za_disable || !__arm_tpidr2_restore) + return 0; + state = __arm_sme_state(); + if (!state.tpidr2) + return 0; + __arm_za_disable(); + return state.tpidr2; +} + +void stackman_sme_resume(void *tpidr2) +{ + if (!tpidr2 || !__arm_tpidr2_restore) + return; + { + /* x0 survives __arm_tpidr2_restore. The helper preserves it. */ + register void *blk asm("x0") = tpidr2; + register void (*fn)(void *) asm("x16") = __arm_tpidr2_restore; + + __asm__ volatile( + ".inst 0xd503457f\n\t" + "blr %1\n\t" + "msr S3_3_C13_C0_5, %0\n\t" + : "+r"(blk) + : "r"(fn) + : "x14", "x15", "x17", "x30", "memory", "cc"); + } +} + +#if defined(__APPLE__) +#define STACKMAN_SME_RESUME "bl _stackman_sme_resume\n\t" +#else +#define STACKMAN_SME_RESUME "bl stackman_sme_resume\n\t" +#endif +#endif /* !__ASSEMBLER__ */ + #if !__ASSEMBLER__ && !defined(STACKMAN_ASSEMBLY_SRC) /* @@ -55,23 +127,26 @@ STACKMAN_LINKAGE_SWITCH_INASM void *STACKMAN_SWITCH_INASM_NAME(stackman_cb_t callback, void *context) { void *sp; + void *tpidr2; /* have gcc save/restore registers for us on the stack */ __asm__ volatile ("" : : : GP_PRESERVE); #ifdef FP_PRESERVE __asm__ volatile ("" : : : FP_PRESERVE); #endif /* FPCR and FPSR are one thread context on Linux and macOS. - * Park both in a 16-byte slot and pass that sp to the callback, - * so the copied image carries them. Sixteen bytes keeps the - * AAPCS64 alignment. x1 and x2 are free until the compiler - * assigns the saved stack pointer. The reload clobbers the - * same registers after the restore callback returns. + * Park both in the bottom of a 32-byte slot and pass that sp + * to the callback, so the copied image carries them. The next + * word is the TPIDR2 block pointer just captured, or null when + * ZA was off or the SME helpers are absent. x1 and x2 are free + * until the compiler assigns the saved stack pointer. */ + tpidr2 = stackman_sme_suspend(); __asm__ volatile ( "mrs x1, fpcr\n\t" "mrs x2, fpsr\n\t" - "stp x1, x2, [sp, #-16]!" - ::: "x1", "x2", "memory"); + "stp x1, x2, [sp, #-32]!\n\t" + "stp %0, xzr, [sp, #16]" + :: "r"(tpidr2) : "x1", "x2", "memory"); /* The compiler's CFI still describes its own frame. This * adjusts by the slot just pushed. It is not in the stp asm: * that string is copied into the shared .S, and Darwin @@ -79,33 +154,45 @@ void *STACKMAN_SWITCH_INASM_NAME(stackman_cb_t callback, void *context) * CFA with CFI_DEF_CFA_OFFSET instead. */ #if !defined(__APPLE__) - __asm__ volatile (".cfi_adjust_cfa_offset 16"); + __asm__ volatile (".cfi_adjust_cfa_offset 32"); #endif /* sp = get stack pointer from assembly code */ __asm__ ("mov %[result], sp" : [result] "=r" (sp)); /* store stack */ sp = callback(context, STACKMAN_OP_SAVE, sp); /* Keep the live frame pointer attached to the new stack. - * x1 is free until the compiler assigns the RESTORE opcode. * x29 is not an asm clobber: a clobber would make gcc save * and restore it around this block and undo the delta. + * Resume a dormant ZA before the callback runs on this stack. + * A null slot means the entry mode was off. The call is inside + * this asm so the compiler does not spill between the sp write + * and the restore. */ - __asm__ ( + __asm__ volatile ( "sub x1, sp, %[result]\n\t" "mov sp, %[result]\n\t" - "sub x29, x29, x1" - : : [result] "r" (sp) : "x1"); + "sub x29, x29, x1\n\t" + "ldr x0, [sp, #16]\n\t" + "cbz x0, .Lsme%=\n\t" + STACKMAN_SME_RESUME + ".Lsme%=:\n\t" + : : [result] "r" (sp) + : "x0", "x1", "x2", "x3", "x4", "x5", "x6", "x7", + "x8", "x9", "x10", "x11", "x12", "x13", "x14", "x15", + "x16", "x17", "x30", + "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", + "memory", "cc"); sp = callback(context, STACKMAN_OP_RESTORE, sp); /* Reload FPCR and FPSR from the stack now current, then release * the slot before the epilogue pops the callee-saved registers. */ __asm__ volatile ( - "ldp x1, x2, [sp], #16\n\t" + "ldp x1, x2, [sp], #32\n\t" "msr fpcr, x1\n\t" "msr fpsr, x2" ::: "x1", "x2", "memory"); #if !defined(__APPLE__) - __asm__ volatile (".cfi_adjust_cfa_offset -16"); + __asm__ volatile (".cfi_adjust_cfa_offset -32"); #endif return sp; }