refactor(cm): couple el2 registers with dependent feature flags

Currently the EL2 part of the context structure (el2_sysregs_t), is
mostly feature dependent.

For instance, CTX_HCRX_EL2 is only needed when FEAT_HCX
(ENABLE_FEAT_HCX=1) is set, but the entry is unconditionally added
in the EL2 context structure and thereby consuming memory even in
build configurations where FEAT_HCX is disabled.

Henceforth, all such context entries should be coupled/tied with
their respective feature enables and be optimized away when unused.
This would reduce the context memory allocation for platforms, that
dont enable/support all the architectural features at once.

Further, converting the assembly context-offset entries into
a c structure relies on garbage collection of the linker
removing unreferenced structures from memory, as well as aiding
in readability and future maintenance.

Change-Id: I0cf49498ee3033cb6f3ee3810331121b26627783
Signed-off-by: Jayanth Dodderi Chidanand <jayanthdodderi.chidanand@arm.com>
diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index d5bd890..44efee5 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -7,6 +7,7 @@
 #ifndef CONTEXT_H
 #define CONTEXT_H
 
+#include <lib/el3_runtime/context_el2.h>
 #include <lib/el3_runtime/cpu_data.h>
 #include <lib/utils_def.h>
 
@@ -157,102 +158,11 @@
  */
 #define CTX_EL1_SYSREGS_END		CTX_MTE_REGS_END
 
-/*
- * EL2 register set
- */
-
-#if CTX_INCLUDE_EL2_REGS
-/* For later discussion
- * ICH_AP0R<n>_EL2
- * ICH_AP1R<n>_EL2
- * AMEVCNTVOFF0<n>_EL2
- * AMEVCNTVOFF1<n>_EL2
- * ICH_LR<n>_EL2
- */
-#define CTX_EL2_SYSREGS_OFFSET	(CTX_EL1_SYSREGS_OFFSET + CTX_EL1_SYSREGS_END)
-
-#define CTX_ACTLR_EL2		U(0x0)
-#define CTX_AFSR0_EL2		U(0x8)
-#define CTX_AFSR1_EL2		U(0x10)
-#define CTX_AMAIR_EL2		U(0x18)
-#define CTX_CNTHCTL_EL2		U(0x20)
-#define CTX_CNTVOFF_EL2		U(0x28)
-#define CTX_CPTR_EL2		U(0x30)
-#define CTX_DBGVCR32_EL2	U(0x38)
-#define CTX_ELR_EL2		U(0x40)
-#define CTX_ESR_EL2		U(0x48)
-#define CTX_FAR_EL2		U(0x50)
-#define CTX_HACR_EL2		U(0x58)
-#define CTX_HCR_EL2		U(0x60)
-#define CTX_HPFAR_EL2		U(0x68)
-#define CTX_HSTR_EL2		U(0x70)
-#define CTX_ICC_SRE_EL2		U(0x78)
-#define CTX_ICH_HCR_EL2		U(0x80)
-#define CTX_ICH_VMCR_EL2	U(0x88)
-#define CTX_MAIR_EL2		U(0x90)
-#define CTX_MDCR_EL2		U(0x98)
-#define CTX_PMSCR_EL2		U(0xa0)
-#define CTX_SCTLR_EL2		U(0xa8)
-#define CTX_SPSR_EL2		U(0xb0)
-#define CTX_SP_EL2		U(0xb8)
-#define CTX_TCR_EL2		U(0xc0)
-#define CTX_TPIDR_EL2		U(0xc8)
-#define CTX_TTBR0_EL2		U(0xd0)
-#define CTX_VBAR_EL2		U(0xd8)
-#define CTX_VMPIDR_EL2		U(0xe0)
-#define CTX_VPIDR_EL2		U(0xe8)
-#define CTX_VTCR_EL2		U(0xf0)
-#define CTX_VTTBR_EL2		U(0xf8)
-
-// Only if MTE registers in use
-#define CTX_TFSR_EL2		U(0x100)
-
-// Starting with Armv8.6
-#define CTX_HDFGRTR_EL2		U(0x108)
-#define CTX_HAFGRTR_EL2		U(0x110)
-#define CTX_HDFGWTR_EL2		U(0x118)
-#define CTX_HFGITR_EL2		U(0x120)
-#define CTX_HFGRTR_EL2		U(0x128)
-#define CTX_HFGWTR_EL2		U(0x130)
-#define CTX_CNTPOFF_EL2		U(0x138)
-
-// Starting with Armv8.4
-#define CTX_CONTEXTIDR_EL2	U(0x140)
-#define CTX_TTBR1_EL2		U(0x148)
-#define CTX_VDISR_EL2		U(0x150)
-#define CTX_VSESR_EL2		U(0x158)
-#define CTX_VNCR_EL2		U(0x160)
-#define CTX_TRFCR_EL2		U(0x168)
-
-// Starting with Armv8.5
-#define CTX_SCXTNUM_EL2		U(0x170)
-
-// Register for FEAT_HCX
-#define CTX_HCRX_EL2            U(0x178)
-
-// Starting with Armv8.9
-#define CTX_TCR2_EL2            U(0x180)
-#define CTX_POR_EL2             U(0x188)
-#define CTX_PIRE0_EL2           U(0x190)
-#define CTX_PIR_EL2             U(0x198)
-#define CTX_S2PIR_EL2		U(0x1a0)
-#define CTX_GCSCR_EL2           U(0x1a8)
-#define CTX_GCSPR_EL2           U(0x1b0)
-
-/* Align to the next 16 byte boundary */
-#define CTX_EL2_SYSREGS_END	U(0x1c0)
-
-#endif /* CTX_INCLUDE_EL2_REGS */
-
 /*******************************************************************************
  * Constants that allow assembler code to access members of and the 'fp_regs'
  * structure at their correct offsets.
  ******************************************************************************/
-#if CTX_INCLUDE_EL2_REGS
-# define CTX_FPREGS_OFFSET	(CTX_EL2_SYSREGS_OFFSET + CTX_EL2_SYSREGS_END)
-#else
 # define CTX_FPREGS_OFFSET	(CTX_EL1_SYSREGS_OFFSET + CTX_EL1_SYSREGS_END)
-#endif
 #if CTX_INCLUDE_FPREGS
 #define CTX_FP_Q0		U(0x0)
 #define CTX_FP_Q1		U(0x10)
@@ -293,10 +203,10 @@
 #define CTX_FPREGS_END		U(0x220) /* Align to the next 16 byte boundary */
 #else
 #define CTX_FPREGS_END		U(0x210) /* Align to the next 16 byte boundary */
-#endif
+#endif /* CTX_INCLUDE_AARCH32_REGS */
 #else
 #define CTX_FPREGS_END		U(0)
-#endif
+#endif /* CTX_INCLUDE_FPREGS */
 
 /*******************************************************************************
  * Registers related to CVE-2018-3639
@@ -373,9 +283,7 @@
 /* Constants to determine the size of individual context structures */
 #define CTX_GPREG_ALL		(CTX_GPREGS_END >> DWORD_SHIFT)
 #define CTX_EL1_SYSREGS_ALL	(CTX_EL1_SYSREGS_END >> DWORD_SHIFT)
