[PATCH 13/20] arm64/fpsimd: Make clone() compatible with ZA lazy saving

[Yury Khrustalev]

On Tue, May 06, 2025 at 04:25:16PM +0100, Mark Rutland wrote:
> Linux is intended to be compatible with userspace written to Arm's
> AAPCS64 procedure call standard [1,2]. For the Scalable Matrix Extension
> (SME), AAPCS64 was extended with a "ZA lazy saving scheme", where SME's
> ZA tile is lazily callee-saved and caller-restored. In this scheme,
> TPIDR2_EL0 indicates whether the ZA tile is live or has been saved by
> pointing to a "TPIDR2 block" in memory, which has a "za_save_buffer"
> pointer. This scheme has been implemented in GCC and LLVM, with
> necessary runtime support implemented in glibc and bionic.
> 
> AAPCS64 does not specify how the ZA lazy saving scheme is expected to
> interact with thread creation mechanisms such as fork() and
> pthread_create(), which would be implemented in terms of the Linux clone
> syscall. The behaviour implemented by Linux and glibc/bionic doesn't
> always compose safely, as explained below.
> 
> Currently the clone syscall is implemented such that PSTATE.ZA and the
> ZA tile are always inherited by the new task, and TPIDR2_EL0 is
> inherited unless the 'flags' argument includes CLONE_SETTLS,
> in which case TPIDR2_EL0 is set to 0/NULL. This doesn't make much sense:
> 
> (a) TPIDR2_EL0 is part of the calling convention, and changes as control
>     is passed between functions. It is *NOT* used for thread local
>     storage, despite superficial similarity to TPIDR_EL0, which is is
>     used as the TLS register.
> 
> (b) TPIDR2_EL0 and PSTATE.ZA are tightly coupled in the procedure call
>     standard, and some combinations of states are illegal. In general,
>     manipulating the two independently is not guaranteed to be safe.
> 
> In practice, code which is compliant with the procedure call standard
> may issue a clone syscall while in the "ZA dormant" state, where
> PSTATE.ZA==1 and TPIDR2_EL0 is non-null and indicates that ZA needs to
> be saved. This can cause a variety of problems, including:
> 
> * If the implementation of pthread_create() passes CLONE_SETTLS, the
>   new thread will start with PSTATE.ZA==1 and TPIDR2==NULL. Per the
>   procedure call standard this is not a legitimate state for most
>   functions. This can cause data corruption (e.g. as code may rely on
>   PSTATE.ZA being 0 to guarantee that an SMSTART ZA instruction will
>   zero the ZA tile contents), and may result in other undefined
>   behaviour.
> 
> * If the implementation of pthread_create() does not pass CLONE_SETTLS, the
>   new thread will start with PSTATE.ZA==1 and TPIDR2 pointing to a
>   TPIDR2 block on the parent thread's stack. This can result in a
>   variety of problems, e.g.
> 
>   - The child may write back to the parent's za_save_buffer, corrupting
>     its contents.
> 
>   - The child may read from the TPIDR2 block after the parent has reused
>     this memory for something else, and consequently the child may abort
>     or clobber arbitrary memory.
> 
> Ideally we'd require that userspace ensures that a task is in the "ZA
> off" state (with PSTATE.ZA==0 and TPIDR2_EL0==NULL) prior to issuing a
> clone syscall, and have the kernel force this state for new threads.
> Unfortunately, contemporary C libraries do not do this, and simply
> forcing this state within the implementation of clone would break
> fork().
> 
> Instead, we can bodge around this by considering the CLONE_VM flag, and
> mainpulate PSTATE.ZA and TPIDR2_EL0 as a pair. CLONE_VM indicates that
> the new task will run in the same address space as its parent, and in
> that case it doesn't make sense to inherit a stale pointer to the
> parent's TPIDR2 block:
> 
> * For fork(), CLONE_VM will not be set, and it is safe to inherit both
>   PSTATE.ZA and TPIDR2_EL0 as the new task will have its own copy of the
>   address space, and cannot clobber its parent's stack.
> 
> * For pthread_create() and vfork(), CLONE_VM will be set, and discarding
>   PSTATE.ZA and TPIDR2_EL0 for the new task doesn't break any existing
>   assumptions in userspace.

Acked-by: Yury Khrustalev <yury.khrustalev at arm.com>

Kind regards,
Yury