SPMD: Adds partially supported EL2 registers.

This patch adds EL2 registers that are supported up to ARMv8.6.
ARM_ARCH_MINOR has to specified to enable save/restore routine.

Note: Following registers are still not covered in save/restore.
 * AMEVCNTVOFF0<n>_EL2
 * AMEVCNTVOFF1<n>_EL2
 * ICH_AP0R<n>_EL2
 * ICH_AP1R<n>_EL2
 * ICH_LR<n>_EL2

Change-Id: I4813f3243e56e21cb297b31ef549a4b38d4876e1
Signed-off-by: Max Shvetsov <maksims.svecovs@arm.com>
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index bcc7eef..30ad7b7 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -29,189 +29,187 @@
 /* -----------------------------------------------------
  * The following function strictly follows the AArch64
  * PCS to use x9-x17 (temporary caller-saved registers)
- * to save EL1 system register context. It assumes that
- * 'x0' is pointing to a 'el1_sys_regs' structure where
+ * 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.
+ *
+ * The following registers are not added.
+ * AMEVCNTVOFF0<n>_EL2
+ * AMEVCNTVOFF1<n>_EL2
+ * ICH_AP0R<n>_EL2
+ * ICH_AP1R<n>_EL2
+ * ICH_LR<n>_EL2
  * -----------------------------------------------------
  */
+
 func el2_sysregs_context_save
-
 	mrs	x9, actlr_el2