-#if CTX_INCLUDE_EL2_REGS
-# define CTX_EL2_SYSREGS_ALL	(CTX_EL2_SYSREGS_END >> DWORD_SHIFT)
-#endif
+
 #if CTX_INCLUDE_FPREGS
 # define CTX_FPREG_ALL		(CTX_FPREGS_END >> DWORD_SHIFT)
 #endif
@@ -403,15 +311,6 @@
  */
 DEFINE_REG_STRUCT(el1_sysregs, CTX_EL1_SYSREGS_ALL);
 
-
-/*
- * AArch64 EL2 system register context structure for preserving the
- * architectural state during world switches.
- */
-#if CTX_INCLUDE_EL2_REGS
-DEFINE_REG_STRUCT(el2_sysregs, CTX_EL2_SYSREGS_ALL);
-#endif
-
 /*
  * AArch64 floating point register context structure for preserving
  * the floating point state during switches from one security state to
@@ -460,19 +359,24 @@
 	gp_regs_t gpregs_ctx;
 	el3_state_t el3state_ctx;
 	el1_sysregs_t el1_sysregs_ctx;
-#if CTX_INCLUDE_EL2_REGS
-	el2_sysregs_t el2_sysregs_ctx;
-#endif
+
 #if CTX_INCLUDE_FPREGS
 	fp_regs_t fpregs_ctx;
 #endif
 	cve_2018_3639_t cve_2018_3639_ctx;
+
 #if CTX_INCLUDE_PAUTH_REGS
 	pauth_t pauth_ctx;
 #endif
+
 #if CTX_INCLUDE_MPAM_REGS
 	mpam_t	mpam_ctx;
 #endif
+
+#if CTX_INCLUDE_EL2_REGS
+	el2_sysregs_t el2_sysregs_ctx;
+#endif
+
 } cpu_context_t;
 
 /*
@@ -512,28 +416,30 @@
  */
 CASSERT(CTX_GPREGS_OFFSET == __builtin_offsetof(cpu_context_t, gpregs_ctx),
 	assert_core_context_gp_offset_mismatch);
+
+CASSERT(CTX_EL3STATE_OFFSET == __builtin_offsetof(cpu_context_t, el3state_ctx),
+	assert_core_context_el3state_offset_mismatch);
+
 CASSERT(CTX_EL1_SYSREGS_OFFSET == __builtin_offsetof(cpu_context_t, el1_sysregs_ctx),
 	assert_core_context_el1_sys_offset_mismatch);
-#if CTX_INCLUDE_EL2_REGS
-CASSERT(CTX_EL2_SYSREGS_OFFSET == __builtin_offsetof(cpu_context_t, el2_sysregs_ctx),
-	assert_core_context_el2_sys_offset_mismatch);
-#endif
+
 #if CTX_INCLUDE_FPREGS
 CASSERT(CTX_FPREGS_OFFSET == __builtin_offsetof(cpu_context_t, fpregs_ctx),
 	assert_core_context_fp_offset_mismatch);
-#endif
-CASSERT(CTX_EL3STATE_OFFSET == __builtin_offsetof(cpu_context_t, el3state_ctx),
-	assert_core_context_el3state_offset_mismatch);
+#endif /* CTX_INCLUDE_FPREGS */
+
 CASSERT(CTX_CVE_2018_3639_OFFSET == __builtin_offsetof(cpu_context_t, cve_2018_3639_ctx),
 	assert_core_context_cve_2018_3639_offset_mismatch);
+
 #if CTX_INCLUDE_PAUTH_REGS
 CASSERT(CTX_PAUTH_REGS_OFFSET == __builtin_offsetof(cpu_context_t, pauth_ctx),
 	assert_core_context_pauth_offset_mismatch);
-#endif
+#endif /* CTX_INCLUDE_PAUTH_REGS */
+
 #if CTX_INCLUDE_MPAM_REGS
 CASSERT(CTX_MPAM_REGS_OFFSET == __builtin_offsetof(cpu_context_t, mpam_ctx),
 	assert_core_context_mpam_offset_mismatch);
