[gcc r16-9318] aarch64: Fix ZA state transition [PR119210]

https://gcc.gnu.org/g:c7c77ec3190e37cc079ef1a88521af485d7224be

commit r16-9318-gc7c77ec3190e37cc079ef1a88521af485d7224be
Author: Alice Carlotti <[email protected]>
Date:   Thu Apr 16 13:02:35 2026 +0100

    aarch64: Fix ZA state transition [PR119210]
    
    In the INACTIVE_CALLER -> INACTIVE LOCAL transition, ensure ZA is active
    and zeroed before setting tpidr2_el0.
    
    gcc/ChangeLog:
    
            PR target/119210
            * config/aarch64/aarch64.cc (aarch64_mode_emit_local_sme_state):
            Add PSTATE.ZA enablement, and zero it if already enabled.
    
    gcc/testsuite/ChangeLog:
    
            PR target/119210
            * gcc.target/aarch64/sme/za_state_8.c: New test.

Diff:
---
 gcc/config/aarch64/aarch64.cc                     | 13 +++++++++++-
 gcc/testsuite/gcc.target/aarch64/sme/za_state_8.c | 25 +++++++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/gcc/config/aarch64/aarch64.cc b/gcc/config/aarch64/aarch64.cc
index 0b9cb043c665..99890a463d4c 100644
--- a/gcc/config/aarch64/aarch64.cc
+++ b/gcc/config/aarch64/aarch64.cc
@@ -31984,7 +31984,8 @@ aarch64_mode_emit_local_sme_state (aarch64_local_sme_state mode,
       emit_insn (gen_aarch64_tpidr2_save ());
       emit_insn (gen_aarch64_clear_tpidr2 ());
       if (mode == aarch64_local_sme_state::ACTIVE_LIVE
-	  || mode == aarch64_local_sme_state::ACTIVE_DEAD)
+	  || mode == aarch64_local_sme_state::ACTIVE_DEAD
+	  || mode == aarch64_local_sme_state::INACTIVE_LOCAL)
 	{
 	  if (aarch64_cfun_has_state ("za"))
 	    emit_insn (gen_aarch64_initial_zero_za ());
@@ -32066,6 +32067,16 @@ aarch64_mode_emit_local_sme_state (aarch64_local_sme_state mode,
 
   if (mode == aarch64_local_sme_state::INACTIVE_LOCAL)
     {
+      if (prev_mode == aarch64_local_sme_state::INACTIVE_CALLER)
+	/* Enable ZA (if it wasn't already enabled on entry).  Enabling ZA has
+	   the side-effect of zeroing ZA.
+
+	   A functionally correct alternative would be to leave TPIDR2_EL0 null
+	   and zero the save buffer.  However, zeroing the save buffer would require
+	   more code and would optimize for the case in which a callee also
+	   initialises private ZA state (which should be a rare event).  */
+	emit_insn (gen_aarch64_smstart_za ());
+
       if (prev_mode == aarch64_local_sme_state::ACTIVE_LIVE
 	  || prev_mode == aarch64_local_sme_state::ACTIVE_DEAD
 	  || prev_mode == aarch64_local_sme_state::INACTIVE_CALLER)
diff --git a/gcc/testsuite/gcc.target/aarch64/sme/za_state_8.c b/gcc/testsuite/gcc.target/aarch64/sme/za_state_8.c
new file mode 100644
index 000000000000..9b7a6ffa69cb
--- /dev/null
+++ b/gcc/testsuite/gcc.target/aarch64/sme/za_state_8.c
@@ -0,0 +1,25 @@
+// { dg-options "-O -fomit-frame-pointer -fno-optimize-sibling-calls" }
+// { dg-final { check-function-bodies "**" "" } }
+
+#include <arm_sme.h>
+
+void callee_ns();
+__arm_streaming __arm_inout("za") void callee_s();
+
+/*
+** foo:
+**	...
+**	smstart	za
+**	...
+**	msr	tpidr2_el0, x\d+
+**	...
+*/
+__arm_locally_streaming __arm_new("za") const float * foo(const float* x) {
+    callee_ns ();
+    const float32_t *x_f_in = x;
+    svzero_za();
+    callee_s ();
+    return x_f_in;
+}
+
+