diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index 512d196..d2d12e5 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2021, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2022, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -217,13 +217,13 @@
 
 // Starting with Armv8.4
 #define CTX_CONTEXTIDR_EL2	U(0x198)
-#define CTX_SDER32_EL2		U(0x1a0)
-#define CTX_TTBR1_EL2		U(0x1a8)
-#define CTX_VDISR_EL2		U(0x1b0)
-#define CTX_VNCR_EL2		U(0x1b8)
-#define CTX_VSESR_EL2		U(0x1c0)
-#define CTX_VSTCR_EL2		U(0x1c8)
-#define CTX_VSTTBR_EL2		U(0x1d0)
+#define CTX_TTBR1_EL2		U(0x1a0)
+#define CTX_VDISR_EL2		U(0x1a8)
+#define CTX_VSESR_EL2		U(0x1b0)
+#define CTX_VSTCR_EL2		U(0x1b8)
+#define CTX_VSTTBR_EL2		U(0x1c0)
+#define CTX_SDER32_EL2		U(0x1c8)
+#define CTX_VNCR_EL2		U(0x1d0)
 #define CTX_TRFCR_EL2		U(0x1d8)
 
 // Starting with Armv8.5
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index c1c0612..0432d96 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -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
  */
@@ -13,14 +13,14 @@
 #if CTX_INCLUDE_EL2_REGS
 	.global	el2_sysregs_context_save
 	.global	el2_sysregs_context_restore
-#endif
+#endif /* CTX_INCLUDE_EL2_REGS */
 
 	.global	el1_sysregs_context_save
 	.global	el1_sysregs_context_restore
 #if CTX_INCLUDE_FPREGS
 	.global	fpregs_context_save
 	.global	fpregs_context_restore
-#endif
+#endif /* CTX_INCLUDE_FPREGS */
 	.global	prepare_el3_entry
 	.global	restore_gp_pmcr_pauth_regs
 	.global save_and_update_ptw_el1_sys_regs