-#endif
+#endif /* CTX_INCLUDE_MPAM_REGS */
 
 /*
  * Helper macro to set the general purpose registers that correspond to
diff --git a/include/lib/el3_runtime/context_el2.h b/include/lib/el3_runtime/context_el2.h
new file mode 100644
index 0000000..4ad9634
--- /dev/null
+++ b/include/lib/el3_runtime/context_el2.h
@@ -0,0 +1,319 @@
+/*
+ * Copyright (c) 2024, Arm Limited and Contributors. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#ifndef CONTEXT_EL2_H
+#define CONTEXT_EL2_H
+
+#ifndef __ASSEMBLER__
+/*******************************************************************************
+ * EL2 Registers:
+ * AArch64 EL2 system register context structure for preserving the
+ * architectural state during world switches.
+ ******************************************************************************/
+#if CTX_INCLUDE_EL2_REGS
+typedef struct el2_common_regs {
+	uint64_t actlr_el2;
+	uint64_t afsr0_el2;
+	uint64_t afsr1_el2;
+	uint64_t amair_el2;
+	uint64_t cnthctl_el2;
+	uint64_t cntvoff_el2;
+	uint64_t cptr_el2;
+	uint64_t dbgvcr32_el2;
+	uint64_t elr_el2;
+	uint64_t esr_el2;
+	uint64_t far_el2;
+	uint64_t hacr_el2;
+	uint64_t hcr_el2;
+	uint64_t hpfar_el2;
+	uint64_t hstr_el2;
+	uint64_t icc_sre_el2;
+	uint64_t ich_hcr_el2;
+	uint64_t ich_vmcr_el2;
+	uint64_t mair_el2;
+	uint64_t mdcr_el2;
+	uint64_t pmscr_el2;
+	uint64_t sctlr_el2;
+	uint64_t spsr_el2;
+	uint64_t sp_el2;
+	uint64_t tcr_el2;
+	uint64_t tpidr_el2;
+	uint64_t ttbr0_el2;
+	uint64_t vbar_el2;
+	uint64_t vmpidr_el2;
+	uint64_t vpidr_el2;
+	uint64_t vtcr_el2;
+	uint64_t vttbr_el2;
+} el2_common_regs_t;
+
+typedef struct el2_mte_regs {
+	uint64_t tfsr_el2;
+} el2_mte_regs_t;
+
+typedef struct el2_fgt_regs {
+	uint64_t hdfgrtr_el2;
+	uint64_t hafgrtr_el2;
+	uint64_t hdfgwtr_el2;
+	uint64_t hfgitr_el2;
+	uint64_t hfgrtr_el2;
+	uint64_t hfgwtr_el2;
+} el2_fgt_regs_t;
+
+typedef struct el2_ecv_regs {
+	uint64_t cntpoff_el2;
+} el2_ecv_regs_t;
+
+typedef struct el2_vhe_regs {
+	uint64_t contextidr_el2;
+	uint64_t ttbr1_el2;
+} el2_vhe_regs_t;
+
+typedef struct el2_ras_regs {
+	uint64_t vdisr_el2;
+	uint64_t vsesr_el2;
+} el2_ras_regs_t;
+
+typedef struct el2_neve_regs {
+	uint64_t vncr_el2;
+} el2_neve_regs_t;
+
+typedef struct el2_trf_regs {
+	uint64_t trfcr_el2;
+} el2_trf_regs_t;
+
+typedef struct el2_csv2_regs {
+	uint64_t scxtnum_el2;
+} el2_csv2_regs_t;
+
+typedef struct el2_hcx_regs {
+	uint64_t hcrx_el2;
+} el2_hcx_regs_t;
+
+typedef struct el2_tcr2_regs {
+	uint64_t tcr2_el2;
+} el2_tcr2_regs_t;
+
+typedef struct el2_sxpoe_regs {
+	uint64_t por_el2;
+} el2_sxpoe_regs_t;
+
+typedef struct el2_sxpie_regs {
+	uint64_t pire0_el2;
+	uint64_t pir_el2;
+} el2_sxpie_regs_t;
+
+typedef struct el2_s2pie_regs {
+	uint64_t s2pir_el2;
+} el2_s2pie_regs_t;
+
+typedef struct el2_gcs_regs {
+	uint64_t gcscr_el2;
+	uint64_t gcspr_el2;
+} el2_gcs_regs_t;
+
+typedef struct el2_sysregs {
+
+	el2_common_regs_t common;
+
+#if ENABLE_FEAT_MTE
+	el2_mte_regs_t mte;
+#endif
+
+#if ENABLE_FEAT_FGT
+	el2_fgt_regs_t fgt;
+#endif
+
+#if ENABLE_FEAT_ECV
+	el2_ecv_regs_t ecv;
+#endif
+
+#if ENABLE_FEAT_VHE
+	el2_vhe_regs_t vhe;
+#endif
+
+#if ENABLE_FEAT_RAS
+	el2_ras_regs_t ras;
+#endif
+
+#if CTX_INCLUDE_NEVE_REGS
+	el2_neve_regs_t neve;
+#endif
+
+#if ENABLE_TRF_FOR_NS
+	el2_trf_regs_t trf;
+#endif
+
+#if ENABLE_FEAT_CSV2_2
+	el2_csv2_regs_t csv2;
+#endif
+
+#if ENABLE_FEAT_HCX
+	el2_hcx_regs_t hcx;
+#endif
+
+#if ENABLE_FEAT_TCR2
+	el2_tcr2_regs_t tcr2;
+#endif
+
+#if (ENABLE_FEAT_S1POE || ENABLE_FEAT_S2POE)
+	el2_sxpoe_regs_t sxpoe;
+#endif
+
+#if (ENABLE_FEAT_S1PIE || ENABLE_FEAT_S2PIE)
+	el2_sxpie_regs_t sxpie;
+#endif
+
+#if ENABLE_FEAT_S2PIE
+	el2_s2pie_regs_t s2pie;
+#endif
+
+#if ENABLE_FEAT_GCS
+	el2_gcs_regs_t gcs;
+#endif
+
+} el2_sysregs_t;
+
+/*
+ * Macros to access members related to individual features of the el2_sysregs_t
+ * structures.
+ */
+#define read_el2_ctx_common(ctx, reg)		(((ctx)->common).reg)
+
+#define write_el2_ctx_common(ctx, reg, val)	((((ctx)->common).reg)	\
+							= (uint64_t) (val))
+
+#if ENABLE_FEAT_MTE
+#define read_el2_ctx_mte(ctx, reg)		(((ctx)->mte).reg)
+#define write_el2_ctx_mte(ctx, reg, val)	((((ctx)->mte).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_mte(ctx, reg)		ULL(0)
+#define write_el2_ctx_mte(ctx, reg, val)
+#endif /* ENABLE_FEAT_MTE */
+
+#if ENABLE_FEAT_FGT
+#define read_el2_ctx_fgt(ctx, reg)		(((ctx)->fgt).reg)
+#define write_el2_ctx_fgt(ctx, reg, val)	((((ctx)->fgt).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_fgt(ctx, reg)		ULL(0)
+#define write_el2_ctx_fgt(ctx, reg, val)
+#endif /* ENABLE_FEAT_FGT */
+
+#if ENABLE_FEAT_ECV
+#define read_el2_ctx_ecv(ctx, reg)		(((ctx)->ecv).reg)
+#define write_el2_ctx_ecv(ctx, reg, val)	((((ctx)->ecv).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_ecv(ctx, reg)		ULL(0)
+#define write_el2_ctx_ecv(ctx, reg, val)
+#endif /* ENABLE_FEAT_ECV */
+
+#if ENABLE_FEAT_VHE
+#define read_el2_ctx_vhe(ctx, reg)		(((ctx)->vhe).reg)
+#define write_el2_ctx_vhe(ctx, reg, val)	((((ctx)->vhe).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_vhe(ctx, reg)		ULL(0)
+#define write_el2_ctx_vhe(ctx, reg, val)
+#endif /* ENABLE_FEAT_VHE */
+
+#if ENABLE_FEAT_RAS
+#define read_el2_ctx_ras(ctx, reg)		(((ctx)->ras).reg)
+#define write_el2_ctx_ras(ctx, reg, val)	((((ctx)->ras).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_ras(ctx, reg)		ULL(0)
+#define write_el2_ctx_ras(ctx, reg, val)
+#endif /* ENABLE_FEAT_RAS */
+
+#if CTX_INCLUDE_NEVE_REGS
+#define read_el2_ctx_neve(ctx, reg)		(((ctx)->neve).reg)
+#define write_el2_ctx_neve(ctx, reg, val)	((((ctx)->neve).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_neve(ctx, reg)		ULL(0)
+#define write_el2_ctx_neve(ctx, reg, val)
+#endif /* CTX_INCLUDE_NEVE_REGS */
+
+#if ENABLE_TRF_FOR_NS
+#define read_el2_ctx_trf(ctx, reg)		(((ctx)->trf).reg)
+#define write_el2_ctx_trf(ctx, reg, val)	((((ctx)->trf).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_trf(ctx, reg)		ULL(0)
+#define write_el2_ctx_trf(ctx, reg, val)
+#endif /* ENABLE_TRF_FOR_NS */
+
+#if ENABLE_FEAT_CSV2_2
+#define read_el2_ctx_csv2_2(ctx, reg)		(((ctx)->csv2).reg)
+#define write_el2_ctx_csv2_2(ctx, reg, val)	((((ctx)->csv2).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_csv2_2(ctx, reg)		ULL(0)
+#define write_el2_ctx_csv2_2(ctx, reg, val)
+#endif /* ENABLE_FEAT_CSV2_2 */
+
+#if ENABLE_FEAT_HCX
+#define read_el2_ctx_hcx(ctx, reg)		(((ctx)->hcx).reg)
+#define write_el2_ctx_hcx(ctx, reg, val)	((((ctx)->hcx).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_hcx(ctx, reg)		ULL(0)
+#define write_el2_ctx_hcx(ctx, reg, val)
+#endif /* ENABLE_FEAT_HCX */
+
+#if ENABLE_FEAT_TCR2
+#define read_el2_ctx_tcr2(ctx, reg)		(((ctx)->tcr2).reg)
+#define write_el2_ctx_tcr2(ctx, reg, val)	((((ctx)->tcr2).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_tcr2(ctx, reg)		ULL(0)
+#define write_el2_ctx_tcr2(ctx, reg, val)
+#endif /* ENABLE_FEAT_TCR2 */
+
+#if (ENABLE_FEAT_S1POE || ENABLE_FEAT_S2POE)
+#define read_el2_ctx_sxpoe(ctx, reg)		(((ctx)->sxpoe).reg)
+#define write_el2_ctx_sxpoe(ctx, reg, val)	((((ctx)->sxpoe).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_sxpoe(ctx, reg)		ULL(0)
+#define write_el2_ctx_sxpoe(ctx, reg, val)
+#endif /*(ENABLE_FEAT_S1POE || ENABLE_FEAT_S2POE) */
+
+#if (ENABLE_FEAT_S1PIE || ENABLE_FEAT_S2PIE)
+#define read_el2_ctx_sxpie(ctx, reg)		(((ctx)->sxpie).reg)
+#define write_el2_ctx_sxpie(ctx, reg, val)	((((ctx)->sxpie).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_sxpie(ctx, reg)		ULL(0)
+#define write_el2_ctx_sxpie(ctx, reg, val)
+#endif /*(ENABLE_FEAT_S1PIE || ENABLE_FEAT_S2PIE) */
+
+#if ENABLE_FEAT_S2PIE
+#define read_el2_ctx_s2pie(ctx, reg)		(((ctx)->s2pie).reg)
+#define write_el2_ctx_s2pie(ctx, reg, val)	((((ctx)->s2pie).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_s2pie(ctx, reg)		ULL(0)
+#define write_el2_ctx_s2pie(ctx, reg, val)
+#endif /* ENABLE_FEAT_S2PIE */
+
+#if ENABLE_FEAT_GCS
+#define read_el2_ctx_gcs(ctx, reg)		(((ctx)->gcs).reg)
+#define write_el2_ctx_gcs(ctx, reg, val)	((((ctx)->gcs).reg)	\
+							= (uint64_t) (val))
+#else
+#define read_el2_ctx_gcs(ctx, reg)		ULL(0)
+#define write_el2_ctx_gcs(ctx, reg, val)
+#endif /* ENABLE_FEAT_GCS */
+
+#endif /* CTX_INCLUDE_EL2_REGS */
+/******************************************************************************/
+
+#endif /* __ASSEMBLER__ */
+
+#endif /* CONTEXT_EL2_H */
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 922b2cf..5c8f03c 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -262,10 +262,10 @@
 	 * Initialize SCTLR_EL2 context register using Endianness value
 	 * taken from the entrypoint attribute.
 	 */
