[SME] Add ZA-compatible interface and routines to save/restore SME state #264

sdesmalen-arm · 2024-06-07T11:09:12Z

This implements requests to add a new "ZA-compatible" interface which can be called with ZA state being either 'off', 'active' or 'dormant', and which can preserve any and all state enabled under PSTATE.ZA.

…ate. This implements requests to add a new "ZA-compatible" interface which can be called with ZA state being either 'off', 'active' or 'dormant', and which can preserve any and all state enabled under PSTATE.ZA.

rsandifo-arm · 2024-08-05T14:44:12Z

aapcs64/aapcs64.rst


 .. _`private-ZA`:
 .. _`shared-ZA`:
+.. _`compatible-ZA`:


Can we call this ZA-agnostic, for compatibility with the ACLE PR?

rsandifo-arm · 2024-08-05T14:45:46Z

aapcs64/aapcs64.rst

-Subroutines with a `private-ZA`_ interface and subroutines with a `shared-ZA`_
-interface can both (at their option) choose to guarantee that they
+The compatible-ZA interface is intended to be called from any function
+without requiring a change to PSTATE.ZA and is generally used in conjunction


I think we should make this stronger, and say that the function must preserve all state associated with PSTATE.ZA.

rsandifo-arm · 2024-08-05T14:46:17Z

aapcs64/aapcs64.rst

+with the expectation that it `preserves ZA`_.
+
+Subroutines with a `private-ZA`_ interface, `shared-ZA`_ interface or
+`compatible-ZA`_ interface can (at their option) choose to guarantee that they


…and then leave this paragraph as it was originally.

rsandifo-arm · 2024-08-05T14:48:14Z

aapcs64/aapcs64.rst

@@ -2069,6 +2077,8 @@ support routines:

   * the current values of TPIDR2_EL0, PSTATE.SM and PSTATE.ZA.

+   * the registers enabled by PSTATE.ZA, if PSTATE.ZA is enabled.


What is this change referring to? It doesn't look like the patch changes __arm_sme_state.

Sorry this should not have been committed. I've removed the change.

rsandifo-arm · 2024-08-05T14:50:24Z

aapcs64/aapcs64.rst

+  The size is guaranteed to be a multiple of 16.
+
+  The layout that corresponds to the calculated size is unspecified,
+  but the assumption is that the size always matches the implementation


I think this is a requirement rather than an assumption.

rsandifo-arm · 2024-08-05T14:57:04Z

aapcs64/aapcs64.rst

+    * If bit 1 of ``OPTIONS`` is 0 and ZT0 is available, then the subroutine
+      stores the full contents of ZT0 to ``PTR->ZT0``.
+
+    * If bit 63 of ``OPTIONS`` is 0, then the subroutine disables PSTATE.ZA.


This means that it isn't compatible-ZA according to the above table, since ZA is in a different state on return.

Also, bit 31 == 0 && bit 2 == 1 would be an invalid combination, since we would potentially have a nonnull (unsaved) TPIDR2_EL0 while PSTATE.ZA==0.

rsandifo-arm · 2024-08-05T14:58:41Z

aapcs64/aapcs64.rst

+    +--------+-----------------------------------------+
+    | bits   | Options                                 |
+    +========+=========================================+
+    | 63     | Preserve ZA using lazy save mechanism   |


Is it worth the complication of supporting two ways of saving ZA? Couldn't we start with the lazy version, and add the non-lazy version as a variation later if necessary?

rsandifo-arm · 2024-08-05T14:59:49Z

aapcs64/aapcs64.rst

+    | 62 - 3 | Zero for this revision of the AAPCS64,  |
+    |        | but reserved for future expansion       |
+    +--------+-----------------------------------------+
+    | 2      | Exclude TPIDR2_EL0                      |


Similarly, I'm not sure it's worth having a separate toggle for this and ZA.

rsandifo-arm · 2024-08-05T15:03:34Z

aapcs64/aapcs64.rst

+
+    * both bit 0 and bit 63 of ``OPTIONS`` are 1.
+
+  * If PSTATE.ZA is 0, then the subroutine enables PSTATE.ZA.


It should only do this if PSTATE.ZA was 1 when __arm_sme_save was called. Similarly for the restorations below.

rsandifo-arm · 2024-08-05T15:05:10Z

aapcs64/aapcs64.rst

