Preserve dormant SME ZA across the AArch64 stack switch.

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.

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

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`
 

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.

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