[compiler-rt] [compiler-rt][AArch64] Exit early from __arm_za_disable. (PR #174942)

[Sander de Smalen via llvm-commits]

https://github.com/sdesmalen-arm created https://github.com/llvm/llvm-project/pull/174942

Because `__arm_za_disable` is a private-ZA function, it's only ever entered with ZA state `off` or `dormant`. If the state is `off` then we can safely return and there is no need to call `__arm_tpidr2_save` or to explicitly set PSTATE.ZA or TPIDR2_EL0 to zero.

>From 2661315ca6879f950d0d4d374bc055e8e6cf9357 Mon Sep 17 00:00:00 2001
From: Sander de Smalen <sander.desmalen at arm.com>
Date: Thu, 8 Jan 2026 09:54:18 +0000
Subject: [PATCH] [compiler-rt][AArch64] Exit early from __arm_za_disable.

Because `__arm_za_disable` is a private-ZA function, it's only
ever entered with ZA state `off` or `dormant`. If the state is `off`
then we can safely return and there is no need to call `__arm_tpidr2_save`
or to explicitly set PSTATE.ZA or TPIDR2_EL0 to zero.
---
 compiler-rt/lib/builtins/aarch64/sme-abi.S | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/compiler-rt/lib/builtins/aarch64/sme-abi.S b/compiler-rt/lib/builtins/aarch64/sme-abi.S
index 1713a5969459a..465f1c763c0d3 100644
--- a/compiler-rt/lib/builtins/aarch64/sme-abi.S
+++ b/compiler-rt/lib/builtins/aarch64/sme-abi.S
@@ -161,6 +161,13 @@ DEFINE_COMPILERRT_FUNCTION(__arm_za_disable)
   ldr x14, [x14, CPU_FEATS_SYMBOL_OFFSET]
   tbz x14, #FEAT_SME_BIT, 0f
 
+  // If TPIDR2_EL0 is zero then nothing needs doing because
+  // __arm_za_disable is a private-ZA function and is only
+  // ever entered with ZA state 'off' or 'dormant', so
+  // we can assume PSTATE.ZA=0.
+  mrs x14, TPIDR2_EL0
+  cbz x14, 0f
+
   // Otherwise, the subroutine behaves as if it did the following:
   // * Call __arm_tpidr2_save.
   stp x29, x30, [sp, #-16]!