-	u_register_t sctlr_el2 = (EP_GET_EE(ep->h.attr) != 0U) ? SCTLR_EE_BIT : 0UL;
-	sctlr_el2 |= SCTLR_EL2_RES1;
-	write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_SCTLR_EL2,
-			sctlr_el2);
+	u_register_t sctlr_el2_val = (EP_GET_EE(ep->h.attr) != 0U) ? SCTLR_EE_BIT : 0UL;
+	sctlr_el2_val |= SCTLR_EL2_RES1;
+	write_el2_ctx_common(get_el2_sysregs_ctx(ctx), sctlr_el2, sctlr_el2_val);
+
 
 	if (is_feat_hcx_supported()) {
 		/*
@@ -276,7 +276,7 @@
 		 * this feature if not properly initialized, especially when
 		 * it comes to those bits that enable/disable traps.
 		 */
-		write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_HCRX_EL2,
+		write_el2_ctx_hcx(get_el2_sysregs_ctx(ctx), hcrx_el2,
 			HCRX_EL2_INIT_VAL);
 	}
 
@@ -286,13 +286,14 @@
 		 * systems unaware of FEAT_FGT do not get trapped due to their lack
 		 * of initialization for this feature.
 		 */
-		write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_HFGITR_EL2,
+		write_el2_ctx_fgt(get_el2_sysregs_ctx(ctx), hfgitr_el2,
 			HFGITR_EL2_INIT_VAL);
-		write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_HFGRTR_EL2,
+		write_el2_ctx_fgt(get_el2_sysregs_ctx(ctx), hfgrtr_el2,
 			HFGRTR_EL2_INIT_VAL);
-		write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_HFGWTR_EL2,
+		write_el2_ctx_fgt(get_el2_sysregs_ctx(ctx), hfgwtr_el2,
 			HFGWTR_EL2_INIT_VAL);
 	}
+
 #endif /* CTX_INCLUDE_EL2_REGS */
 
 	manage_extensions_nonsecure(ctx);
@@ -330,9 +331,9 @@
 	 * These bits are set in the gicv3 driver. Losing them (especially the
 	 * SRE bit) is problematic for all worlds. Henceforth recreate them.
 	 */
-	u_register_t icc_sre_el2 = ICC_SRE_DIB_BIT | ICC_SRE_DFB_BIT |
+	u_register_t icc_sre_el2_val = ICC_SRE_DIB_BIT | ICC_SRE_DFB_BIT |
 				   ICC_SRE_EN_BIT | ICC_SRE_SRE_BIT;
-	write_ctx_reg(el2_ctx, CTX_ICC_SRE_EL2, icc_sre_el2);
+	write_el2_ctx_common(el2_ctx, icc_sre_el2, icc_sre_el2_val);
 #endif /* CTX_INCLUDE_EL2_REGS */
 
 	/* Start with a clean SCR_EL3 copy as all relevant values are set */
@@ -987,7 +988,6 @@
 			}
 		}
 