+    ``PTR->TPIDR2_EL0`` at unspecified offsets in the buffer pointed to by PTR:
+
+    * If bit 63 of ``OPTIONS`` is 1 and TPIDR2_EL0 is null, then the function
+      copies ``PTR->BLK`` to X0 and calls ``__arm_tpidr2_restore``.


Suggested change

copies ``PTR->BLK`` to X0 and calls ``__arm_tpidr2_restore``.

points X0 at ``PTR->BLK`` and calls ``__arm_tpidr2_restore``.

* Renamed `compatible ZA` -> `agnostic ZA` * Changed __arm_sme_save/restore such that the save routine should record whether ZA or ZT0 is saved and such that the restore routine should check whether ZA or ZT0 was saved. This removes the (previously) implicit assumption in `__arm_sme-save` that PSTATE.ZA must be 1 if PTR is not nullptr. * ZA is now always saved/restored using the lazy-save mechanism. * Changed __arm_sme_save/restore to have a custom ZA interface instead of 'agnostic-ZA' which was incorrect.

rsandifo-arm

It isn't clear to me whether we expect callers of __arm_sme_state_size to check for zero returns, and skip the save and restore calls for that case. If so (“case A”), I think this makes things unnecessarily complex for the callers. If not (“case B”), the buffer should have nonzero size even if PSTATE.ZA==0, so that it can hold at least the SAVED_ZA and SAVED_ZT0 fields.

And for case B, the save and restore functions should be callable even in threads that don't have access to SME.

In a similar vein, how about making __arm_sme_restore decide what to do based only on the SAVED_* fields, rather than passing OPTIONS again? The caller of the save and restore functions is expected to be the same, so it should know what needs to be preserved.

I wonder how useful the OPTIONS will be in practice, given that __arm_sme_save must leave PSTATE.ZA in a dormant or off state. It seems unlikely that a caller would know enough about its position in the call hierarchy/overall program to know that certain state doesn't need to be preserved, but at the same time not know enough for it to have a specific __arm_in/out/inout/preserves annotation. No objection to keeping OPTIONS if you prefer though.

rsandifo-arm · 2024-08-22T09:19:55Z

aapcs64/aapcs64.rst

+  as well as any other state required for `__arm_sme_save`_ and
+  `__arm_sme_restore`_.
+
+  `__arm_sme_state_size`_ assumes that ZA is saved lazily and will account


I think this is potentially confusing, since it isn't clear who is assumed to do the lazy saving. I think what this paragraph is saying comes under the previous sentence — “as well as any other state required for __arm_sme_save_ and __arm_sme_restore_” — and so it might be better just to delete this.

rsandifo-arm · 2024-08-22T09:51:03Z

aapcs64/aapcs64.rst

+    ``PTR->BLK``, ``PTR->ZA`` and ``PTR->TPIDR2_EL0`` at unspecified offsets
+    in the buffer pointed to by ``PTR``:
+
+    * The full contents of ``TPIDR2_EL0`` are written to ``PTR->TPIDR2_EL0``.


Suggested change

* The full contents of ``TPIDR2_EL0`` are written to ``PTR->TPIDR2_EL0``.

* The contents of ``TPIDR2_EL0`` are written to ``PTR->TPIDR2_EL0``.

Using “full” sounded odd in reference to a system register that contains a pointer.

rsandifo-arm · 2024-08-22T09:57:12Z

aapcs64/aapcs64.rst

+
+    * The value 1 is written to ``PTR->SAVED_ZA``.
+
+  * If bit 1 of ``OPTIONS`` is 0 and ZT0 is available, then for the addresses


I assume that if bit 0 is 1, the function should turn PSTATE.ZA off before returning. In other words, ZA must not be active on return from this function.

Good spot, you're right!

sdesmalen-arm · 2024-08-22T10:41:21Z

It isn't clear to me whether we expect callers of __arm_sme_state_size to check for zero returns, and skip the save and restore calls for that case. If so (“case A”), I think this makes things unnecessarily complex for the callers. If not (“case B”), the buffer should have nonzero size even if PSTATE.ZA==0, so that it can hold at least the SAVED_ZA and SAVED_ZT0 fields.

And for case B, the save and restore functions should be callable even in threads that don't have access to SME.

The initial thinking was "case A", i.e. the intrinsics are only called when there is something worth saving. But now that I've added SAVED_ZA and SAVED_ZT0, it probably makes sense to remove that assumption.

