refactor(context mgmt): refactor EL2 context save and restore functions

This patch splits the el2_sysregs_context_save/restore functions
into multiple functions based on features. This will allow us to
selectively save and restore EL2 context registers based on
features enabled for a particular configuration.

For now feature build flags are used to decide which registers
to save and restore. The long term plan is to dynamically check
for features that are enabled and then save/restore registers
accordingly. Splitting el2_sysregs_context_save/restore functions
into smaller assembly functions makes that task easier. For more
information please take a look at:
https://trustedfirmware-a.readthedocs.io/en/latest/design_documents/context_mgmt_rework.html

Signed-off-by: Zelalem Aweke <zelalem.aweke@arm.com>
Change-Id: I1819a9de8b70fa35c8f45568908025f790c4808c
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 69acc2f..acfef80 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -11,8 +11,52 @@
 #include <el3_common_macros.S>
 
 #if CTX_INCLUDE_EL2_REGS
-	.global	el2_sysregs_context_save
-	.global	el2_sysregs_context_restore
+	.global	el2_sysregs_context_save_common
+	.global	el2_sysregs_context_restore_common
+#if ENABLE_SPE_FOR_LOWER_ELS
+	.global	el2_sysregs_context_save_spe
+	.global	el2_sysregs_context_restore_spe
+#endif /* ENABLE_SPE_FOR_LOWER_ELS */
+#if CTX_INCLUDE_MTE_REGS
+	.global	el2_sysregs_context_save_mte
+	.global	el2_sysregs_context_restore_mte
+#endif /* CTX_INCLUDE_MTE_REGS */
+#if ENABLE_MPAM_FOR_LOWER_ELS
+	.global	el2_sysregs_context_save_mpam
+	.global	el2_sysregs_context_restore_mpam
+#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
+#if ENABLE_FEAT_FGT
+	.global	el2_sysregs_context_save_fgt
+	.global	el2_sysregs_context_restore_fgt
+#endif /* ENABLE_FEAT_FGT */
+#if ENABLE_FEAT_ECV
+	.global	el2_sysregs_context_save_ecv
+	.global	el2_sysregs_context_restore_ecv
+#endif /* ENABLE_FEAT_ECV */
+#if ENABLE_FEAT_VHE
+	.global	el2_sysregs_context_save_vhe
+	.global	el2_sysregs_context_restore_vhe
+#endif /* ENABLE_FEAT_VHE */
+#if RAS_EXTENSION
+	.global	el2_sysregs_context_save_ras
+	.global	el2_sysregs_context_restore_ras
+#endif /* RAS_EXTENSION */
+#if CTX_INCLUDE_NEVE_REGS
+	.global	el2_sysregs_context_save_nv2
+	.global	el2_sysregs_context_restore_nv2
+#endif /* CTX_INCLUDE_NEVE_REGS */
+#if ENABLE_TRF_FOR_NS
+	.global	el2_sysregs_context_save_trf
+	.global	el2_sysregs_context_restore_trf
+#endif /* ENABLE_TRF_FOR_NS */
+#if ENABLE_FEAT_CSV2_2
+	.global	el2_sysregs_context_save_csv2
+	.global	el2_sysregs_context_restore_csv2
+#endif /* ENABLE_FEAT_CSV2_2 */
+#if ENABLE_FEAT_HCX
+	.global	el2_sysregs_context_save_hcx
+	.global	el2_sysregs_context_restore_hcx
+#endif /* ENABLE_FEAT_HCX */
 #endif /* CTX_INCLUDE_EL2_REGS */
 
 	.global	el1_sysregs_context_save
@@ -29,11 +73,16 @@
 #if CTX_INCLUDE_EL2_REGS
 
 /* -----------------------------------------------------
- * The following function strictly follows the AArch64
+ * The following functions strictly follow the AArch64
  * PCS to use x9-x16 (temporary caller-saved registers)
- * to save EL2 system register context. It assumes that
- * 'x0' is pointing to a 'el2_sys_regs' structure where
- * the register context will be saved.
+ * to save/restore EL2 system register context.
+ * el2_sysregs_context_save/restore_common functions
+ * save and restore registers that are common to all
+ * configurations. The rest of the functions save and
+ * restore EL2 system registers that are present when a
+ * particular feature is enabled. All functions assume
+ * that 'x0' is pointing to a 'el2_sys_regs' structure
+ * where the register context will be saved/restored.
  *
  * The following registers are not added.
  * AMEVCNTVOFF0<n>_EL2
@@ -43,7 +92,7 @@
  * ICH_LR<n>_EL2
  * -----------------------------------------------------
  */