-
 		if ((scr_el3 & SCR_HCE_BIT) != 0U) {
 			/* Use SCTLR_EL1.EE value to initialise sctlr_el2 */
 			sctlr_elx = read_ctx_reg(get_el1_sysregs_ctx(ctx),
@@ -1016,26 +1016,26 @@
 
 static void el2_sysregs_context_save_fgt(el2_sysregs_t *ctx)
 {
-	write_ctx_reg(ctx, CTX_HDFGRTR_EL2, read_hdfgrtr_el2());
+	write_el2_ctx_fgt(ctx, hdfgrtr_el2, read_hdfgrtr_el2());
 	if (is_feat_amu_supported()) {
-		write_ctx_reg(ctx, CTX_HAFGRTR_EL2, read_hafgrtr_el2());
+		write_el2_ctx_fgt(ctx, hafgrtr_el2, read_hafgrtr_el2());
 	}
-	write_ctx_reg(ctx, CTX_HDFGWTR_EL2, read_hdfgwtr_el2());
-	write_ctx_reg(ctx, CTX_HFGITR_EL2, read_hfgitr_el2());
-	write_ctx_reg(ctx, CTX_HFGRTR_EL2, read_hfgrtr_el2());
-	write_ctx_reg(ctx, CTX_HFGWTR_EL2, read_hfgwtr_el2());
+	write_el2_ctx_fgt(ctx, hdfgwtr_el2, read_hdfgwtr_el2());
+	write_el2_ctx_fgt(ctx, hfgitr_el2, read_hfgitr_el2());
+	write_el2_ctx_fgt(ctx, hfgrtr_el2, read_hfgrtr_el2());
+	write_el2_ctx_fgt(ctx, hfgwtr_el2, read_hfgwtr_el2());
 }
 
 static void el2_sysregs_context_restore_fgt(el2_sysregs_t *ctx)
 {
-	write_hdfgrtr_el2(read_ctx_reg(ctx, CTX_HDFGRTR_EL2));
+	write_hdfgrtr_el2(read_el2_ctx_fgt(ctx, hdfgrtr_el2));
 	if (is_feat_amu_supported()) {
-		write_hafgrtr_el2(read_ctx_reg(ctx, CTX_HAFGRTR_EL2));
+		write_hafgrtr_el2(read_el2_ctx_fgt(ctx, hafgrtr_el2));
 	}
-	write_hdfgwtr_el2(read_ctx_reg(ctx, CTX_HDFGWTR_EL2));
-	write_hfgitr_el2(read_ctx_reg(ctx, CTX_HFGITR_EL2));
-	write_hfgrtr_el2(read_ctx_reg(ctx, CTX_HFGRTR_EL2));
-	write_hfgwtr_el2(read_ctx_reg(ctx, CTX_HFGWTR_EL2));
+	write_hdfgwtr_el2(read_el2_ctx_fgt(ctx, hdfgwtr_el2));
+	write_hfgitr_el2(read_el2_ctx_fgt(ctx, hfgitr_el2));
+	write_hfgrtr_el2(read_el2_ctx_fgt(ctx, hfgrtr_el2));
+	write_hfgwtr_el2(read_el2_ctx_fgt(ctx, hfgwtr_el2));
 }
 
 #if CTX_INCLUDE_MPAM_REGS
@@ -1148,38 +1148,37 @@
 static void el2_sysregs_context_save_gic(el2_sysregs_t *ctx)
 {
 #if defined(SPD_spmd) && SPMD_SPM_AT_SEL2
-	write_ctx_reg(ctx, CTX_ICC_SRE_EL2, read_icc_sre_el2());
+	write_el2_ctx_common(ctx, icc_sre_el2, read_icc_sre_el2());
 #else
 	u_register_t scr_el3 = read_scr_el3();
 	write_scr_el3(scr_el3 | SCR_NS_BIT);
 	isb();
 
-	write_ctx_reg(ctx, CTX_ICC_SRE_EL2, read_icc_sre_el2());
+	write_el2_ctx_common(ctx, icc_sre_el2, read_icc_sre_el2());
 
 	write_scr_el3(scr_el3);
 	isb();
-
 #endif
-	write_ctx_reg(ctx, CTX_ICH_HCR_EL2, read_ich_hcr_el2());
-	write_ctx_reg(ctx, CTX_ICH_VMCR_EL2, read_ich_vmcr_el2());
+	write_el2_ctx_common(ctx, ich_hcr_el2, read_ich_hcr_el2());
+	write_el2_ctx_common(ctx, ich_vmcr_el2, read_ich_vmcr_el2());
 }
 
 static void el2_sysregs_context_restore_gic(el2_sysregs_t *ctx)
 {
 #if defined(SPD_spmd) && SPMD_SPM_AT_SEL2
-	write_icc_sre_el2(read_ctx_reg(ctx, CTX_ICC_SRE_EL2));
+	write_icc_sre_el2(read_el2_ctx_common(ctx, icc_sre_el2));
 #else
 	u_register_t scr_el3 = read_scr_el3();
 	write_scr_el3(scr_el3 | SCR_NS_BIT);
 	isb();
 
-	write_icc_sre_el2(read_ctx_reg(ctx, CTX_ICC_SRE_EL2));
+	write_icc_sre_el2(read_el2_ctx_common(ctx, icc_sre_el2));
 
 	write_scr_el3(scr_el3);
 	isb();
 #endif
-	write_ich_hcr_el2(read_ctx_reg(ctx, CTX_ICH_HCR_EL2));
-	write_ich_vmcr_el2(read_ctx_reg(ctx, CTX_ICH_VMCR_EL2));
+	write_ich_hcr_el2(read_el2_ctx_common(ctx, ich_hcr_el2));
+	write_ich_vmcr_el2(read_el2_ctx_common(ctx, ich_vmcr_el2));
 }
 
 /* -----------------------------------------------------
@@ -1190,70 +1189,70 @@
  */
 static void el2_sysregs_context_save_common(el2_sysregs_t *ctx)
 {
-	write_ctx_reg(ctx, CTX_ACTLR_EL2, read_actlr_el2());
-	write_ctx_reg(ctx, CTX_AFSR0_EL2, read_afsr0_el2());
-	write_ctx_reg(ctx, CTX_AFSR1_EL2, read_afsr1_el2());
-	write_ctx_reg(ctx, CTX_AMAIR_EL2, read_amair_el2());
-	write_ctx_reg(ctx, CTX_CNTHCTL_EL2, read_cnthctl_el2());
-	write_ctx_reg(ctx, CTX_CNTVOFF_EL2, read_cntvoff_el2());
-	write_ctx_reg(ctx, CTX_CPTR_EL2, read_cptr_el2());
+	write_el2_ctx_common(ctx, actlr_el2, read_actlr_el2());
+	write_el2_ctx_common(ctx, afsr0_el2, read_afsr0_el2());
+	write_el2_ctx_common(ctx, afsr1_el2, read_afsr1_el2());
+	write_el2_ctx_common(ctx, amair_el2, read_amair_el2());
+	write_el2_ctx_common(ctx, cnthctl_el2, read_cnthctl_el2());
+	write_el2_ctx_common(ctx, cntvoff_el2, read_cntvoff_el2());
+	write_el2_ctx_common(ctx, cptr_el2, read_cptr_el2());
 	if (CTX_INCLUDE_AARCH32_REGS) {
-		write_ctx_reg(ctx, CTX_DBGVCR32_EL2, read_dbgvcr32_el2());
+		write_el2_ctx_common(ctx, dbgvcr32_el2, read_dbgvcr32_el2());
 	}
-	write_ctx_reg(ctx, CTX_ELR_EL2, read_elr_el2());
-	write_ctx_reg(ctx, CTX_ESR_EL2, read_esr_el2());
-	write_ctx_reg(ctx, CTX_FAR_EL2, read_far_el2());
-	write_ctx_reg(ctx, CTX_HACR_EL2, read_hacr_el2());
-	write_ctx_reg(ctx, CTX_HCR_EL2, read_hcr_el2());
-	write_ctx_reg(ctx, CTX_HPFAR_EL2, read_hpfar_el2());
-	write_ctx_reg(ctx, CTX_HSTR_EL2, read_hstr_el2());
-	write_ctx_reg(ctx, CTX_MAIR_EL2, read_mair_el2());
-	write_ctx_reg(ctx, CTX_MDCR_EL2, read_mdcr_el2());
-	write_ctx_reg(ctx, CTX_SCTLR_EL2, read_sctlr_el2());
-	write_ctx_reg(ctx, CTX_SPSR_EL2, read_spsr_el2());
-	write_ctx_reg(ctx, CTX_SP_EL2, read_sp_el2());
-	write_ctx_reg(ctx, CTX_TCR_EL2, read_tcr_el2());
-	write_ctx_reg(ctx, CTX_TPIDR_EL2, read_tpidr_el2());
-	write_ctx_reg(ctx, CTX_TTBR0_EL2, read_ttbr0_el2());
-	write_ctx_reg(ctx, CTX_VBAR_EL2, read_vbar_el2());
-	write_ctx_reg(ctx, CTX_VMPIDR_EL2, read_vmpidr_el2());
-	write_ctx_reg(ctx, CTX_VPIDR_EL2, read_vpidr_el2());
-	write_ctx_reg(ctx, CTX_VTCR_EL2, read_vtcr_el2());
-	write_ctx_reg(ctx, CTX_VTTBR_EL2, read_vttbr_el2());
+	write_el2_ctx_common(ctx, elr_el2, read_elr_el2());
+	write_el2_ctx_common(ctx, esr_el2, read_esr_el2());
+	write_el2_ctx_common(ctx, far_el2, read_far_el2());
+	write_el2_ctx_common(ctx, hacr_el2, read_hacr_el2());
+	write_el2_ctx_common(ctx, hcr_el2, read_hcr_el2());
+	write_el2_ctx_common(ctx, hpfar_el2, read_hpfar_el2());
+	write_el2_ctx_common(ctx, hstr_el2, read_hstr_el2());
+	write_el2_ctx_common(ctx, mair_el2, read_mair_el2());
+	write_el2_ctx_common(ctx, mdcr_el2, read_mdcr_el2());
+	write_el2_ctx_common(ctx, sctlr_el2, read_sctlr_el2());
+	write_el2_ctx_common(ctx, spsr_el2, read_spsr_el2());
+	write_el2_ctx_common(ctx, sp_el2, read_sp_el2());
+	write_el2_ctx_common(ctx, tcr_el2, read_tcr_el2());
+	write_el2_ctx_common(ctx, tpidr_el2, read_tpidr_el2());
+	write_el2_ctx_common(ctx, ttbr0_el2, read_ttbr0_el2());
+	write_el2_ctx_common(ctx, vbar_el2, read_vbar_el2());
+	write_el2_ctx_common(ctx, vmpidr_el2, read_vmpidr_el2());
+	write_el2_ctx_common(ctx, vpidr_el2, read_vpidr_el2());
+	write_el2_ctx_common(ctx, vtcr_el2, read_vtcr_el2());
+	write_el2_ctx_common(ctx, vttbr_el2, read_vttbr_el2());
 }
 
 static void el2_sysregs_context_restore_common(el2_sysregs_t *ctx)
 {
-	write_actlr_el2(read_ctx_reg(ctx, CTX_ACTLR_EL2));
-	write_afsr0_el2(read_ctx_reg(ctx, CTX_AFSR0_EL2));
-	write_afsr1_el2(read_ctx_reg(ctx, CTX_AFSR1_EL2));
-	write_amair_el2(read_ctx_reg(ctx, CTX_AMAIR_EL2));
-	write_cnthctl_el2(read_ctx_reg(ctx, CTX_CNTHCTL_EL2));
-	write_cntvoff_el2(read_ctx_reg(ctx, CTX_CNTVOFF_EL2));
-	write_cptr_el2(read_ctx_reg(ctx, CTX_CPTR_EL2));
+	write_actlr_el2(read_el2_ctx_common(ctx, actlr_el2));
+	write_afsr0_el2(read_el2_ctx_common(ctx, afsr0_el2));
+	write_afsr1_el2(read_el2_ctx_common(ctx, afsr1_el2));
+	write_amair_el2(read_el2_ctx_common(ctx, amair_el2));
+	write_cnthctl_el2(read_el2_ctx_common(ctx, cnthctl_el2));
+	write_cntvoff_el2(read_el2_ctx_common(ctx, cntvoff_el2));
+	write_cptr_el2(read_el2_ctx_common(ctx, cptr_el2));
 	if (CTX_INCLUDE_AARCH32_REGS) {
-		write_dbgvcr32_el2(read_ctx_reg(ctx, CTX_DBGVCR32_EL2));
+		write_dbgvcr32_el2(read_el2_ctx_common(ctx, dbgvcr32_el2));
 	}
-	write_elr_el2(read_ctx_reg(ctx, CTX_ELR_EL2));
-	write_esr_el2(read_ctx_reg(ctx, CTX_ESR_EL2));
-	write_far_el2(read_ctx_reg(ctx, CTX_FAR_EL2));
-	write_hacr_el2(read_ctx_reg(ctx, CTX_HACR_EL2));
-	write_hcr_el2(read_ctx_reg(ctx, CTX_HCR_EL2));
-	write_hpfar_el2(read_ctx_reg(ctx, CTX_HPFAR_EL2));
-	write_hstr_el2(read_ctx_reg(ctx, CTX_HSTR_EL2));
-	write_mair_el2(read_ctx_reg(ctx, CTX_MAIR_EL2));
-	write_mdcr_el2(read_ctx_reg(ctx, CTX_MDCR_EL2));
-	write_sctlr_el2(read_ctx_reg(ctx, CTX_SCTLR_EL2));
-	write_spsr_el2(read_ctx_reg(ctx, CTX_SPSR_EL2));
-	write_sp_el2(read_ctx_reg(ctx, CTX_SP_EL2));
-	write_tcr_el2(read_ctx_reg(ctx, CTX_TCR_EL2));
-	write_tpidr_el2(read_ctx_reg(ctx, CTX_TPIDR_EL2));
-	write_ttbr0_el2(read_ctx_reg(ctx, CTX_TTBR0_EL2));
-	write_vbar_el2(read_ctx_reg(ctx, CTX_VBAR_EL2));
-	write_vmpidr_el2(read_ctx_reg(ctx, CTX_VMPIDR_EL2));
-	write_vpidr_el2(read_ctx_reg(ctx, CTX_VPIDR_EL2));
-	write_vtcr_el2(read_ctx_reg(ctx, CTX_VTCR_EL2));
-	write_vttbr_el2(read_ctx_reg(ctx, CTX_VTTBR_EL2));
+	write_elr_el2(read_el2_ctx_common(ctx, elr_el2));
+	write_esr_el2(read_el2_ctx_common(ctx, esr_el2));
+	write_far_el2(read_el2_ctx_common(ctx, far_el2));
+	write_hacr_el2(read_el2_ctx_common(ctx, hacr_el2));
+	write_hcr_el2(read_el2_ctx_common(ctx, hcr_el2));
+	write_hpfar_el2(read_el2_ctx_common(ctx, hpfar_el2));
+	write_hstr_el2(read_el2_ctx_common(ctx, hstr_el2));
+	write_mair_el2(read_el2_ctx_common(ctx, mair_el2));
+	write_mdcr_el2(read_el2_ctx_common(ctx, mdcr_el2));
+	write_sctlr_el2(read_el2_ctx_common(ctx, sctlr_el2));
+	write_spsr_el2(read_el2_ctx_common(ctx, spsr_el2));
+	write_sp_el2(read_el2_ctx_common(ctx, sp_el2));
+	write_tcr_el2(read_el2_ctx_common(ctx, tcr_el2));
+	write_tpidr_el2(read_el2_ctx_common(ctx, tpidr_el2));
+	write_ttbr0_el2(read_el2_ctx_common(ctx, ttbr0_el2));
+	write_vbar_el2(read_el2_ctx_common(ctx, vbar_el2));
+	write_vmpidr_el2(read_el2_ctx_common(ctx, vmpidr_el2));
+	write_vpidr_el2(read_el2_ctx_common(ctx, vpidr_el2));
+	write_vtcr_el2(read_el2_ctx_common(ctx, vtcr_el2));
+	write_vttbr_el2(read_el2_ctx_common(ctx, vttbr_el2));
 }
 
 /*******************************************************************************
@@ -1272,8 +1271,8 @@
 	el2_sysregs_context_save_common(el2_sysregs_ctx);
 	el2_sysregs_context_save_gic(el2_sysregs_ctx);
 
-	if (is_feat_mte2_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_TFSR_EL2, read_tfsr_el2());
+	if (is_feat_mte_supported()) {
+		write_el2_ctx_mte(el2_sysregs_ctx, tfsr_el2, read_tfsr_el2());
 	}
 
 #if CTX_INCLUDE_MPAM_REGS
@@ -1288,51 +1287,57 @@
 	}
 
 	if (is_feat_ecv_v2_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_CNTPOFF_EL2, read_cntpoff_el2());
+		write_el2_ctx_ecv(el2_sysregs_ctx, cntpoff_el2, read_cntpoff_el2());
 	}
 
 	if (is_feat_vhe_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_CONTEXTIDR_EL2, read_contextidr_el2());
-		write_ctx_reg(el2_sysregs_ctx, CTX_TTBR1_EL2, read_ttbr1_el2());
+		write_el2_ctx_vhe(el2_sysregs_ctx, contextidr_el2,
+					read_contextidr_el2());
+		write_el2_ctx_vhe(el2_sysregs_ctx, ttbr1_el2, read_ttbr1_el2());
 	}
 
 	if (is_feat_ras_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_VDISR_EL2, read_vdisr_el2());
-		write_ctx_reg(el2_sysregs_ctx, CTX_VSESR_EL2, read_vsesr_el2());
+		write_el2_ctx_ras(el2_sysregs_ctx, vdisr_el2, read_vdisr_el2());
+		write_el2_ctx_ras(el2_sysregs_ctx, vsesr_el2, read_vsesr_el2());
 	}
 
 	if (is_feat_nv2_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_VNCR_EL2, read_vncr_el2());
+		write_el2_ctx_neve(el2_sysregs_ctx, vncr_el2, read_vncr_el2());
 	}
 
 	if (is_feat_trf_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_TRFCR_EL2, read_trfcr_el2());
+		write_el2_ctx_trf(el2_sysregs_ctx, trfcr_el2, read_trfcr_el2());
 	}
 
-	/* CSV2 version 2 and above */
 	if (is_feat_csv2_2_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_SCXTNUM_EL2, read_scxtnum_el2());
+		write_el2_ctx_csv2_2(el2_sysregs_ctx, scxtnum_el2,
+					read_scxtnum_el2());
 	}
 
 	if (is_feat_hcx_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_HCRX_EL2, read_hcrx_el2());
+		write_el2_ctx_hcx(el2_sysregs_ctx, hcrx_el2, read_hcrx_el2());
 	}
