[PATCH] arm64: suspend: Save and restore TPIDR2_EL0 across CPU suspend

[Wei Wang <wvw@google.com>]

On systems with FEAT_SME, TPIDR2_EL0 is part of userspace thread state:
the AAPCS64 SME ABI uses it as the ZA lazy-save pointer, and the kernel
context-switches it eagerly in tls_thread_switch(). While a task is
running, the live register is the only copy of its value.

The architecture defines TPIDR2_EL0 as UNKNOWN after a warm reset, so
a CPU power-down through PSCI CPU_SUSPEND or SYSTEM_SUSPEND can lose it.
cpu_do_suspend() and cpu_do_resume() already preserve TPIDR_EL0 and
TPIDRRO_EL0, but not TPIDR2_EL0.

For CPU idle this is harmless: current is the idle task, and the next
context switch reloads TPIDR2_EL0 from next->thread.tpidr2_el0. For
suspend-to-RAM and hibernation, however, cpu_suspend() runs in the
context of the userspace task that wrote /sys/power/state, and that
task returns to EL0 with an UNKNOWN TPIDR2_EL0. On hardware where the
reset value is non-zero, the next SME ABI lazy-save check in that
thread faults. For example, bionic's fork() -> clone() ->
__arm_za_disable() -> __arm_tpidr2_save() loads from the garbage
pointer:

  # echo mem > /sys/power/state; true
  signal 11 (SIGSEGV), code 1 (SEGV_MAPERR), fault addr 0xa3bad11bcaf72e93
      x16 a3bad11bcaf72e89
  #00 pc ... libc.so (__arm_tpidr2_save+24)
  #01 pc ... libc.so (__arm_za_disable+32)
  #02 pc ... libc.so (clone+480)
  #03 pc ... libc.so (fork+60)

Having firmware zero TPIDR2_EL0 on warm reset avoids the crash. It is
not enough on its own, though: a thread that enters the kernel with a
lazy save pending (TPIDR2_EL0 pointing at a valid TPIDR2 block) would
see TPIDR2_EL0 == 0 on return. It would then wrongly conclude that ZA
had been committed and restore stale data from its save buffer. The
kernel owns this state and must preserve it.

Save TPIDR2_EL0 in cpu_do_suspend() and restore it in cpu_do_resume()
when SME is present. Placing this in proc.S covers both cpu_suspend()
and hibernation. The new slot uses ctx_regs[14]; sp moves into what was
previously tail padding, so sizeof(struct cpu_suspend_ctx) stays 128
bytes. sp is only accessed through the CPU_CTX_SP asm-offset.

Fixes: a9d691585950 ("arm64/sme: Implement support for TPIDR2")
Signed-off-by: Wei Wang <wvw@google.com>
---
 arch/arm64/include/asm/suspend.h | 2 +-
 arch/arm64/mm/proc.S             | 8 ++++++++
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/arch/arm64/include/asm/suspend.h b/arch/arm64/include/asm/suspend.h
index eb60c9735553..2cc5fb405769 100644
--- a/arch/arm64/include/asm/suspend.h
+++ b/arch/arm64/include/asm/suspend.h
@@ -2,7 +2,7 @@
 #ifndef __ASM_SUSPEND_H
 #define __ASM_SUSPEND_H
 
-#define NR_CTX_REGS 14
+#define NR_CTX_REGS 15
 #define NR_CALLEE_SAVED_REGS 12
 
 /*
diff --git a/arch/arm64/mm/proc.S b/arch/arm64/mm/proc.S
index 7afbece885a3..c7b7fa6491c3 100644
--- a/arch/arm64/mm/proc.S
+++ b/arch/arm64/mm/proc.S
@@ -103,6 +103,10 @@ SYM_FUNC_START(cpu_do_suspend)
 alternative_if ARM64_HAS_TCR2
 	mrs	x2, REG_TCR2_EL1
 	str	x2, [x0, #104]
+alternative_else_nop_endif
+alternative_if ARM64_SME
+	mrs_s	x2, SYS_TPIDR2_EL0
+	str	x2, [x0, #112]
 alternative_else_nop_endif
 	ret
 SYM_FUNC_END(cpu_do_suspend)
@@ -142,6 +146,10 @@ alternative_if ARM64_HAS_TCR2
 	ldr	x2, [x0, #104]
 	msr	REG_TCR2_EL1, x2
 alternative_else_nop_endif
+alternative_if ARM64_SME
+	ldr	x2, [x0, #112]
+	msr_s	SYS_TPIDR2_EL0, x2
+alternative_else_nop_endif
 
 	msr	sctlr_el1, x12
 	set_this_cpu_offset x13
-- 
2.56.0.385.gac671c819a-goog