-func el2_sysregs_context_save
+func el2_sysregs_context_save_common
 	mrs	x9, actlr_el2
 	mrs	x10, afsr0_el2
 	stp	x9, x10, [x0, #CTX_ACTLR_EL2]
@@ -88,11 +137,6 @@
 	mrs	x12, mdcr_el2
 	stp	x11, x12, [x0, #CTX_MAIR_EL2]
 
-#if ENABLE_SPE_FOR_LOWER_ELS
-	mrs	x13, PMSCR_EL2
-	str	x13, [x0, #CTX_PMSCR_EL2]
-#endif /* ENABLE_SPE_FOR_LOWER_ELS */
-
 	mrs	x14, sctlr_el2
 	str	x14, [x0, #CTX_SCTLR_EL2]
 
@@ -115,128 +159,10 @@
 	mrs	x15, vtcr_el2
 	mrs	x16, vttbr_el2
 	stp	x15, x16, [x0, #CTX_VTCR_EL2]
-
-#if CTX_INCLUDE_MTE_REGS
-	mrs	x9, TFSR_EL2
-	str	x9, [x0, #CTX_TFSR_EL2]
-#endif /* CTX_INCLUDE_MTE_REGS */
-
-#if ENABLE_MPAM_FOR_LOWER_ELS
-	mrs	x10, MPAM2_EL2
-	str	x10, [x0, #CTX_MPAM2_EL2]
-
-	mrs	x11, MPAMHCR_EL2
-	mrs	x12, MPAMVPM0_EL2
-	stp	x11, x12, [x0, #CTX_MPAMHCR_EL2]
-
-	mrs	x13, MPAMVPM1_EL2
-	mrs	x14, MPAMVPM2_EL2
-	stp	x13, x14, [x0, #CTX_MPAMVPM1_EL2]
-
-	mrs	x15, MPAMVPM3_EL2
-	mrs	x16, MPAMVPM4_EL2
-	stp	x15, x16, [x0, #CTX_MPAMVPM3_EL2]
-
-	mrs	x9, MPAMVPM5_EL2
-	mrs	x10, MPAMVPM6_EL2
-	stp	x9, x10, [x0, #CTX_MPAMVPM5_EL2]
-
-	mrs	x11, MPAMVPM7_EL2
-	mrs	x12, MPAMVPMV_EL2
-	stp	x11, x12, [x0, #CTX_MPAMVPM7_EL2]
-#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
-
-#if ENABLE_FEAT_FGT
-	mrs	x13, HDFGRTR_EL2
-#if ENABLE_FEAT_AMUv1
-   	mrs	x14, HAFGRTR_EL2
-   	stp	x13, x14, [x0, #CTX_HDFGRTR_EL2]
-#else
-   	str	x13, [x0, #CTX_HDFGRTR_EL2]
-#endif /* ENABLE_FEAT_AMUv1 */
-	mrs	x15, HDFGWTR_EL2
-	mrs	x16, HFGITR_EL2
-	stp	x15, x16, [x0, #CTX_HDFGWTR_EL2]
-
-	mrs	x9, HFGRTR_EL2
-	mrs	x10, HFGWTR_EL2
-	stp	x9, x10, [x0, #CTX_HFGRTR_EL2]
-#endif /* ENABLE_FEAT_FGT */
-
-#if ENABLE_FEAT_ECV
-	mrs	x11, CNTPOFF_EL2
-	str	x11, [x0, #CTX_CNTPOFF_EL2]
-#endif /* ENABLE_FEAT_ECV */
-
-#if ENABLE_FEAT_VHE
-	/*
-	 * CONTEXTIDR_EL2 register is saved only when FEAT_VHE or
-	 * FEAT_Debugv8p2 (currently not in TF-A) is supported.
-	 */
-	mrs	x9, contextidr_el2
-	mrs	x10, ttbr1_el2
-	stp	x9, x10, [x0, #CTX_CONTEXTIDR_EL2]
-#endif /* ENABLE_FEAT_VHE */
-
-#if RAS_EXTENSION
-	/*
-	 * VDISR_EL2 and VSESR_EL2 registers are saved only when
-	 * FEAT_RAS is supported.
-	 */
-	mrs	x11, vdisr_el2
-	mrs	x12, vsesr_el2
-	stp	x11, x12, [x0, #CTX_VDISR_EL2]
-#endif /* RAS_EXTENSION */
-
-#if CTX_INCLUDE_NEVE_REGS
-	/*
-	 * VNCR_EL2 register is saved only when FEAT_NV2 is supported.
-	 */
-	mrs	x16, vncr_el2
-	str	x16, [x0, #CTX_VNCR_EL2]
-#endif /* CTX_INCLUDE_NEVE_REGS */
-
-#if ENABLE_TRF_FOR_NS
-	/*
-	 * TRFCR_EL2 register is saved only when FEAT_TRF is supported.
-	 */
-	mrs	x12, TRFCR_EL2
-	str	x12, [x0, #CTX_TRFCR_EL2]
-#endif /* ENABLE_TRF_FOR_NS */
-
-#if ENABLE_FEAT_CSV2_2
-	/*
-	 * SCXTNUM_EL2 register is saved only when FEAT_CSV2_2 is supported.
-	 */
-	mrs	x13, scxtnum_el2
-	str	x13, [x0, #CTX_SCXTNUM_EL2]
-#endif /* ENABLE_FEAT_CSV2_2 */
-
-#if ENABLE_FEAT_HCX
-	mrs	x14, hcrx_el2
-	str	x14, [x0, #CTX_HCRX_EL2]
-#endif /* ENABLE_FEAT_HCX */
-
 	ret
-endfunc el2_sysregs_context_save
+endfunc el2_sysregs_context_save_common
 
-
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x16 (temporary caller-saved registers)
- * to restore EL2 system register context.  It assumes
- * that 'x0' is pointing to a 'el2_sys_regs' structure
- * from where the register context will be restored
-
- * The following registers are not restored
- * AMEVCNTVOFF0<n>_EL2
- * AMEVCNTVOFF1<n>_EL2
- * ICH_AP0R<n>_EL2
- * ICH_AP1R<n>_EL2
- * ICH_LR<n>_EL2
- * -----------------------------------------------------
- */
-func el2_sysregs_context_restore
+func el2_sysregs_context_restore_common
 	ldp	x9, x10, [x0, #CTX_ACTLR_EL2]
 	msr	actlr_el2, x9
 	msr	afsr0_el2, x10
@@ -281,11 +207,6 @@
 	msr	mair_el2, x11
 	msr	mdcr_el2, x12
 
-#if ENABLE_SPE_FOR_LOWER_ELS
-	ldr	x13, [x0, #CTX_PMSCR_EL2]
-	msr	PMSCR_EL2, x13
-#endif /* ENABLE_SPE_FOR_LOWER_ELS */
-
 	ldr	x14, [x0, #CTX_SCTLR_EL2]
 	msr	sctlr_el2, x14
 
@@ -308,13 +229,65 @@
 	ldp	x15, x16, [x0, #CTX_VTCR_EL2]
 	msr	vtcr_el2, x15
 	msr	vttbr_el2, x16
+	ret
+endfunc el2_sysregs_context_restore_common
+
+#if ENABLE_SPE_FOR_LOWER_ELS
+func el2_sysregs_context_save_spe
+	mrs	x13, PMSCR_EL2
+	str	x13, [x0, #CTX_PMSCR_EL2]
+	ret
+endfunc el2_sysregs_context_save_spe
+
+func el2_sysregs_context_restore_spe
+	ldr	x13, [x0, #CTX_PMSCR_EL2]
+	msr	PMSCR_EL2, x13
+	ret
+endfunc el2_sysregs_context_restore_spe
+#endif /* ENABLE_SPE_FOR_LOWER_ELS */
 
 #if CTX_INCLUDE_MTE_REGS
+func el2_sysregs_context_save_mte
+	mrs	x9, TFSR_EL2
+	str	x9, [x0, #CTX_TFSR_EL2]
+	ret
+endfunc el2_sysregs_context_save_mte
+
+func el2_sysregs_context_restore_mte
 	ldr	x9, [x0, #CTX_TFSR_EL2]
 	msr	TFSR_EL2, x9
+	ret
+endfunc el2_sysregs_context_restore_mte
 #endif /* CTX_INCLUDE_MTE_REGS */
 
 #if ENABLE_MPAM_FOR_LOWER_ELS
+func el2_sysregs_context_save_mpam
+	mrs	x10, MPAM2_EL2
+	str	x10, [x0, #CTX_MPAM2_EL2]
+
+	mrs	x11, MPAMHCR_EL2
+	mrs	x12, MPAMVPM0_EL2
+	stp	x11, x12, [x0, #CTX_MPAMHCR_EL2]
+
+	mrs	x13, MPAMVPM1_EL2
+	mrs	x14, MPAMVPM2_EL2
+	stp	x13, x14, [x0, #CTX_MPAMVPM1_EL2]
+
+	mrs	x15, MPAMVPM3_EL2
+	mrs	x16, MPAMVPM4_EL2
+	stp	x15, x16, [x0, #CTX_MPAMVPM3_EL2]
+
+	mrs	x9, MPAMVPM5_EL2
+	mrs	x10, MPAMVPM6_EL2
+	stp	x9, x10, [x0, #CTX_MPAMVPM5_EL2]
+
+	mrs	x11, MPAMVPM7_EL2
+	mrs	x12, MPAMVPMV_EL2
+	stp	x11, x12, [x0, #CTX_MPAMVPM7_EL2]
+	ret
+endfunc func el2_sysregs_context_save_mpam
+
+func el2_sysregs_context_restore_mpam
 	ldr	x10, [x0, #CTX_MPAM2_EL2]
 	msr	MPAM2_EL2, x10
 
@@ -337,10 +310,31 @@
 	ldp	x11, x12, [x0, #CTX_MPAMVPM7_EL2]
 	msr	MPAMVPM7_EL2, x11
 	msr	MPAMVPMV_EL2, x12
+	ret
+endfunc el2_sysregs_context_restore_mpam
 #endif /* ENABLE_MPAM_FOR_LOWER_ELS */
 
 #if ENABLE_FEAT_FGT
+func el2_sysregs_context_save_fgt
+	mrs	x13, HDFGRTR_EL2
 #if ENABLE_FEAT_AMUv1
+	mrs	x14, HAFGRTR_EL2
+	stp	x13, x14, [x0, #CTX_HDFGRTR_EL2]
+#else
+	str	x13, [x0, #CTX_HDFGRTR_EL2]
+#endif /* ENABLE_FEAT_AMUv1 */
+	mrs	x15, HDFGWTR_EL2
+	mrs	x16, HFGITR_EL2
+	stp	x15, x16, [x0, #CTX_HDFGWTR_EL2]
+
+	mrs	x9, HFGRTR_EL2
+	mrs	x10, HFGWTR_EL2
+	stp	x9, x10, [x0, #CTX_HFGRTR_EL2]
+	ret
+endfunc el2_sysregs_context_save_fgt
+
+func el2_sysregs_context_restore_fgt
+	#if ENABLE_FEAT_AMUv1
 	ldp	x13, x14, [x0, #CTX_HDFGRTR_EL2]
 	msr	HAFGRTR_EL2, x14
 #else
@@ -355,14 +349,37 @@
 	ldp	x9, x10, [x0, #CTX_HFGRTR_EL2]
 	msr	HFGRTR_EL2, x9
 	msr	HFGWTR_EL2, x10
+	ret
+endfunc el2_sysregs_context_restore_fgt
 #endif /* ENABLE_FEAT_FGT */
 
 #if ENABLE_FEAT_ECV
+func el2_sysregs_context_save_ecv
+	mrs	x11, CNTPOFF_EL2
+	str	x11, [x0, #CTX_CNTPOFF_EL2]
+	ret
+endfunc el2_sysregs_context_save_ecv
+
+func el2_sysregs_context_restore_ecv
 	ldr	x11, [x0, #CTX_CNTPOFF_EL2]
 	msr	CNTPOFF_EL2, x11
+	ret
+endfunc el2_sysregs_context_restore_ecv
 #endif /* ENABLE_FEAT_ECV */
 
 #if ENABLE_FEAT_VHE
+func el2_sysregs_context_save_vhe
+	/*
+	 * CONTEXTIDR_EL2 register is saved only when FEAT_VHE or
+	 * FEAT_Debugv8p2 (currently not in TF-A) is supported.
+	 */
+	mrs	x9, contextidr_el2
+	mrs	x10, ttbr1_el2
+	stp	x9, x10, [x0, #CTX_CONTEXTIDR_EL2]
+	ret
+endfunc el2_sysregs_context_save_vhe
+
+func el2_sysregs_context_restore_vhe
 	/*
 	 * CONTEXTIDR_EL2 register is restored only when FEAT_VHE or
 	 * FEAT_Debugv8p2 (currently not in TF-A) is supported.
@@ -370,9 +387,23 @@
 	ldp	x9, x10, [x0, #CTX_CONTEXTIDR_EL2]
 	msr	contextidr_el2, x9
 	msr	ttbr1_el2, x10
+	ret
+endfunc el2_sysregs_context_restore_vhe
 #endif /* ENABLE_FEAT_VHE */
 
 #if RAS_EXTENSION
+func el2_sysregs_context_save_ras
+	/*
+	 * VDISR_EL2 and VSESR_EL2 registers are saved only when
+	 * FEAT_RAS is supported.
+	 */
+	mrs	x11, vdisr_el2
+	mrs	x12, vsesr_el2
+	stp	x11, x12, [x0, #CTX_VDISR_EL2]
+	ret
+endfunc el2_sysregs_context_save_ras
+
+func el2_sysregs_context_restore_ras
 	/*
 	 * VDISR_EL2 and VSESR_EL2 registers are restored only when FEAT_RAS
 	 * is supported.
@@ -380,40 +411,83 @@
 	ldp	x11, x12, [x0, #CTX_VDISR_EL2]
 	msr	vdisr_el2, x11
 	msr	vsesr_el2, x12
+	ret
+endfunc el2_sysregs_context_restore_ras
 #endif /* RAS_EXTENSION */
 
 #if CTX_INCLUDE_NEVE_REGS
+func el2_sysregs_context_save_nv2
+	/*
+	 * VNCR_EL2 register is saved only when FEAT_NV2 is supported.
+	 */
+	mrs	x16, vncr_el2
+	str	x16, [x0, #CTX_VNCR_EL2]
+	ret
+endfunc el2_sysregs_context_save_nv2
+
+func el2_sysregs_context_restore_nv2
 	/*
 	 * VNCR_EL2 register is restored only when FEAT_NV2 is supported.
 	 */
 	ldr	x16, [x0, #CTX_VNCR_EL2]
 	msr	vncr_el2, x16
+	ret
+endfunc el2_sysregs_context_restore_nv2
 #endif /* CTX_INCLUDE_NEVE_REGS */
 
 #if ENABLE_TRF_FOR_NS
+func el2_sysregs_context_save_trf
+	/*
+	 * TRFCR_EL2 register is saved only when FEAT_TRF is supported.
+	 */
+	mrs	x12, TRFCR_EL2
+	str	x12, [x0, #CTX_TRFCR_EL2]
+	ret
+endfunc el2_sysregs_context_save_trf
+
+func el2_sysregs_context_restore_trf
 	/*
 	 * TRFCR_EL2 register is restored only when FEAT_TRF is supported.
 	 */
 	ldr	x12, [x0, #CTX_TRFCR_EL2]
 	msr	TRFCR_EL2, x12
+	ret
+endfunc el2_sysregs_context_restore_trf
 #endif /* ENABLE_TRF_FOR_NS */
 
 #if ENABLE_FEAT_CSV2_2
+func el2_sysregs_context_save_csv2
+	/*
+	 * SCXTNUM_EL2 register is saved only when FEAT_CSV2_2 is supported.
+	 */
+	mrs	x13, scxtnum_el2
+	str	x13, [x0, #CTX_SCXTNUM_EL2]
+	ret
+endfunc el2_sysregs_context_save_csv2
+
+func el2_sysregs_context_restore_csv2
 	/*
 	 * SCXTNUM_EL2 register is restored only when FEAT_CSV2_2 is supported.
 	 */
 	ldr	x13, [x0, #CTX_SCXTNUM_EL2]
 	msr	scxtnum_el2, x13
+	ret
+endfunc el2_sysregs_context_restore_csv2
 #endif /* ENABLE_FEAT_CSV2_2 */
 
 #if ENABLE_FEAT_HCX
+func el2_sysregs_context_save_hcx
+	mrs	x14, hcrx_el2
+	str	x14, [x0, #CTX_HCRX_EL2]
+	ret
+endfunc el2_sysregs_context_save_hcx
+
+func el2_sysregs_context_restore_hcx
 	ldr	x14, [x0, #CTX_HCRX_EL2]
 	msr	hcrx_el2, x14
-#endif /* ENABLE_FEAT_HCX */
-
 	ret
-endfunc el2_sysregs_context_restore
-
+endfunc el2_sysregs_context_restore_hcx
+#endif /* ENABLE_FEAT_HCX */
 #endif /* CTX_INCLUDE_EL2_REGS */
 
 /* ------------------------------------------------------------------
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index e393493..da610d0 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2022, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -790,11 +790,47 @@
 	if ((security_state != SECURE) ||
 	    ((security_state == SECURE) && ((scr_el3 & SCR_EEL2_BIT) != 0U))) {
 		cpu_context_t *ctx;
+		el2_sysregs_t *el2_sysregs_ctx;
 
 		ctx = cm_get_context(security_state);
 		assert(ctx != NULL);
 
-		el2_sysregs_context_save(get_el2_sysregs_ctx(ctx));
+		el2_sysregs_ctx = get_el2_sysregs_ctx(ctx);
+
+		el2_sysregs_context_save_common(el2_sysregs_ctx);
+#if ENABLE_SPE_FOR_LOWER_ELS
+		el2_sysregs_context_save_spe(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_MTE_REGS
+		el2_sysregs_context_save_mte(el2_sysregs_ctx);
+#endif
+#if ENABLE_MPAM_FOR_LOWER_ELS
+		el2_sysregs_context_save_mpam(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_FGT
+		el2_sysregs_context_save_fgt(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_ECV
+		el2_sysregs_context_save_ecv(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_VHE
+		el2_sysregs_context_save_vhe(el2_sysregs_ctx);
+#endif
+#if RAS_EXTENSION
+		el2_sysregs_context_save_ras(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_NEVE_REGS
+		el2_sysregs_context_save_nv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_TRF_FOR_NS
+		el2_sysregs_context_save_trf(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_CSV2_2
+		el2_sysregs_context_save_csv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_HCX
+		el2_sysregs_context_save_hcx(el2_sysregs_ctx);
+#endif
 	}
 }
 
@@ -812,11 +848,47 @@
 	if ((security_state != SECURE) ||
 	    ((security_state == SECURE) && ((scr_el3 & SCR_EEL2_BIT) != 0U))) {
 		cpu_context_t *ctx;
+		el2_sysregs_t *el2_sysregs_ctx;
 
 		ctx = cm_get_context(security_state);
 		assert(ctx != NULL);
 
-		el2_sysregs_context_restore(get_el2_sysregs_ctx(ctx));
+		el2_sysregs_ctx = get_el2_sysregs_ctx(ctx);
+
+		el2_sysregs_context_restore_common(el2_sysregs_ctx);
+#if ENABLE_SPE_FOR_LOWER_ELS
+		el2_sysregs_context_restore_spe(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_MTE_REGS
+		el2_sysregs_context_restore_mte(el2_sysregs_ctx);
+#endif
+#if ENABLE_MPAM_FOR_LOWER_ELS
+		el2_sysregs_context_restore_mpam(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_FGT
+		el2_sysregs_context_restore_fgt(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_ECV
+		el2_sysregs_context_restore_ecv(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_VHE
+		el2_sysregs_context_restore_vhe(el2_sysregs_ctx);
+#endif
+#if RAS_EXTENSION
+		el2_sysregs_context_restore_ras(el2_sysregs_ctx);
+#endif
+#if CTX_INCLUDE_NEVE_REGS
+		el2_sysregs_context_restore_nv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_TRF_FOR_NS
+		el2_sysregs_context_restore_trf(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_CSV2_2
+		el2_sysregs_context_restore_csv2(el2_sysregs_ctx);
+#endif
+#if ENABLE_FEAT_HCX
+		el2_sysregs_context_restore_hcx(el2_sysregs_ctx);
+#endif
 	}
 }
 #endif /* CTX_INCLUDE_EL2_REGS */