@@ -62,7 +62,7 @@
 #if CTX_INCLUDE_AARCH32_REGS
 	mrs	x16, dbgvcr32_el2
 	str	x16, [x0, #CTX_DBGVCR32_EL2]
-#endif
+#endif /* CTX_INCLUDE_AARCH32_REGS */
 
 	mrs	x9, elr_el2
 	mrs	x10, esr_el2
@@ -91,7 +91,8 @@
 #if ENABLE_SPE_FOR_LOWER_ELS
 	mrs	x13, PMSCR_EL2
 	str	x13, [x0, #CTX_PMSCR_EL2]
-#endif
+#endif /* ENABLE_SPE_FOR_LOWER_ELS */
+
 	mrs	x14, sctlr_el2
 	str	x14, [x0, #CTX_SCTLR_EL2]
 
@@ -118,7 +119,7 @@
 #if CTX_INCLUDE_MTE_REGS
 	mrs	x9, TFSR_EL2
 	str	x9, [x0, #CTX_TFSR_EL2]
-#endif
+#endif /* CTX_INCLUDE_MTE_REGS */
 
 #if ENABLE_MPAM_FOR_LOWER_ELS
 	mrs	x10, MPAM2_EL2
@@ -143,7 +144,7 @@
 	mrs	x11, MPAMVPM7_EL2
 	mrs	x12, MPAMVPMV_EL2
 	stp	x11, x12, [x0, #CTX_MPAMVPM7_EL2]
-#endif
+#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
 
 #if ENABLE_FEAT_FGT
 	mrs	x13, HDFGRTR_EL2
@@ -152,7 +153,7 @@
    	stp	x13, x14, [x0, #CTX_HDFGRTR_EL2]
 #else
    	str	x13, [x0, #CTX_HDFGRTR_EL2]
-#endif
+#endif /* ENABLE_FEAT_AMUv1 */
 	mrs	x15, HDFGWTR_EL2
 	mrs	x16, HFGITR_EL2
 	stp	x15, x16, [x0, #CTX_HDFGWTR_EL2]
@@ -160,48 +161,80 @@
 	mrs	x9, HFGRTR_EL2
 	mrs	x10, HFGWTR_EL2
 	stp	x9, x10, [x0, #CTX_HFGRTR_EL2]
-#endif
+#endif /* ENABLE_FEAT_FGT */
 
 #if ENABLE_FEAT_ECV
 	mrs	x11, CNTPOFF_EL2
 	str	x11, [x0, #CTX_CNTPOFF_EL2]
-#endif
+#endif /* ENABLE_FEAT_ECV */
 
-#if ARM_ARCH_AT_LEAST(8, 4)
-	mrs	x12, contextidr_el2
-	str	x12, [x0, #CTX_CONTEXTIDR_EL2]
+#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 CTX_INCLUDE_AARCH32_REGS
-	mrs	x13, sder32_el2
-	str	x13, [x0, #CTX_SDER32_EL2]
-#endif
-	mrs	x14, ttbr1_el2
-	mrs	x15, vdisr_el2
-	stp	x14, x15, [x0, #CTX_TTBR1_EL2]
+#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 ENABLE_FEAT_SEL2
+	/*
+	 * VSTCR_EL2 and VSTTBR_EL2 registers are saved only
+	 * when FEAT_SEL2 is supported.
+	 */
+	mrs	x13, vstcr_el2
+	mrs	x14, vsttbr_el2
+	stp	x13, x14, [x0, #CTX_VSTCR_EL2]
+#endif /* ENABLE_FEAT_SEL2 */
+
+#if CTX_INCLUDE_AARCH32_REGS && ENABLE_FEAT_SEL2
+	/*
+	 * SDER32_EL2 register is saved only when EL2 and EL1
+	 * capable of using Aarch32 and FEAT_SEL2 is supported.
+	 */
+	mrs	x15, sder32_el2
+	str	x15, [x0, #CTX_SDER32_EL2]
+#endif /* CTX_INCLUDE_AARCH32_REGS && ENABLE_FEAT_SEL2 */
 
 #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
+#endif /* CTX_INCLUDE_NEVE_REGS */
 
-	mrs	x9, vsesr_el2
-	mrs	x10, vstcr_el2
-	stp	x9, x10, [x0, #CTX_VSESR_EL2]
-
-	mrs	x11, vsttbr_el2
+#if ENABLE_TRF_FOR_NS
+	/*
+	 * TRFCR_EL2 register is saved only when FEAT_TRF is supported.
+	 */
 	mrs	x12, TRFCR_EL2
-	stp	x11, x12, [x0, #CTX_VSTTBR_EL2]
-#endif
+	str	x12, [x0, #CTX_TRFCR_EL2]
+#endif /* ENABLE_TRF_FOR_NS */
 
-#if ARM_ARCH_AT_LEAST(8, 5)
+#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
+#endif /* ENABLE_FEAT_CSV2_2 */
 
 #if ENABLE_FEAT_HCX
 	mrs	x14, hcrx_el2
 	str	x14, [x0, #CTX_HCRX_EL2]
-#endif
+#endif /* ENABLE_FEAT_HCX */
 
 	ret
 endfunc el2_sysregs_context_save
@@ -241,7 +274,7 @@
 #if CTX_INCLUDE_AARCH32_REGS
 	ldr	x16, [x0, #CTX_DBGVCR32_EL2]
 	msr	dbgvcr32_el2, x16
-#endif
+#endif /* CTX_INCLUDE_AARCH32_REGS */
 
 	ldp	x9, x10, [x0, #CTX_ELR_EL2]
 	msr	elr_el2, x9
@@ -270,7 +303,8 @@
 #if ENABLE_SPE_FOR_LOWER_ELS
 	ldr	x13, [x0, #CTX_PMSCR_EL2]
 	msr	PMSCR_EL2, x13
-#endif
+#endif /* ENABLE_SPE_FOR_LOWER_ELS */
+
 	ldr	x14, [x0, #CTX_SCTLR_EL2]
 	msr	sctlr_el2, x14
 
@@ -297,7 +331,7 @@
 #if CTX_INCLUDE_MTE_REGS
 	ldr	x9, [x0, #CTX_TFSR_EL2]
 	msr	TFSR_EL2, x9
-#endif
+#endif /* CTX_INCLUDE_MTE_REGS */
 
 #if ENABLE_MPAM_FOR_LOWER_ELS
 	ldr	x10, [x0, #CTX_MPAM2_EL2]
@@ -322,7 +356,7 @@
 	ldp	x11, x12, [x0, #CTX_MPAMVPM7_EL2]
 	msr	MPAMVPM7_EL2, x11
 	msr	MPAMVPMV_EL2, x12
-#endif
+#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
 
 #if ENABLE_FEAT_FGT
 #if ENABLE_FEAT_AMUv1
@@ -330,7 +364,7 @@
 	msr	HAFGRTR_EL2, x14
 #else
 	ldr	x13, [x0, #CTX_HDFGRTR_EL2]
-#endif
+#endif /* ENABLE_FEAT_AMUv1 */
 	msr	HDFGRTR_EL2, x13
 
 	ldp	x15, x16, [x0, #CTX_HDFGWTR_EL2]
@@ -340,48 +374,80 @@
 	ldp	x9, x10, [x0, #CTX_HFGRTR_EL2]
 	msr	HFGRTR_EL2, x9
 	msr	HFGWTR_EL2, x10
-#endif
+#endif /* ENABLE_FEAT_FGT */
 
 #if ENABLE_FEAT_ECV
 	ldr	x11, [x0, #CTX_CNTPOFF_EL2]
 	msr	CNTPOFF_EL2, x11
-#endif
+#endif /* ENABLE_FEAT_ECV */
+
+#if ENABLE_FEAT_VHE
+	/*
+	 * CONTEXTIDR_EL2 register is restored only when FEAT_VHE or
+	 * FEAT_Debugv8p2 (currently not in TF-A) is supported.
+	 */
+	ldp	x9, x10, [x0, #CTX_CONTEXTIDR_EL2]
+	msr	contextidr_el2, x9
+	msr	ttbr1_el2, x10
+#endif /* ENABLE_FEAT_VHE */
 
-#if ARM_ARCH_AT_LEAST(8, 4)
-	ldr	x12, [x0, #CTX_CONTEXTIDR_EL2]
-	msr	contextidr_el2, x12
+#if RAS_EXTENSION
+	/*
+	 * VDISR_EL2 and VSESR_EL2 registers are restored only when FEAT_RAS
+	 * is supported.
+	 */
+	ldp	x11, x12, [x0, #CTX_VDISR_EL2]
+	msr	vdisr_el2, x11
+	msr	vsesr_el2, x12
+#endif /* RAS_EXTENSION */
 
-#if CTX_INCLUDE_AARCH32_REGS
-	ldr	x13, [x0, #CTX_SDER32_EL2]
-	msr	sder32_el2, x13
-#endif
-	ldp	x14, x15, [x0, #CTX_TTBR1_EL2]
-	msr	ttbr1_el2, x14
-	msr	vdisr_el2, x15
+#if ENABLE_FEAT_SEL2
+	/*
+	 * VSTCR_EL2 and VSTTBR_EL2 registers are restored only when FEAT_SEL2
+	 * is supported.
+	 */
+	ldp	x13, x14, [x0, #CTX_VSTCR_EL2]
+	msr	vstcr_el2, x13
+	msr	vsttbr_el2, x14
+#endif /* ENABLE_FEAT_SEL2 */
+
+#if CTX_INCLUDE_AARCH32_REGS && ENABLE_FEAT_SEL2
+	/*
+	 * SDER32_EL2 register is restored only when EL2 and EL1 capable of using
+	 * Aarch32 and FEAT_SEL2 is supported.
+	 */
+	ldr	x15, [x0, #CTX_SDER32_EL2]
+	msr	sder32_el2, x15
+#endif /* CTX_INCLUDE_AARCH32_REGS && ENABLE_FEAT_SEL2 */
 
 #if CTX_INCLUDE_NEVE_REGS
+	/*
+	 * VNCR_EL2 register is restored only when FEAT_NV2 is supported.
+	 */
 	ldr	x16, [x0, #CTX_VNCR_EL2]
 	msr	vncr_el2, x16
-#endif
-
-	ldp	x9, x10, [x0, #CTX_VSESR_EL2]
-	msr	vsesr_el2, x9
-	msr	vstcr_el2, x10
+#endif /* CTX_INCLUDE_NEVE_REGS */
 
-	ldp	x11, x12, [x0, #CTX_VSTTBR_EL2]
-	msr	vsttbr_el2, x11
+#if ENABLE_TRF_FOR_NS
+	/*
+	 * TRFCR_EL2 register is restored only when FEAT_TRF is supported.
+	 */
+	ldr	x12, [x0, #CTX_TRFCR_EL2]
 	msr	TRFCR_EL2, x12
-#endif
+#endif /* ENABLE_TRF_FOR_NS */
 
-#if ARM_ARCH_AT_LEAST(8, 5)
+#if ENABLE_FEAT_CSV2_2
+	/*
+	 * SCXTNUM_EL2 register is restored only when FEAT_CSV2_2 is supported.
+	 */
 	ldr	x13, [x0, #CTX_SCXTNUM_EL2]
 	msr	scxtnum_el2, x13
-#endif
+#endif /* ENABLE_FEAT_CSV2_2 */
 
 #if ENABLE_FEAT_HCX
 	ldr	x14, [x0, #CTX_HCRX_EL2]
 	msr	hcrx_el2, x14
-#endif
+#endif /* ENABLE_FEAT_HCX */
 
 	ret
 endfunc el2_sysregs_context_restore
@@ -405,7 +471,7 @@
 	mrs	x15, sctlr_el1
 	mrs	x16, tcr_el1
 	stp	x15, x16, [x0, #CTX_SCTLR_EL1]
-#endif
+#endif /* ERRATA_SPECULATIVE_AT */
 
 	mrs	x17, cpacr_el1
 	mrs	x9, csselr_el1
@@ -456,7 +522,7 @@
 	mrs	x15, dacr32_el2
 	mrs	x16, ifsr32_el2
 	stp	x15, x16, [x0, #CTX_DACR32_EL2]
-#endif
+#endif /* CTX_INCLUDE_AARCH32_REGS */
 
 	/* Save NS timer registers if the build has instructed so */
 #if NS_TIMER_SWITCH
@@ -470,7 +536,7 @@
 
 	mrs	x14, cntkctl_el1
 	str	x14, [x0, #CTX_CNTKCTL_EL1]
-#endif
+#endif /* NS_TIMER_SWITCH */
 
 	/* Save MTE system registers if the build has instructed so */
 #if CTX_INCLUDE_MTE_REGS
@@ -481,7 +547,7 @@
 	mrs	x9, RGSR_EL1
 	mrs	x10, GCR_EL1
 	stp	x9, x10, [x0, #CTX_RGSR_EL1]
-#endif
+#endif /* CTX_INCLUDE_MTE_REGS */
 
 	ret
 endfunc el1_sysregs_context_save
@@ -504,7 +570,7 @@
 	ldp	x15, x16, [x0, #CTX_SCTLR_EL1]
 	msr	sctlr_el1, x15
 	msr	tcr_el1, x16
-#endif
+#endif /* ERRATA_SPECULATIVE_AT */
 
 	ldp	x17, x9, [x0, #CTX_CPACR_EL1]
 	msr	cpacr_el1, x17
@@ -555,7 +621,8 @@
 	ldp	x15, x16, [x0, #CTX_DACR32_EL2]
 	msr	dacr32_el2, x15
 	msr	ifsr32_el2, x16
-#endif
+#endif /* CTX_INCLUDE_AARCH32_REGS */
+
 	/* Restore NS timer registers if the build has instructed so */
 #if NS_TIMER_SWITCH
 	ldp	x10, x11, [x0, #CTX_CNTP_CTL_EL0]
@@ -568,7 +635,8 @@
 
 	ldr	x14, [x0, #CTX_CNTKCTL_EL1]
 	msr	cntkctl_el1, x14
-#endif
+#endif /* NS_TIMER_SWITCH */
+
 	/* Restore MTE system registers if the build has instructed so */
 #if CTX_INCLUDE_MTE_REGS
 	ldp	x11, x12, [x0, #CTX_TFSRE0_EL1]
@@ -578,7 +646,7 @@
 	ldp	x13, x14, [x0, #CTX_RGSR_EL1]
 	msr	RGSR_EL1, x13
 	msr	GCR_EL1, x14
-#endif
+#endif /* CTX_INCLUDE_MTE_REGS */
 
 	/* No explict ISB required here as ERET covers it */
 	ret
@@ -626,7 +694,7 @@
 #if CTX_INCLUDE_AARCH32_REGS
 	mrs	x11, fpexc32_el2
 	str	x11, [x0, #CTX_FP_FPEXC32_EL2]
-#endif
+#endif /* CTX_INCLUDE_AARCH32_REGS */
 	ret
 endfunc fpregs_context_save
 
@@ -671,7 +739,8 @@
 #if CTX_INCLUDE_AARCH32_REGS
 	ldr	x11, [x0, #CTX_FP_FPEXC32_EL2]
 	msr	fpexc32_el2, x11
-#endif
+#endif /* CTX_INCLUDE_AARCH32_REGS */
+
 	/*
 	 * No explict ISB required here as ERET to
 	 * switch to secure EL1 or non-secure world
@@ -688,13 +757,13 @@
 	 * in ARM DDI 0487F.c page J1-7635 to a default value.
 	 */
 	.macro set_unset_pstate_bits
-        /*
-         * If Data Independent Timing (DIT) functionality is implemented,
-         * always enable DIT in EL3
-         */
+	/*
+	 * If Data Independent Timing (DIT) functionality is implemented,
+	 * always enable DIT in EL3
+	 */
 #if ENABLE_FEAT_DIT
-        mov     x8, #DIT_BIT
-        msr     DIT, x8
+	mov     x8, #DIT_BIT
+	msr     DIT, x8
 #endif /* ENABLE_FEAT_DIT */
 	.endm /* set_unset_pstate_bits */
 
@@ -933,7 +1002,7 @@
 	mrs	x17, spsel
 	cmp	x17, #MODE_SP_EL0
 	ASM_ASSERT(eq)
-#endif
+#endif /* ENABLE_ASSERTIONS */
 
 	/* ----------------------------------------------------------
 	 * Save the current SP_EL0 i.e. the EL3 runtime stack which
@@ -971,7 +1040,7 @@
 	isb
 	msr	S3_6_C1_C2_0, x20 /* zcr_el3 */
 sve_not_enabled:
-#endif
+#endif /* IMAGE_BL31 */
 
 #if IMAGE_BL31 && DYNAMIC_WORKAROUND_CVE_2018_3639
 	/* ----------------------------------------------------------
@@ -982,7 +1051,8 @@
 	cbz	x17, 1f
 	blr	x17
 1:
-#endif
+#endif /* IMAGE_BL31 && DYNAMIC_WORKAROUND_CVE_2018_3639 */
+
 	restore_ptw_el1_sys_regs
 
 	/* ----------------------------------------------------------
@@ -1005,10 +1075,12 @@
 	esb
 #else
 	dsb	sy
-#endif
+#endif /* IMAGE_BL31 && RAS_EXTENSION */
+
 #ifdef IMAGE_BL31
 	str	xzr, [sp, #CTX_EL3STATE_OFFSET + CTX_IS_IN_EL3]
-#endif
+#endif /* IMAGE_BL31 */
+
 	exception_return
 
 endfunc el3_exit