I wonder how useful the OPTIONS will be in practice, given that __arm_sme_save must leave PSTATE.ZA in a dormant or off state. It seems unlikely that a caller would know enough about its position in the call hierarchy/overall program to know that certain state doesn't need to be preserved, but at the same time not know enough for it to have a specific __arm_in/out/inout/preserves annotation. No objection to keeping OPTIONS if you prefer though.

My reasons for adding OPTIONS was to give an interface that allows the routines to be used more generically, e.g. to support implementing the following case more easily using these intrinsics:

void bar() __arm_inout("zt0");
void foo() __arm_inout("za", "zt0") { bar(); }

where only "za" needs preserving across the call, but not "zt0". This could be done using explicit allocas, calling the __arm_tpidr2_* routines, but this functionality may become more involved in the future should more state be added to PSTATE.ZA.

In a similar vein, how about making __arm_sme_restore decide what to do based only on the SAVED_* fields, rather than passing OPTIONS again? The caller of the save and restore functions is expected to be the same, so it should know what needs to be preserved.

The reasoning for that was that one might only want to restore part of the saved state at one point, and restore the other part at another point. The compiler or user writing asm could use these routines to save/restore any state, e.g.

void bar();
void use_zt0() __arm_inout("zt0");
void use_za() __arm_inout("za");

void foo() __arm_inout("za", "zt0") {
               // preserve both zt0 and za, could use the __arm_sme_save routine for this.
  bar();
               // reload only zt0
  use_zt0();
               // save only zt0 (za was already saved)
  bar();
               // restore za
  use_za();
               // restore zt0
}

rsandifo-arm · 2024-08-22T12:52:34Z

I wonder how useful the OPTIONS will be in practice, given that __arm_sme_save must leave PSTATE.ZA in a dormant or off state. It seems unlikely that a caller would know enough about its position in the call hierarchy/overall program to know that certain state doesn't need to be preserved, but at the same time not know enough for it to have a specific __arm_in/out/inout/preserves annotation. No objection to keeping OPTIONS if you prefer though.

My reasons for adding OPTIONS was to give an interface that allows the routines to be used more generically, e.g. to support implementing the following case more easily using these intrinsics:
void bar() __arm_inout("zt0");
void foo() __arm_inout("za", "zt0") { bar(); }
where only "za" needs preserving across the call, but not "zt0". This could be done using explicit allocas, calling the __arm_tpidr2_* routines, but this functionality may become more involved in the future should more state be added to PSTATE.ZA.

I don't think the new routines work for that case though. The routines rely on saving ZA lazily, whereas the above needs to save ZA eagerly. bar is entitled to assume that ZA is active rather than dormant on entry.

sdesmalen-arm · 2024-08-22T12:58:23Z

I wonder how useful the OPTIONS will be in practice, given that __arm_sme_save must leave PSTATE.ZA in a dormant or off state. It seems unlikely that a caller would know enough about its position in the call hierarchy/overall program to know that certain state doesn't need to be preserved, but at the same time not know enough for it to have a specific __arm_in/out/inout/preserves annotation. No objection to keeping OPTIONS if you prefer though.

My reasons for adding OPTIONS was to give an interface that allows the routines to be used more generically, e.g. to support implementing the following case more easily using these intrinsics:
void bar() __arm_inout("zt0");
void foo() __arm_inout("za", "zt0") { bar(); }
where only "za" needs preserving across the call, but not "zt0". This could be done using explicit allocas, calling the __arm_tpidr2_* routines, but this functionality may become more involved in the future should more state be added to PSTATE.ZA.
I don't think the new routines work for that case though. The routines rely on saving ZA lazily, whereas the above needs to save ZA eagerly. bar is entitled to assume that ZA is active rather than dormant on entry.

Yes, that was the reason I initially added the "lazily" as an option to the struct. I figured it should still be possible to set up the lazy-save using these routines and then commit the lazy-save manually before the call to bar and then restore afterwards, if ZA needs saving eagerly. Or is there some subtlety that I'm missing?

rsandifo-arm · 2024-08-22T13:20:23Z

Yes, that was the reason I initially added the "lazily" as an option to the struct. I figured it should still be possible to set up the lazy-save using these routines and then commit the lazy-save manually before the call to bar and then restore afterwards, if ZA needs saving eagerly. Or is there some subtlety that I'm missing?

The difference for me is that the new routines are fundamentally about processing the live ZA state (and any future SME state) such that it's safe to call a normal function. In some cases this will involve turning PSTATE.ZA off. This is different from what the mixed-sharing example needs: there the intention is explicitly to keep ZA on (and active).