+
 	if (is_feat_tcr2_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_TCR2_EL2, read_tcr2_el2());
+		write_el2_ctx_tcr2(el2_sysregs_ctx, tcr2_el2, read_tcr2_el2());
 	}
+
 	if (is_feat_sxpie_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_PIRE0_EL2, read_pire0_el2());
-		write_ctx_reg(el2_sysregs_ctx, CTX_PIR_EL2, read_pir_el2());
+		write_el2_ctx_sxpie(el2_sysregs_ctx, pire0_el2, read_pire0_el2());
+		write_el2_ctx_sxpie(el2_sysregs_ctx, pir_el2, read_pir_el2());
 	}
-	if (is_feat_s2pie_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_S2PIR_EL2, read_s2pir_el2());
-	}
+
 	if (is_feat_sxpoe_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_POR_EL2, read_por_el2());
+		write_el2_ctx_sxpoe(el2_sysregs_ctx, por_el2, read_por_el2());
 	}
+
+	if (is_feat_s2pie_supported()) {
+		write_el2_ctx_s2pie(el2_sysregs_ctx, s2pir_el2, read_s2pir_el2());
+	}
+
 	if (is_feat_gcs_supported()) {
-		write_ctx_reg(el2_sysregs_ctx, CTX_GCSPR_EL2, read_gcspr_el2());
-		write_ctx_reg(el2_sysregs_ctx, CTX_GCSCR_EL2, read_gcscr_el2());
+		write_el2_ctx_gcs(el2_sysregs_ctx, gcscr_el2, read_gcspr_el2());
+		write_el2_ctx_gcs(el2_sysregs_ctx, gcspr_el2, read_gcscr_el2());
 	}
 }
 