-	str	x9, [x0, #CTX_ACTLR_EL2]
+	mrs	x10, afsr0_el2
+	stp	x9, x10, [x0, #CTX_ACTLR_EL2]
 
-	mrs	x9, afsr0_el2
-	str	x9, [x0, #CTX_AFSR0_EL2]
+	mrs	x11, afsr1_el2
+	mrs	x12, amair_el2
+	stp	x11, x12, [x0, #CTX_AFSR1_EL2]
 
-	mrs	x9, afsr1_el2
-	str	x9, [x0, #CTX_AFSR1_EL2]
+	mrs	x13, cnthctl_el2
+	mrs	x14, cnthp_ctl_el2
+	stp	x13, x14, [x0, #CTX_CNTHCTL_EL2]
 
-	mrs	x9, amair_el2
-	str	x9, [x0, #CTX_AMAIR_EL2]
+	mrs	x15, cnthp_cval_el2
+	mrs	x16, cnthp_tval_el2
+	stp	x15, x16, [x0, #CTX_CNTHP_CVAL_EL2]
 
-	mrs	x9, cnthctl_el2
-	str	x9, [x0, #CTX_CNTHCTL_EL2]
-
-	mrs	x9, cnthp_ctl_el2
-	str	x9, [x0, #CTX_CNTHP_CTL_EL2]
-
-	mrs	x9, cnthp_cval_el2
-	str	x9, [x0, #CTX_CNTHP_CVAL_EL2]
-
-	mrs	x9, cnthp_tval_el2
-	str	x9, [x0, #CTX_CNTHP_TVAL_EL2]
-
-	mrs	x9, CNTPOFF_EL2
-	str	x9, [x0, #CTX_CNTPOFF_EL2]
-
-	mrs	x9, cntvoff_el2
-	str	x9, [x0, #CTX_CNTVOFF_EL2]
-
+	mrs	x17, cntvoff_el2
 	mrs	x9, cptr_el2
-	str	x9, [x0, #CTX_CPTR_EL2]
+	stp	x17, x9, [x0, #CTX_CNTVOFF_EL2]
 
-	mrs	x9, dbgvcr32_el2
-	str	x9, [x0, #CTX_DBGVCR32_EL2]
+	mrs	x10, dbgvcr32_el2
+	mrs	x11, elr_el2
+	stp	x10, x11, [x0, #CTX_DBGVCR32_EL2]
 
-	mrs	x9, elr_el2
-	str	x9, [x0, #CTX_ELR_EL2]
+	mrs	x14, esr_el2
+	mrs	x15, far_el2
+	stp	x14, x15, [x0, #CTX_ESR_EL2]
 
-	mrs	x9, esr_el2
-	str	x9, [x0, #CTX_ESR_EL2]
-
-	mrs	x9, far_el2
-	str	x9, [x0, #CTX_FAR_EL2]
-
-	mrs	x9, fpexc32_el2
-	str	x9, [x0, #CTX_FPEXC32_EL2]
-
-	mrs	x9, hacr_el2
-	str	x9, [x0, #CTX_HACR_EL2]
-
-	mrs	x9, HAFGRTR_EL2
-	str	x9, [x0, #CTX_HAFGRTR_EL2]
+	mrs	x16, fpexc32_el2
+	mrs	x17, hacr_el2
+	stp	x16, x17, [x0, #CTX_FPEXC32_EL2]
 
 	mrs	x9, hcr_el2
-	str	x9, [x0, #CTX_HCR_EL2]
+	mrs	x10, hpfar_el2
+	stp	x9, x10, [x0, #CTX_HCR_EL2]
 
-	mrs	x9, HDFGRTR_EL2
-	str	x9, [x0, #CTX_HDFGRTR_EL2]
+	mrs	x11, hstr_el2
+	mrs	x12, ICC_SRE_EL2
+	stp	x11, x12, [x0, #CTX_HSTR_EL2]
 
-	mrs	x9, HDFGWTR_EL2
-	str	x9, [x0, #CTX_HDFGWTR_EL2]
+	mrs	x13, ICH_HCR_EL2
+	mrs	x14, ICH_VMCR_EL2
+	stp	x13, x14, [x0, #CTX_ICH_HCR_EL2]
 
-	mrs	x9, HFGITR_EL2
-	str	x9, [x0, #CTX_HFGITR_EL2]
+	mrs	x15, mair_el2
+	mrs	x16, mdcr_el2
+	stp	x15, x16, [x0, #CTX_MAIR_EL2]
 
-	mrs	x9, HFGRTR_EL2
-	str	x9, [x0, #CTX_HFGRTR_EL2]
-
-	mrs	x9, HFGWTR_EL2
-	str	x9, [x0, #CTX_HFGWTR_EL2]
-
-	mrs	x9, hpfar_el2
-	str	x9, [x0, #CTX_HPFAR_EL2]
-
-	mrs	x9, hstr_el2
-	str	x9, [x0, #CTX_HSTR_EL2]
-
-	mrs	x9, ICC_SRE_EL2
-	str	x9, [x0, #CTX_ICC_SRE_EL2]
-
-	mrs	x9, ICH_EISR_EL2
-	str	x9, [x0, #CTX_ICH_EISR_EL2]
-
-	mrs	x9, ICH_ELRSR_EL2
-	str	x9, [x0, #CTX_ICH_ELRSR_EL2]
-
-	mrs	x9, ICH_HCR_EL2
-	str	x9, [x0, #CTX_ICH_HCR_EL2]
-
-	mrs	x9, ICH_MISR_EL2
-	str	x9, [x0, #CTX_ICH_MISR_EL2]
-
-	mrs	x9, ICH_VMCR_EL2
-	str	x9, [x0, #CTX_ICH_VMCR_EL2]
-
-	mrs	x9, ICH_VTR_EL2
-	str	x9, [x0, #CTX_ICH_VTR_EL2]
-
-	mrs	x9, mair_el2
-	str	x9, [x0, #CTX_MAIR_EL2]
-
-	mrs	x9, mdcr_el2
-	str	x9, [x0, #CTX_MDCR_EL2]
-
-	mrs	x9, MPAM2_EL2
-	str	x9, [x0, #CTX_MPAM2_EL2]
-
-	mrs	x9, MPAMHCR_EL2
-	str	x9, [x0, #CTX_MPAMHCR_EL2]
-
-	mrs	x9, MPAMVPM0_EL2
-	str	x9, [x0, #CTX_MPAMVPM0_EL2]
-
-	mrs	x9, MPAMVPM1_EL2
-	str	x9, [x0, #CTX_MPAMVPM1_EL2]
-
-	mrs	x9, MPAMVPM2_EL2
-	str	x9, [x0, #CTX_MPAMVPM2_EL2]
-
-	mrs	x9, MPAMVPM3_EL2
-	str	x9, [x0, #CTX_MPAMVPM3_EL2]
-
-	mrs	x9, MPAMVPM4_EL2
-	str	x9, [x0, #CTX_MPAMVPM4_EL2]
-
-	mrs	x9, MPAMVPM5_EL2
-	str	x9, [x0, #CTX_MPAMVPM5_EL2]
-
-	mrs	x9, MPAMVPM6_EL2
-	str	x9, [x0, #CTX_MPAMVPM6_EL2]
-
-	mrs	x9, MPAMVPM7_EL2
-	str	x9, [x0, #CTX_MPAMVPM7_EL2]
-
-	mrs	x9, MPAMVPMV_EL2
-	str	x9, [x0, #CTX_MPAMVPMV_EL2]
-
-	mrs	x9, rmr_el2
-	str	x9, [x0, #CTX_RMR_EL2]
-
+	mrs	x17, PMSCR_EL2
 	mrs	x9, sctlr_el2
-	str	x9, [x0, #CTX_SCTLR_EL2]
+	stp	x17, x9, [x0, #CTX_PMSCR_EL2]
 
-	mrs	x9, spsr_el2
-	str	x9, [x0, #CTX_SPSR_EL2]
+	mrs	x10, spsr_el2
+	mrs	x11, sp_el2
+	stp	x10, x11, [x0, #CTX_SPSR_EL2]
 
-	mrs	x9, sp_el2
-	str	x9, [x0, #CTX_SP_EL2]
+	mrs	x12, tcr_el2
+	mrs	x13, TRFCR_EL2
+	stp	x12, x13, [x0, #CTX_TCR_EL2]
 
-	mrs	x9, tcr_el2
-	str	x9, [x0, #CTX_TCR_EL2]
+	mrs	x14, ttbr0_el2
+	mrs	x15, vbar_el2
+	stp	x14, x15, [x0, #CTX_TTBR0_EL2]
 
-	mrs	x9, tpidr_el2
-	str	x9, [x0, #CTX_TPIDR_EL2]
-
-	mrs	x9, ttbr0_el2
-	str	x9, [x0, #CTX_TTBR0_EL2]
-
-	mrs	x9, vbar_el2
-	str	x9, [x0, #CTX_VBAR_EL2]
-
-	mrs	x9, vmpidr_el2
-	str	x9, [x0, #CTX_VMPIDR_EL2]
-
-	mrs	x9, vpidr_el2
-	str	x9, [x0, #CTX_VPIDR_EL2]
+	mrs	x16, vmpidr_el2
+	mrs	x17, vpidr_el2
+	stp	x16, x17, [x0, #CTX_VMPIDR_EL2]
 
 	mrs	x9, vtcr_el2
-	str	x9, [x0, #CTX_VTCR_EL2]
+	mrs	x10, vttbr_el2
+	stp	x9, x10, [x0, #CTX_VTCR_EL2]
 
-	mrs	x9, vttbr_el2
-	str	x9, [x0, #CTX_VTTBR_EL2]
+#if CTX_INCLUDE_MTE_REGS
+	mrs	x11, TFSR_EL2
+	str	x11, [x0, #CTX_TFSR_EL2]
+#endif
 
-	mrs	x9, ZCR_EL2
-	str	x9, [x0, #CTX_ZCR_EL2]
+#if ENABLE_MPAM_FOR_LOWER_ELS
+	mrs	x9, MPAM2_EL2
+	mrs	x10, MPAMHCR_EL2
+	stp	x9, x10, [x0, #CTX_MPAM2_EL2]
+
+	mrs	x11, MPAMVPM0_EL2
+	mrs	x12, MPAMVPM1_EL2
+	stp	x11, x12, [x0, #CTX_MPAMVPM0_EL2]
+
+	mrs	x13, MPAMVPM2_EL2
+	mrs	x14, MPAMVPM3_EL2
+	stp	x13, x14, [x0, #CTX_MPAMVPM2_EL2]
+
+	mrs	x15, MPAMVPM4_EL2
+	mrs	x16, MPAMVPM5_EL2
+	stp	x15, x16, [x0, #CTX_MPAMVPM4_EL2]
+
+	mrs	x17, MPAMVPM6_EL2
+	mrs	x9, MPAMVPM7_EL2
+	stp	x17, x9, [x0, #CTX_MPAMVPM6_EL2]
+
+	mrs	x10, MPAMVPMV_EL2
+	str	x10, [x0, #CTX_MPAMVPMV_EL2]
+#endif
+
+
+#if ARM_ARCH_AT_LEAST(8, 6)
+	mrs	x11, HAFGRTR_EL2
+	mrs	x12, HDFGRTR_EL2
+	stp	x11, x12, [x0, #CTX_HAFGRTR_EL2]
+
+	mrs	x13, HDFGWTR_EL2
+	mrs	x14, HFGITR_EL2
+	stp	x13, x14, [x0, #CTX_HDFGWTR_EL2]
+
+	mrs	x15, HFGRTR_EL2
+	mrs	x16, HFGWTR_EL2
+	stp	x15, x16, [x0, #CTX_HFGRTR_EL2]
+
+	mrs	x17, CNTPOFF_EL2
+	str	x17, [x0, #CTX_CNTPOFF_EL2]
+#endif
+
+#if ARM_ARCH_AT_LEAST(8, 4)
+	mrs	x9, cnthps_ctl_el2
+	mrs	x10, cnthps_cval_el2
+	stp	x9, x10, [x0, #CTX_CNTHPS_CTL_EL2]
+
+	mrs	x11, cnthps_tval_el2
+	mrs	x12, cnthvs_ctl_el2
+	stp	x11, x12, [x0, #CTX_CNTHPS_TVAL_EL2]
+
+	mrs	x13, cnthvs_cval_el2
+	mrs	x14, cnthvs_tval_el2
+	stp	x13, x14, [x0, #CTX_CNTHVS_CVAL_EL2]
+
+	mrs	x15, cnthv_ctl_el2
+	mrs	x16, cnthv_cval_el2
+	stp	x15, x16, [x0, #CTX_CNTHV_CTL_EL2]
+
+	mrs	x17, cnthv_tval_el2
+	mrs	x9, contextidr_el2
+	stp	x17, x9, [x0, #CTX_CNTHV_TVAL_EL2]
+
+	mrs	x10, sder32_el2
+	str	x10, [x0, #CTX_SDER32_EL2]
+
+	mrs	x11, ttbr1_el2
+	str	x11, [x0, #CTX_TTBR1_EL2]
+
+	mrs	x12, vdisr_el2
+	str	x12, [x0, #CTX_VDISR_EL2]
+
+	mrs	x13, vncr_el2
+	str	x13, [x0, #CTX_VNCR_EL2]
+
+	mrs	x14, vsesr_el2
+	str	x14, [x0, #CTX_VSESR_EL2]
+
+	mrs	x15, vstcr_el2
+	str	x15, [x0, #CTX_VSTCR_EL2]
+
+	mrs	x16, vsttbr_el2
+	str	x16, [x0, #CTX_VSTTBR_EL2]
+#endif
+
+#if ARM_ARCH_AT_LEAST(8, 5)
+	mrs	x17, scxtnum_el2
+	str	x17, [x0, #CTX_SCXTNUM_EL2]
+#endif
 
 	ret
 endfunc el2_sysregs_context_save
@@ -219,189 +217,186 @@
 /* -----------------------------------------------------
  * The following function strictly follows the AArch64
  * PCS to use x9-x17 (temporary caller-saved registers)
- * to restore EL1 system register context.  It assumes
- * that 'x0' is pointing to a 'el1_sys_regs' structure
+ * 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
 
-	ldr	x9, [x0, #CTX_ACTLR_EL2]
+	ldp	x9, x10, [x0, #CTX_ACTLR_EL2]
 	msr	actlr_el2, x9
+	msr	afsr0_el2, x10
 
-	ldr	x9, [x0, #CTX_AFSR0_EL2]
-	msr	afsr0_el2, x9
+	ldp	x11, x12, [x0, #CTX_AFSR1_EL2]
+	msr	afsr1_el2, x11
+	msr	amair_el2, x12
 
-	ldr	x9, [x0, #CTX_AFSR1_EL2]
-	msr	afsr1_el2, x9
+	ldp	x13, x14, [x0, #CTX_CNTHCTL_EL2]
+	msr	cnthctl_el2, x13
+	msr	cnthp_ctl_el2, x14
 
-	ldr	x9, [x0, #CTX_AMAIR_EL2]
-	msr	amair_el2, x9
+	ldp	x15, x16, [x0, #CTX_CNTHP_CVAL_EL2]
+	msr	cnthp_cval_el2, x15
+	msr	cnthp_tval_el2, x16
 
-	ldr	x9, [x0, #CTX_CNTHCTL_EL2]
-	msr	cnthctl_el2, x9
-
-	ldr	x9, [x0, #CTX_CNTHP_CTL_EL2]
-	msr	cnthp_ctl_el2, x9
-
-	ldr	x9, [x0, #CTX_CNTHP_CVAL_EL2]
-	msr	cnthp_cval_el2, x9
-
-	ldr	x9, [x0, #CTX_CNTHP_TVAL_EL2]
-	msr	cnthp_tval_el2, x9
-
-	ldr	x9, [x0, #CTX_CNTPOFF_EL2]
-	msr	CNTPOFF_EL2, x9
-
-	ldr	x9, [x0, #CTX_CNTVOFF_EL2]
-	msr	cntvoff_el2, x9
-
-	ldr	x9, [x0, #CTX_CPTR_EL2]
+	ldp	x17, x9, [x0, #CTX_CNTVOFF_EL2]
+	msr	cntvoff_el2, x17
 	msr	cptr_el2, x9
 
-	ldr	x9, [x0, #CTX_DBGVCR32_EL2]
-	msr	dbgvcr32_el2, x9
+	ldp	x10, x11, [x0, #CTX_DBGVCR32_EL2]
+	msr	dbgvcr32_el2, x10
+	msr	elr_el2, x11
 
-	ldr	x9, [x0, #CTX_ELR_EL2]
-	msr	elr_el2, x9
+	ldp	x14, x15, [x0, #CTX_ESR_EL2]
+	msr	esr_el2, x14
+	msr	far_el2, x15
 
-	ldr	x9, [x0, #CTX_ESR_EL2]
-	msr	esr_el2, x9
+	ldp	x16, x17, [x0, #CTX_FPEXC32_EL2]
+	msr	fpexc32_el2, x16
+	msr	hacr_el2, x17
 
-	ldr	x9, [x0, #CTX_FAR_EL2]
-	msr	far_el2, x9
-
-	ldr	x9, [x0, #CTX_FPEXC32_EL2]
-	msr	fpexc32_el2, x9
-
-	ldr	x9, [x0, #CTX_HACR_EL2]
-	msr	hacr_el2, x9
-
-	ldr	x9, [x0, #CTX_HAFGRTR_EL2]
-	msr	HAFGRTR_EL2, x9
-
-	ldr	x9, [x0, #CTX_HCR_EL2]
+	ldp	x9, x10, [x0, #CTX_HCR_EL2]
 	msr	hcr_el2, x9
+	msr	hpfar_el2, x10
 
-	ldr	x9, [x0, #CTX_HDFGRTR_EL2]
-	msr	HDFGRTR_EL2, x9
+	ldp	x11, x12, [x0, #CTX_HSTR_EL2]
+	msr	hstr_el2, x11
+	msr	ICC_SRE_EL2, x12
 
-	ldr	x9, [x0, #CTX_HDFGWTR_EL2]
-	msr	HDFGWTR_EL2, x9
+	ldp	x13, x14, [x0, #CTX_ICH_HCR_EL2]
+	msr	ICH_HCR_EL2, x13
+	msr	ICH_VMCR_EL2, x14
 
-	ldr	x9, [x0, #CTX_HFGITR_EL2]
-	msr	HFGITR_EL2, x9
+	ldp	x15, x16, [x0, #CTX_MAIR_EL2]
+	msr	mair_el2, x15
+	msr	mdcr_el2, x16
 
-	ldr	x9, [x0, #CTX_HFGRTR_EL2]
-	msr	HFGRTR_EL2, x9
-
-	ldr	x9, [x0, #CTX_HFGWTR_EL2]
-	msr	HFGWTR_EL2, x9
-
-	ldr	x9, [x0, #CTX_HPFAR_EL2]
-	msr	hpfar_el2, x9
-
-	ldr	x9, [x0, #CTX_HSTR_EL2]
-	msr	hstr_el2, x9
-
-	ldr	x9, [x0, #CTX_ICC_SRE_EL2]
-	msr	ICC_SRE_EL2, x9
-
-	ldr	x9, [x0, #CTX_ICH_EISR_EL2]
-	msr	ICH_EISR_EL2, x9
-
-	ldr	x9, [x0, #CTX_ICH_ELRSR_EL2]
-	msr	ICH_ELRSR_EL2, x9
-
-	ldr	x9, [x0, #CTX_ICH_HCR_EL2]
-	msr	ICH_HCR_EL2, x9
-
-	ldr	x9, [x0, #CTX_ICH_MISR_EL2]
-	msr	ICH_MISR_EL2, x9
-
-	ldr	x9, [x0, #CTX_ICH_VMCR_EL2]
-	msr	ICH_VMCR_EL2, x9
-
-	ldr	x9, [x0, #CTX_ICH_VTR_EL2]
-	msr	ICH_VTR_EL2, x9
-
-	ldr	x9, [x0, #CTX_MAIR_EL2]
-	msr	mair_el2, x9
-
-	ldr	x9, [x0, #CTX_MDCR_EL2]
-	msr	mdcr_el2, x9
-
-	ldr	x9, [x0, #CTX_MPAM2_EL2]
-	msr	MPAM2_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMHCR_EL2]
-	msr	MPAMHCR_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM0_EL2]
-	msr	MPAMVPM0_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM1_EL2]
-	msr	MPAMVPM1_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM2_EL2]
-	msr	MPAMVPM2_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM3_EL2]
-	msr	MPAMVPM3_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM4_EL2]
-	msr	MPAMVPM4_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM5_EL2]
-	msr	MPAMVPM5_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM6_EL2]
-	msr	MPAMVPM6_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPM7_EL2]
-	msr	MPAMVPM7_EL2, x9
-
-	ldr	x9, [x0, #CTX_MPAMVPMV_EL2]
-	msr	MPAMVPMV_EL2, x9
-
-	ldr	x9, [x0, #CTX_RMR_EL2]
-	msr	rmr_el2, x9
-
-	ldr	x9, [x0, #CTX_SCTLR_EL2]
+	ldp	x17, x9, [x0, #CTX_PMSCR_EL2]
+	msr	PMSCR_EL2, x17
 	msr	sctlr_el2, x9
 
-	ldr	x9, [x0, #CTX_SPSR_EL2]
-	msr	spsr_el2, x9
+	ldp	x10, x11, [x0, #CTX_SPSR_EL2]
+	msr	spsr_el2, x10
+	msr	sp_el2, x11
 
-	ldr	x9, [x0, #CTX_SP_EL2]
-	msr	sp_el2, x9
+	ldp	x12, x13, [x0, #CTX_TCR_EL2]
+	msr	tcr_el2, x12
+	msr	TRFCR_EL2, x13
 
-	ldr	x9, [x0, #CTX_TCR_EL2]
-	msr	tcr_el2, x9
+	ldp	x14, x15, [x0, #CTX_TTBR0_EL2]
+	msr	ttbr0_el2, x14
+	msr	vbar_el2, x15
 
-	ldr	x9, [x0, #CTX_TPIDR_EL2]
-	msr	tpidr_el2, x9
+	ldp	x16, x17, [x0, #CTX_VMPIDR_EL2]
+	msr	vmpidr_el2, x16
+	msr	vpidr_el2, x17
 
-	ldr	x9, [x0, #CTX_TTBR0_EL2]
-	msr	ttbr0_el2, x9
-
-	ldr	x9, [x0, #CTX_VBAR_EL2]
-	msr	vbar_el2, x9
-
-	ldr	x9, [x0, #CTX_VMPIDR_EL2]
-	msr	vmpidr_el2, x9
-
-	ldr	x9, [x0, #CTX_VPIDR_EL2]
-	msr	vpidr_el2, x9
-
-	ldr	x9, [x0, #CTX_VTCR_EL2]
+	ldp	x9, x10, [x0, #CTX_VTCR_EL2]
 	msr	vtcr_el2, x9
+	msr	vttbr_el2, x10
 
-	ldr	x9, [x0, #CTX_VTTBR_EL2]
-	msr	vttbr_el2, x9
+#if CTX_INCLUDE_MTE_REGS
+	ldr	x11, [x0, #CTX_TFSR_EL2]
+	msr	TFSR_EL2, x11
+#endif
 
-	ldr	x9, [x0, #CTX_ZCR_EL2]
-	msr	ZCR_EL2, x9
+#if ENABLE_MPAM_FOR_LOWER_ELS
+	ldp	x9, x10, [x0, #CTX_MPAM2_EL2]
+	msr	MPAM2_EL2, x9
+	msr	MPAMHCR_EL2, x10
+
+	ldp	x11, x12, [x0, #CTX_MPAMVPM0_EL2]
+	msr	MPAMVPM0_EL2, x11
+	msr	MPAMVPM1_EL2, x12
+
+	ldp	x13, x14, [x0, #CTX_MPAMVPM2_EL2]
+	msr	MPAMVPM2_EL2, x13
+	msr	MPAMVPM3_EL2, x14
+
+	ldp	x15, x16, [x0, #CTX_MPAMVPM4_EL2]
+	msr	MPAMVPM4_EL2, x15
+	msr	MPAMVPM5_EL2, x16
+
+	ldp	x17, x9, [x0, #CTX_MPAMVPM6_EL2]
+	msr	MPAMVPM6_EL2, x17
+	msr	MPAMVPM7_EL2, x9
+
+	ldr	x10, [x0, #CTX_MPAMVPMV_EL2]
+	msr	MPAMVPMV_EL2, x10
+#endif
+
+#if ARM_ARCH_AT_LEAST(8, 6)
+	ldp	x11, x12, [x0, #CTX_HAFGRTR_EL2]
+	msr	HAFGRTR_EL2, x11
+	msr	HDFGRTR_EL2, x12
+
+	ldp	x13, x14, [x0, #CTX_HDFGWTR_EL2]
+	msr	HDFGWTR_EL2, x13
+	msr	HFGITR_EL2, x14
+
+	ldp	x15, x16, [x0, #CTX_HFGRTR_EL2]
+	msr	HFGRTR_EL2, x15
+	msr	HFGWTR_EL2, x16
+
+	ldr	x17, [x0, #CTX_CNTPOFF_EL2]
+	msr	CNTPOFF_EL2, x17
+#endif
+
+#if ARM_ARCH_AT_LEAST(8, 4)
+	ldp	x9, x10, [x0, #CTX_CNTHPS_CTL_EL2]
+	msr	cnthps_ctl_el2, x9
+	msr	cnthps_cval_el2, x10
+
+	ldp	x11, x12, [x0, #CTX_CNTHPS_TVAL_EL2]
+	msr	cnthps_tval_el2, x11
+	msr	cnthvs_ctl_el2, x12
+
+	ldp	x13, x14, [x0, #CTX_CNTHVS_CVAL_EL2]
+	msr	cnthvs_cval_el2, x13
+	msr	cnthvs_tval_el2, x14
+
+	ldp	x15, x16, [x0, #CTX_CNTHV_CTL_EL2]
+	msr	cnthv_ctl_el2, x15
+	msr	cnthv_cval_el2, x16
+
+	ldp	x17, x9, [x0, #CTX_CNTHV_TVAL_EL2]
+	msr	cnthv_tval_el2, x17
+	msr	contextidr_el2, x9
+
+	ldr	x10, [x0, #CTX_SDER32_EL2]
+	msr	sder32_el2, x10
+
+	ldr	x11, [x0, #CTX_TTBR1_EL2]
+	msr	ttbr1_el2, x11
+
+	ldr	x12, [x0, #CTX_VDISR_EL2]
+	msr	vdisr_el2, x12
+
+	ldr	x13, [x0, #CTX_VNCR_EL2]
+	msr	vncr_el2, x13
+
+	ldr	x14, [x0, #CTX_VSESR_EL2]
+	msr	vsesr_el2, x14
+
+	ldr	x15, [x0, #CTX_VSTCR_EL2]
+	msr	vstcr_el2, x15
+
+	ldr	x16, [x0, #CTX_VSTTBR_EL2]
+	msr	vsttbr_el2, x16
+#endif
+
+#if ARM_ARCH_AT_LEAST(8, 5)
+	ldr	x17, [x0, #CTX_SCXTNUM_EL2]
+	msr	scxtnum_el2, x17
+#endif
 
 	ret
 endfunc el2_sysregs_context_restore
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index f59bcfc..0314a85 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -234,7 +234,7 @@
 	 * and other EL2 registers are set up by cm_prepare_ns_entry() as they
 	 * are not part of the stored cpu_context.
 	 */
-	write_ctx_reg(get_sysregs_ctx(ctx), CTX_SCTLR_EL1, sctlr_elx);
+	write_ctx_reg(get_el1_sysregs_ctx(ctx), CTX_SCTLR_EL1, sctlr_elx);
 
 	/*
 	 * Base the context ACTLR_EL1 on the current value, as it is
@@ -244,7 +244,7 @@
 	 * be zero.
 	 */
 	actlr_elx = read_actlr_el1();
-	write_ctx_reg((get_sysregs_ctx(ctx)), (CTX_ACTLR_EL1), (actlr_elx));
+	write_ctx_reg((get_el1_sysregs_ctx(ctx)), (CTX_ACTLR_EL1), (actlr_elx));
 
 	/*
 	 * Populate EL3 state so that we've the right context
@@ -336,7 +336,7 @@
 						 CTX_SCR_EL3);
 		if ((scr_el3 & SCR_HCE_BIT) != 0U) {
 			/* Use SCTLR_EL1.EE value to initialise sctlr_el2 */
-			sctlr_elx = read_ctx_reg(get_sysregs_ctx(ctx),
+			sctlr_elx = read_ctx_reg(get_el1_sysregs_ctx(ctx),
 							   CTX_SCTLR_EL1);
 			sctlr_elx &= SCTLR_EE_BIT;
 			sctlr_elx |= SCTLR_EL2_RES1;
@@ -549,7 +549,7 @@
 		ctx = cm_get_context(security_state);
 		assert(ctx != NULL);
 
-		el2_sysregs_context_save(get_sysregs_ctx(ctx));
+		el2_sysregs_context_save(get_el2_sysregs_ctx(ctx));
 	}
 }
 
@@ -571,7 +571,7 @@
 		ctx = cm_get_context(security_state);
 		assert(ctx != NULL);
 
-		el2_sysregs_context_restore(get_sysregs_ctx(ctx));
+		el2_sysregs_context_restore(get_el2_sysregs_ctx(ctx));
 	}
 }
 #endif /* CTX_INCLUDE_EL2_REGS */
@@ -588,7 +588,7 @@
 	ctx = cm_get_context(security_state);
 	assert(ctx != NULL);
 
-	el1_sysregs_context_save(get_sysregs_ctx(ctx));
+	el1_sysregs_context_save(get_el1_sysregs_ctx(ctx));
 
 #if IMAGE_BL31
 	if (security_state == SECURE)
@@ -605,7 +605,7 @@
 	ctx = cm_get_context(security_state);
 	assert(ctx != NULL);
 
-	el1_sysregs_context_restore(get_sysregs_ctx(ctx));
+	el1_sysregs_context_restore(get_el1_sysregs_ctx(ctx));
 
 #if IMAGE_BL31
 	if (security_state == SECURE)