The new routines are also about providing an interface that can handle unknown future state. Thus the assumption is that extra state must be stored unless explicitly turned off (via OPTIONS). This is different from the example of moving between different sharing types, since there the compiler (by necessity) knows the state that is live in the caller and the state that is shared with the callee. The compiler only wants to store the difference between the two, and not any future unknown state.

When the callee shares some ZA state, but shares less state than the caller, the caller has to do an eager save and restore of the remaining state. I think we should just treat this as being conceptually like any other caller save. Admittedly the save and restore are quite heavy operations for ZA, but in concept they're not much different from a register save and restore. Treating them like that also helps with your later example involving calls to separate __arm_inout("zt0") and __arm_inout("za") functions.

@rsandifo-arm

As @rsandifo-arm explained, the same routines cannot be used for general use of saving/restoring of state enabled by PSTATE.ZA because of the different expectations on the ZA interface when partially saving/restoring state enabled by PSTATE.ZA. Removing the option, drastically simplifies the logic. Note that I have removed the two booleans (internal to the save/restore routines) that distinguish between having saved 'ZA' and 'ZT0', and replaced that with a single 'VALID' bit, because I think we can assume that the save/restore routines are called by PEs that have the same SME state.

rsandifo-arm

Sorry for only raising this now, but I wonder if we should reconsider the case in which ZA is dormant on entry to an agnostic-ZA function. Rather than say that agnostic-ZA functions must preserve the state (and so leave ZA dormant), should we instead say that agnostic-ZA behaves like private-ZA for that case? No other PSTATE.ZA state can be live if ZA is dormant. In particular, if ZA is dormant then ZT0 must be dead.

In other words, agnostic-ZA is like private-ZA except that ZA can be active on entry. If ZA is active on entry, it must be unchanged on return.

(If in future there is any other SME state that needs to be preserved, and that is not controlled by PSTATE.ZA, then that would still need to be saved and restored.)

The dormant-on-entry case isn't too interesting for direct calls between streaming code and agnostic-ZA functions, since the streaming code wouldn't be expected to set up a lazy save in that case. But perhaps it is more relevant if streaming code calls a normal function and that normal function calls a private-ZA function. (Ideally, the normal function would be changed to private-ZA, but that might not always be possible for ABI reasons.)

Other than that, it looks good to me. Some minor comments below.

rsandifo-arm · 2024-08-27T11:19:50Z

aapcs64/aapcs64.rst

+  * For addresses ``PTR->BLK``, ``PTR->ZA`` and ``PTR->TPIDR2_EL0`` at
+    unspecified offsets in the buffer pointed to by ``PTR``:
+
+    * The subroutine aborts in some platform-specific manner if the current


We can remove this, since:

If the thread has access to SME then it must also have access to TPIDR2_EL0.

rsandifo-arm · 2024-08-27T11:24:08Z

aapcs64/aapcs64.rst

+
+    * The current thread does not have access to SME.
+
+    * The current thread does not have access to ``TPIDR2_EL0`` when PSTATE.ZA


Similarly, this can be deleted.

rsandifo-arm · 2024-08-27T11:44:03Z

aapcs64/aapcs64.rst

+  * X1-X15, X19-X29 and SP are call-preserved.
+  * Z0-Z31 are call-preserved.
+  * P0-P15 are call-preserved.
+  * the subroutine `preserves ZA`_.


Might be better to delete this, since “agnostic-ZA” has its own requirements about the state of ZA on return.

rsandifo-arm · 2024-08-27T11:47:49Z

aapcs64/aapcs64.rst

+
+* The custom ``ZA`` interface has the following properties:
+
+  * If ZA state is 'off' on entry and SME state needs restoring, then it is


active or dormant, since if ZA was dormant on entry to __arm_sme_save, and if the lazy save is later committed, we'll end up restoring the previous dormant state. Similarly for the final bullet point.

But I think we move this down immediately before “The subroutine behaves as follows:” and then just say something like:

The ZA state on normal return is the same as the ZA state on entry to the __arm_sme_save call that initialized PTR.

(Sentence feels a bit clunky though, so improvements definitely welcome.)

rsandifo-arm · 2024-08-27T11:55:10Z

aapcs64/aapcs64.rst

+* The subroutine takes the following arguments:
+
+  PTR
+    a 64-bit data pointer passed in X0 that points to a buffer which


How about instead saying:

a 64-bit data pointer in X0 that must previously have been passed to __arm_sme_save. The buffer that it points to must still contain the data written by that call.

?