@@ -1352,8 +1357,8 @@
 	el2_sysregs_context_restore_common(el2_sysregs_ctx);
 	el2_sysregs_context_restore_gic(el2_sysregs_ctx);
 
-	if (is_feat_mte2_supported()) {
-		write_tfsr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_TFSR_EL2));
+	if (is_feat_mte_supported()) {
+		write_tfsr_el2(read_el2_ctx_mte(el2_sysregs_ctx, tfsr_el2));
 	}
 
 #if CTX_INCLUDE_MPAM_REGS
@@ -1368,50 +1373,57 @@
 	}
 
 	if (is_feat_ecv_v2_supported()) {
-		write_cntpoff_el2(read_ctx_reg(el2_sysregs_ctx, CTX_CNTPOFF_EL2));
+		write_cntpoff_el2(read_el2_ctx_ecv(el2_sysregs_ctx, cntpoff_el2));
 	}
 
 	if (is_feat_vhe_supported()) {
-		write_contextidr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_CONTEXTIDR_EL2));
-		write_ttbr1_el2(read_ctx_reg(el2_sysregs_ctx, CTX_TTBR1_EL2));
+		write_contextidr_el2(read_el2_ctx_vhe(el2_sysregs_ctx,
+					contextidr_el2));
+		write_ttbr1_el2(read_el2_ctx_vhe(el2_sysregs_ctx, ttbr1_el2));
 	}
 
 	if (is_feat_ras_supported()) {
-		write_vdisr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_VDISR_EL2));
-		write_vsesr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_VSESR_EL2));
+		write_vdisr_el2(read_el2_ctx_ras(el2_sysregs_ctx, vdisr_el2));
+		write_vsesr_el2(read_el2_ctx_ras(el2_sysregs_ctx, vsesr_el2));
 	}
 
 	if (is_feat_nv2_supported()) {
-		write_vncr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_VNCR_EL2));
+		write_vncr_el2(read_el2_ctx_neve(el2_sysregs_ctx, vncr_el2));
 	}
+
 	if (is_feat_trf_supported()) {
-		write_trfcr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_TRFCR_EL2));
+		write_trfcr_el2(read_el2_ctx_trf(el2_sysregs_ctx, trfcr_el2));
 	}
 
-	/* CSV2 version 2 and above */
 	if (is_feat_csv2_2_supported()) {
-		write_scxtnum_el2(read_ctx_reg(el2_sysregs_ctx, CTX_SCXTNUM_EL2));
+		write_scxtnum_el2(read_el2_ctx_csv2_2(el2_sysregs_ctx,
+					scxtnum_el2));
 	}
 
 	if (is_feat_hcx_supported()) {
-		write_hcrx_el2(read_ctx_reg(el2_sysregs_ctx, CTX_HCRX_EL2));
+		write_hcrx_el2(read_el2_ctx_hcx(el2_sysregs_ctx, hcrx_el2));
 	}
+
 	if (is_feat_tcr2_supported()) {
-		write_tcr2_el2(read_ctx_reg(el2_sysregs_ctx, CTX_TCR2_EL2));
+		write_tcr2_el2(read_el2_ctx_tcr2(el2_sysregs_ctx, tcr2_el2));
 	}
