diff --git a/changelog.yaml b/changelog.yaml
index 1a11c9b..dd60020 100644
--- a/changelog.yaml
+++ b/changelog.yaml
@@ -623,6 +623,9 @@
       - title: PSA
         scope: psa
 
+      - title: Context Management
+        scope: context mgmt
+
   - title: Drivers
 
     subsections:
diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index 3a09383..6c13166 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -509,9 +509,53 @@
 void el1_sysregs_context_restore(el1_sysregs_t *regs);
 
 #if CTX_INCLUDE_EL2_REGS
-void el2_sysregs_context_save(el2_sysregs_t *regs);
-void el2_sysregs_context_restore(el2_sysregs_t *regs);
-#endif
+void el2_sysregs_context_save_common(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_common(el2_sysregs_t *regs);
+#if ENABLE_SPE_FOR_LOWER_ELS
+void el2_sysregs_context_save_spe(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_spe(el2_sysregs_t *regs);
+#endif /* ENABLE_SPE_FOR_LOWER_ELS */
+#if CTX_INCLUDE_MTE_REGS
+void el2_sysregs_context_save_mte(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_mte(el2_sysregs_t *regs);
+#endif /* CTX_INCLUDE_MTE_REGS */
+#if ENABLE_MPAM_FOR_LOWER_ELS
+void el2_sysregs_context_save_mpam(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_mpam(el2_sysregs_t *regs);
+#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
+#if ENABLE_FEAT_FGT
+void el2_sysregs_context_save_fgt(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_fgt(el2_sysregs_t *regs);
+#endif /* ENABLE_FEAT_FGT */
+#if ENABLE_FEAT_ECV
+void el2_sysregs_context_save_ecv(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_ecv(el2_sysregs_t *regs);
+#endif /* ENABLE_FEAT_ECV */
+#if ENABLE_FEAT_VHE
+void el2_sysregs_context_save_vhe(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_vhe(el2_sysregs_t *regs);
+#endif /* ENABLE_FEAT_VHE */
+#if RAS_EXTENSION
+void el2_sysregs_context_save_ras(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_ras(el2_sysregs_t *regs);
+#endif /* RAS_EXTENSION */
+#if CTX_INCLUDE_NEVE_REGS
+void el2_sysregs_context_save_nv2(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_nv2(el2_sysregs_t *regs);
+#endif /* CTX_INCLUDE_NEVE_REGS */
+#if ENABLE_TRF_FOR_NS
+void el2_sysregs_context_save_trf(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_trf(el2_sysregs_t *regs);
+#endif /* ENABLE_TRF_FOR_NS */
+#if ENABLE_FEAT_CSV2_2
+void el2_sysregs_context_save_csv2(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_csv2(el2_sysregs_t *regs);
+#endif /* ENABLE_FEAT_CSV2_2 */
+#if ENABLE_FEAT_HCX
+void el2_sysregs_context_save_hcx(el2_sysregs_t *regs);
+void el2_sysregs_context_restore_hcx(el2_sysregs_t *regs);
+#endif /* ENABLE_FEAT_HCX */
+#endif /* CTX_INCLUDE_EL2_REGS */
 
 #if CTX_INCLUDE_FPREGS
 void fpregs_context_save(fp_regs_t *regs);
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 */