fix(lib): save and restore smcr_el2 when cpu suspend and resume

Change-Id: I0c0f054e025eaafeebbe0adb5d0699def5c4fc60
Signed-off-by: Jing Han <jing.han@arm.com>
diff --git a/lib/power_management/suspend/aarch64/asm_tftf_suspend.S b/lib/power_management/suspend/aarch64/asm_tftf_suspend.S
index e715e49..c643475 100644
--- a/lib/power_management/suspend/aarch64/asm_tftf_suspend.S
+++ b/lib/power_management/suspend/aarch64/asm_tftf_suspend.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -81,8 +81,18 @@
 	stp	x3, x4, [x0, #SUSPEND_CTX_TTBR0_OFFSET]
 	stp	x5, x6, [x0, #SUSPEND_CTX_VBAR_OFFSET]
 	mrs	x1, hcr_el2
+	/*
+	 * Check if the processor supports SME
+	 */
+	mrs	x2, id_aa64pfr1_el1
+	tst	x2, #(1 << 25)
+	bne	3f
 	str	x1, [x0, #SUSPEND_CTX_HCR_OFFSET]
 	ret
+
+3:	mrs	x2, SMCR_EL2
+	stp	x1, x2, [x0, #SUSPEND_CTX_HCR_OFFSET]
+	ret
 endfunc __tftf_save_arch_context
 
 /*
@@ -120,13 +130,24 @@
 	msr	ttbr0_el2, x3
 	msr	tcr_el2, x4
 	msr	vbar_el2, x5
+        /*
+	 * Check if the processor supports SME
+	 */
+	mrs	x2, id_aa64pfr1_el1
+	tst	x2, #(1 << 25)
+	bne	3f
 	ldr	x1, [x0, #SUSPEND_CTX_HCR_OFFSET]
 	msr	hcr_el2, x1
+	b	4f
+
+3:	ldp	x1, x2, [x0, #SUSPEND_CTX_HCR_OFFSET]
+	msr     hcr_el2, x1
+	msr	SMCR_EL2, x2
 
 	/*
 	 * TLB invalidations need to be completed before enabling MMU
 	 */
-	dsb	nsh
+4:	dsb	nsh
 	msr	sctlr_el2, x6
 	/* Ensure the MMU enable takes effect immediately */
 	isb
diff --git a/lib/power_management/suspend/suspend_private.h b/lib/power_management/suspend/suspend_private.h
index 2ccf7a3..debd84f 100644
--- a/lib/power_management/suspend/suspend_private.h
+++ b/lib/power_management/suspend/suspend_private.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -10,8 +10,9 @@
 /*
  * Number of system registers we need to save/restore across a CPU suspend:
  * EL1: MAIR, CPACR, TTBR0, TCR, VBAR, SCTLR
- * EL2: MAIR, CPTR, TTBR0, TCR, VBAR, SCTLR, HCR
+ * EL2: MAIR, CPTR, TTBR0, TCR, VBAR, SCTLR, HCR, SMCR
  * APIAKeyLo_EL1 and APIAKeyHi_EL1 (if enabled).
+ * On need basis other registers can be included in tftf_suspend_context.
  */
 #if ENABLE_PAUTH
 #define NR_CTX_REGS 10