+
 	if (is_feat_sxpie_supported()) {
-		write_pire0_el2(read_ctx_reg(el2_sysregs_ctx, CTX_PIRE0_EL2));
-		write_pir_el2(read_ctx_reg(el2_sysregs_ctx, CTX_PIR_EL2));
-	}
-	if (is_feat_s2pie_supported()) {
-		write_s2pir_el2(read_ctx_reg(el2_sysregs_ctx, CTX_S2PIR_EL2));
+		write_pire0_el2(read_el2_ctx_sxpie(el2_sysregs_ctx, pire0_el2));
+		write_pir_el2(read_el2_ctx_sxpie(el2_sysregs_ctx, pir_el2));
 	}
+
 	if (is_feat_sxpoe_supported()) {
-		write_por_el2(read_ctx_reg(el2_sysregs_ctx, CTX_POR_EL2));
+		write_por_el2(read_el2_ctx_sxpoe(el2_sysregs_ctx, por_el2));
 	}
+
+	if (is_feat_s2pie_supported()) {
+		write_s2pir_el2(read_el2_ctx_s2pie(el2_sysregs_ctx, s2pir_el2));
+	}
+
 	if (is_feat_gcs_supported()) {
-		write_gcscr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_GCSCR_EL2));
-		write_gcspr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_GCSPR_EL2));
+		write_gcscr_el2(read_el2_ctx_gcs(el2_sysregs_ctx, gcscr_el2));
+		write_gcspr_el2(read_el2_ctx_gcs(el2_sysregs_ctx, gcspr_el2));
 	}
 }
 #endif /* CTX_INCLUDE_EL2_REGS */
diff --git a/lib/extensions/pmuv3/aarch64/pmuv3.c b/lib/extensions/pmuv3/aarch64/pmuv3.c
index 61fc47d..71aa303 100644
--- a/lib/extensions/pmuv3/aarch64/pmuv3.c
+++ b/lib/extensions/pmuv3/aarch64/pmuv3.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2023-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -24,11 +24,11 @@
 void pmuv3_enable(cpu_context_t *ctx)
 {
 #if CTX_INCLUDE_EL2_REGS
-	u_register_t mdcr_el2;
+	u_register_t mdcr_el2_val;
 
-	mdcr_el2 = read_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_MDCR_EL2);
-	mdcr_el2 = init_mdcr_el2_hpmn(mdcr_el2);
-	write_ctx_reg(get_el2_sysregs_ctx(ctx), CTX_MDCR_EL2, mdcr_el2);
+	mdcr_el2_val = read_el2_ctx_common(get_el2_sysregs_ctx(ctx), mdcr_el2);
+	mdcr_el2_val = init_mdcr_el2_hpmn(mdcr_el2_val);
+	write_el2_ctx_common(get_el2_sysregs_ctx(ctx), mdcr_el2, mdcr_el2_val);
 #endif /* CTX_INCLUDE_EL2_REGS */
 }
 
diff --git a/plat/arm/board/neoverse_rd/common/ras/nrd_ras_cpu.c b/plat/arm/board/neoverse_rd/common/ras/nrd_ras_cpu.c
index 7f1c376..371bde6 100644
--- a/plat/arm/board/neoverse_rd/common/ras/nrd_ras_cpu.c
+++ b/plat/arm/board/neoverse_rd/common/ras/nrd_ras_cpu.c
@@ -95,39 +95,39 @@
 						  CTX_TTBR1_EL1);
 
 #if CTX_INCLUDE_EL2_REGS
-	cpu_info->ErrCtxEl2Reg[0]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_ELR_EL2);
-	cpu_info->ErrCtxEl2Reg[1]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_ESR_EL2);
-	cpu_info->ErrCtxEl2Reg[2]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_FAR_EL2);
-	cpu_info->ErrCtxEl2Reg[3]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_HACR_EL2);
-	cpu_info->ErrCtxEl2Reg[4]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_HCR_EL2);
-	cpu_info->ErrCtxEl2Reg[5]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_HPFAR_EL2);
-	cpu_info->ErrCtxEl2Reg[6]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_MAIR_EL2);
-	cpu_info->ErrCtxEl2Reg[7]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_SCTLR_EL2);
-	cpu_info->ErrCtxEl2Reg[8]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_SP_EL2);
-	cpu_info->ErrCtxEl2Reg[9]   = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_SPSR_EL2);
-	cpu_info->ErrCtxEl2Reg[10]  = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_TCR_EL2);
-	cpu_info->ErrCtxEl2Reg[11]  = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_TPIDR_EL2);
-	cpu_info->ErrCtxEl2Reg[12]  = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_TTBR0_EL2);
-	cpu_info->ErrCtxEl2Reg[13]  = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_VTCR_EL2);
-	cpu_info->ErrCtxEl2Reg[14]  = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_VTTBR_EL2);
-	cpu_info->ErrCtxEl2Reg[15]  = read_ctx_reg(get_el2_sysregs_ctx(ctx),
-						   CTX_ESR_EL2);
-#endif
+	cpu_info->ErrCtxEl2Reg[0]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						elr_el2);
+	cpu_info->ErrCtxEl2Reg[1]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						esr_el2);
+	cpu_info->ErrCtxEl2Reg[2]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						far_el2);
+	cpu_info->ErrCtxEl2Reg[3]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						hacr_el2);
+	cpu_info->ErrCtxEl2Reg[4]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						hcr_el2);
+	cpu_info->ErrCtxEl2Reg[5]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						hpfar_el2);
+	cpu_info->ErrCtxEl2Reg[6]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						mair_el2);
+	cpu_info->ErrCtxEl2Reg[7]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						sctlr_el2);
+	cpu_info->ErrCtxEl2Reg[8]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						sp_el2);
+	cpu_info->ErrCtxEl2Reg[9]   = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						spsr_el2);
+	cpu_info->ErrCtxEl2Reg[10]  = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						tcr_el2);
+	cpu_info->ErrCtxEl2Reg[11]  = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						tpidr_el2);
+	cpu_info->ErrCtxEl2Reg[12]  = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						ttbr0_el2);
+	cpu_info->ErrCtxEl2Reg[13]  = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						vtcr_el2);
+	cpu_info->ErrCtxEl2Reg[14]  = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						vttbr_el2);
+	cpu_info->ErrCtxEl2Reg[15]  = read_el2_ctx_common(get_el2_sysregs_ctx(ctx),
+						esr_el2);
+#endif /* CTX_INCLUDE_EL2_REGS */
 
 	cpu_info->ErrCtxEl3Reg[0]   = read_ctx_reg(get_el3state_ctx(ctx),
 						   CTX_ELR_EL3);
diff --git a/services/std_svc/rmmd/rmmd_main.c b/services/std_svc/rmmd/rmmd_main.c
index fb096bc..b3b2a69 100644
--- a/services/std_svc/rmmd/rmmd_main.c
+++ b/services/std_svc/rmmd/rmmd_main.c
@@ -109,8 +109,8 @@
 
 static void rmm_el2_context_init(el2_sysregs_t *regs)
 {
-	regs->ctx_regs[CTX_SPSR_EL2 >> 3] = REALM_SPSR_EL2;
-	regs->ctx_regs[CTX_SCTLR_EL2 >> 3] = SCTLR_EL2_RES1;
+	write_el2_ctx_common(regs, spsr_el2, REALM_SPSR_EL2);
+	write_el2_ctx_common(regs, sctlr_el2, SCTLR_EL2_RES1);
 }
 
 /*******************************************************************************