rsandifo-arm · 2024-08-27T11:57:35Z

aapcs64/aapcs64.rst

+* The subroutine behaves as follows:
+
+  * If ``PTR`` does not point to a valid buffer with the required size, the
+    behaviour of calling this routine is undefined.


behavior (and similarly throughout).

This changes the 'dormant on entry' case to be 'no change'.

The specification of these routines can be found here: ARM-software/abi-aa#264

This implements the lowering of calls from agnostic-ZA functions to non-agnostic-ZA functions, using the ABI routines `__arm_sme_state_size`, `__arm_sme_save` and `__arm_sme_restore`. This implements the proposal described in the following PRs: * ARM-software/acle#336 * ARM-software/abi-aa#264

smithp35 requested a review from rsandifo-arm June 10, 2024 10:02

rsandifo-arm reviewed Aug 5, 2024

View reviewed changes

stuij requested a review from sallyarmneale August 12, 2024 13:01

sdesmalen-arm force-pushed the sme-za-compatible branch from efe901f to 6f1bba8 Compare August 22, 2024 08:56

Remove redundant change

34eeecc

rsandifo-arm reviewed Aug 22, 2024

View reviewed changes

rsandifo-arm reviewed Aug 27, 2024

View reviewed changes

sdesmalen-arm added 2 commits August 28, 2024 09:15

Address review comments.

447378a

This changes the 'dormant on entry' case to be 'no change'.

Add section about dynamic symbols to test for ZT0

82c84bd

sdesmalen-arm added a commit to sdesmalen-arm/llvm-project that referenced this pull request Sep 9, 2024

[Compiler-rt] Add AArch64 routines for __arm_agnostic("sme_za_state")

9cae4ef

The specification of these routines can be found here: ARM-software/abi-aa#264

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[SME] Add ZA-compatible interface and routines to save/restore SME state #264

[SME] Add ZA-compatible interface and routines to save/restore SME state #264

sdesmalen-arm commented Jun 7, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

sdesmalen-arm Aug 22, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm Aug 5, 2024

rsandifo-arm left a comment

rsandifo-arm Aug 22, 2024

rsandifo-arm Aug 22, 2024

rsandifo-arm Aug 22, 2024

sdesmalen-arm Aug 22, 2024

sdesmalen-arm commented Aug 22, 2024

rsandifo-arm commented Aug 22, 2024

sdesmalen-arm commented Aug 22, 2024

rsandifo-arm commented Aug 22, 2024

rsandifo-arm left a comment

rsandifo-arm Aug 27, 2024

rsandifo-arm Aug 27, 2024

rsandifo-arm Aug 27, 2024

rsandifo-arm Aug 27, 2024

rsandifo-arm Aug 27, 2024

rsandifo-arm Aug 27, 2024

		@@ -2069,6 +2077,8 @@ support routines:

		* the current values of TPIDR2_EL0, PSTATE.SM and PSTATE.ZA.

		* the registers enabled by PSTATE.ZA, if PSTATE.ZA is enabled.


		* both bit 0 and bit 63 of ``OPTIONS`` are 1.

		* If PSTATE.ZA is 0, then the subroutine enables PSTATE.ZA.

	copies ``PTR->BLK`` to X0 and calls ``__arm_tpidr2_restore``.
	points X0 at ``PTR->BLK`` and calls ``__arm_tpidr2_restore``.

	* The full contents of ``TPIDR2_EL0`` are written to ``PTR->TPIDR2_EL0``.
	* The contents of ``TPIDR2_EL0`` are written to ``PTR->TPIDR2_EL0``.


		* The value 1 is written to ``PTR->SAVED_ZA``.

		* If bit 1 of ``OPTIONS`` is 0 and ZT0 is available, then for the addresses


		* The current thread does not have access to SME.

		* The current thread does not have access to ``TPIDR2_EL0`` when PSTATE.ZA


		* The custom ``ZA`` interface has the following properties:

		* If ZA state is 'off' on entry and SME state needs restoring, then it is

[SME] Add ZA-compatible interface and routines to save/restore SME state #264

Are you sure you want to change the base?

[SME] Add ZA-compatible interface and routines to save/restore SME state #264

Conversation

sdesmalen-arm commented Jun 7, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rsandifo-arm left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

sdesmalen-arm commented Aug 22, 2024

rsandifo-arm commented Aug 22, 2024

sdesmalen-arm commented Aug 22, 2024

rsandifo-arm commented Aug 22, 2024

rsandifo-arm left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment