diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 37bb12c..1101425 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -14,61 +14,16 @@
 	.global	fpregs_context_save
 	.global	fpregs_context_restore
 #endif
-#if CTX_INCLUDE_PAUTH_REGS
-	.global	pauth_context_restore
-	.global	pauth_context_save
-#endif
-#if ENABLE_PAUTH
-	.global	pauth_load_bl_apiakey
-#endif
-	.global	save_gp_registers
-	.global	restore_gp_registers
-	.global	restore_gp_registers_eret
-	.global	save_pmcr_disable_pmu
+	.global	save_gp_pmcr_pauth_regs
+	.global	restore_gp_pmcr_pauth_regs
 	.global	el3_exit
 
-/* -----------------------------------------------------
- * If ARMv8.5-PMU is implemented, cycle counting is
- * disabled by seting MDCR_EL3.SCCD to 1.
- * -----------------------------------------------------
- */
-func save_pmcr_disable_pmu
-	/* -----------------------------------------------------
-	 * Check if earlier initialization MDCR_EL3.SCCD to 1
-	 * failed, meaning that ARMv8-PMU is not implemented and
-	 * PMCR_EL0 should be saved in non-secure context.
-	 * -----------------------------------------------------
-	 */
-	mrs	x9, mdcr_el3
-	tst	x9, #MDCR_SCCD_BIT
-	bne	1f
-
-	/* Secure Cycle Counter is not disabled */
-	mrs	x9, pmcr_el0
-
-	/* Check caller's security state */
-	mrs	x10, scr_el3
-	tst	x10, #SCR_NS_BIT
-	beq	2f
-
-	/* Save PMCR_EL0 if called from Non-secure state */
-	str	x9, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
-
-	/* Disable cycle counter when event counting is prohibited */
-2:	orr	x9, x9, #PMCR_EL0_DP_BIT
-	msr	pmcr_el0, x9
-
-	isb
-1:	ret
-endfunc save_pmcr_disable_pmu
-
-/* -----------------------------------------------------
- * 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
- * the register context will be saved.
- * -----------------------------------------------------
+/* ------------------------------------------------------------------
+ * 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 the register context will be saved.
+ * ------------------------------------------------------------------
  */
 func el1_sysregs_context_save
 
@@ -159,13 +114,13 @@
 	ret
 endfunc el1_sysregs_context_save
 
-/* -----------------------------------------------------
- * 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
- * from where the register context will be restored
- * -----------------------------------------------------
+/* ------------------------------------------------------------------
+ * 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 from where the register context will be
+ * restored
+ * ------------------------------------------------------------------
  */
 func el1_sysregs_context_restore
 
@@ -255,21 +210,19 @@
 	ret
 endfunc el1_sysregs_context_restore
 
-/* -----------------------------------------------------
- * The following function follows the aapcs_64 strictly
- * to use x9-x17 (temporary caller-saved registers
- * according to AArch64 PCS) to save floating point
- * register context. It assumes that 'x0' is pointing to
- * a 'fp_regs' structure where the register context will
+/* ------------------------------------------------------------------
+ * The following function follows the aapcs_64 strictly to use
+ * x9-x17 (temporary caller-saved registers according to AArch64 PCS)
+ * to save floating point register context. It assumes that 'x0' is
+ * pointing to a 'fp_regs' structure where the register context will
  * be saved.
  *
- * Access to VFP registers will trap if CPTR_EL3.TFP is
- * set.  However currently we don't use VFP registers
- * nor set traps in Trusted Firmware, and assume it's
- * cleared
+ * Access to VFP registers will trap if CPTR_EL3.TFP is set.
+ * However currently we don't use VFP registers nor set traps in
+ * Trusted Firmware, and assume it's cleared.
  *
  * TODO: Revisit when VFP is used in secure world
- * -----------------------------------------------------
+ * ------------------------------------------------------------------
  */
 #if CTX_INCLUDE_FPREGS
 func fpregs_context_save
@@ -303,21 +256,19 @@
 	ret
 endfunc fpregs_context_save
 
-/* -----------------------------------------------------
- * The following function follows the aapcs_64 strictly
- * to use x9-x17 (temporary caller-saved registers
- * according to AArch64 PCS) to restore floating point
- * register context. It assumes that 'x0' is pointing to
- * a 'fp_regs' structure from where the register context
+/* ------------------------------------------------------------------
+ * The following function follows the aapcs_64 strictly to use x9-x17
+ * (temporary caller-saved registers according to AArch64 PCS) to
+ * restore floating point register context. It assumes that 'x0' is
+ * pointing to a 'fp_regs' structure from where the register context
  * will be restored.
  *
- * Access to VFP registers will trap if CPTR_EL3.TFP is
- * set.  However currently we don't use VFP registers
- * nor set traps in Trusted Firmware, and assume it's
- * cleared
+ * Access to VFP registers will trap if CPTR_EL3.TFP is set.
+ * However currently we don't use VFP registers nor set traps in
+ * Trusted Firmware, and assume it's cleared.
  *
  * TODO: Revisit when VFP is used in secure world
- * -----------------------------------------------------
+ * ------------------------------------------------------------------
  */
 func fpregs_context_restore
 	ldp	q0, q1, [x0, #CTX_FP_Q0]
@@ -357,109 +308,23 @@
 endfunc fpregs_context_restore
 #endif /* CTX_INCLUDE_FPREGS */
 
-#if CTX_INCLUDE_PAUTH_REGS
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to save the ARMv8.3-PAuth register context. It assumes
- * that 'sp' is pointing to a 'cpu_context_t' structure
- * to where the register context will be saved.
- * -----------------------------------------------------
- */
-func pauth_context_save
-	add	x11, sp, #CTX_PAUTH_REGS_OFFSET
-
-	mrs	x9, APIAKeyLo_EL1
-	mrs	x10, APIAKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACIAKEY_LO]
-
-	mrs	x9, APIBKeyLo_EL1
-	mrs	x10, APIBKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACIBKEY_LO]
-
-	mrs	x9, APDAKeyLo_EL1
-	mrs	x10, APDAKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACDAKEY_LO]
-
-	mrs	x9, APDBKeyLo_EL1
-	mrs	x10, APDBKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACDBKEY_LO]
-
-	mrs	x9, APGAKeyLo_EL1
-	mrs	x10, APGAKeyHi_EL1
-	stp	x9, x10, [x11, #CTX_PACGAKEY_LO]
-
-	ret
-endfunc pauth_context_save
-
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to restore the ARMv8.3-PAuth register context. It assumes
- * that 'sp' is pointing to a 'cpu_context_t' structure
- * from where the register context will be restored.
- * -----------------------------------------------------
- */
-func pauth_context_restore
-	add	x11, sp, #CTX_PAUTH_REGS_OFFSET
-
-	ldp	x9, x10, [x11, #CTX_PACIAKEY_LO]
-	msr	APIAKeyLo_EL1, x9
-	msr	APIAKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACIBKEY_LO]
-	msr	APIBKeyLo_EL1, x9
-	msr	APIBKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACDAKEY_LO]
-	msr	APDAKeyLo_EL1, x9
-	msr	APDAKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACDBKEY_LO]
-	msr	APDBKeyLo_EL1, x9
-	msr	APDBKeyHi_EL1, x10
-
-	ldp	x9, x10, [x11, #CTX_PACGAKEY_LO]
-	msr	APGAKeyLo_EL1, x9
-	msr	APGAKeyHi_EL1, x10
-
-	ret
-endfunc pauth_context_restore
-#endif /* CTX_INCLUDE_PAUTH_REGS */
-
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
- * to load the APIA key used by the firmware.
- * -----------------------------------------------------
- */
-#if ENABLE_PAUTH
-func pauth_load_bl_apiakey
-	/* Load instruction key A used by the Trusted Firmware. */
-	adrp	x11, plat_apiakey
-	add	x11, x11, :lo12:plat_apiakey
-	ldp	x9, x10, [x11, #0]
-
-	msr	APIAKeyLo_EL1, x9
-	msr	APIAKeyHi_EL1, x10
-
-	ret
-endfunc pauth_load_bl_apiakey
-#endif /* ENABLE_PAUTH */
-
-/* -----------------------------------------------------
- * The following functions are used to save and restore
- * all the general purpose registers. Ideally we would
- * only save and restore the callee saved registers when
- * a world switch occurs but that type of implementation
- * is more complex. So currently we will always save and
- * restore these registers on entry and exit of EL3.
- * These are not macros to ensure their invocation fits
- * within the 32 instructions per exception vector.
+/* ------------------------------------------------------------------
+ * The following function is used to save and restore all the general
+ * purpose and ARMv8.3-PAuth (if enabled) registers.
+ * It also checks if Secure Cycle Counter is not disabled in MDCR_EL3
+ * when ARMv8.5-PMU is implemented, and if called from Non-secure
+ * state saves PMCR_EL0 and disables Cycle Counter.
+ *
+ * Ideally we would only save and restore the callee saved registers
+ * when a world switch occurs but that type of implementation is more
+ * complex. So currently we will always save and restore these
+ * registers on entry and exit of EL3.
+ * These are not macros to ensure their invocation fits within the 32
+ * instructions per exception vector.
  * clobbers: x18
- * -----------------------------------------------------
+ * ------------------------------------------------------------------
  */
-func save_gp_registers
+func save_gp_pmcr_pauth_regs
 	stp	x0, x1, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X0]
 	stp	x2, x3, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X2]
 	stp	x4, x5, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X4]
@@ -477,15 +342,114 @@
 	stp	x28, x29, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X28]
 	mrs	x18, sp_el0
 	str	x18, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_SP_EL0]
-	ret
-endfunc save_gp_registers
 
-/* -----------------------------------------------------
- * This function restores all general purpose registers except x30 from the
- * CPU context. x30 register must be explicitly restored by the caller.
- * -----------------------------------------------------
+	/* ----------------------------------------------------------
+	 * Check if earlier initialization MDCR_EL3.SCCD to 1 failed,
+	 * meaning that ARMv8-PMU is not implemented and PMCR_EL0
+	 * should be saved in non-secure context.
+	 * ----------------------------------------------------------
+	 */
+	mrs	x9, mdcr_el3
+	tst	x9, #MDCR_SCCD_BIT
+	bne	1f
+
+	/* Secure Cycle Counter is not disabled */
+	mrs	x9, pmcr_el0
+
+	/* Check caller's security state */
+	mrs	x10, scr_el3
+	tst	x10, #SCR_NS_BIT
+	beq	2f
+
+	/* Save PMCR_EL0 if called from Non-secure state */
+	str	x9, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
+
+	/* Disable cycle counter when event counting is prohibited */
+2:	orr	x9, x9, #PMCR_EL0_DP_BIT
+	msr	pmcr_el0, x9
+	isb
+1:
+#if CTX_INCLUDE_PAUTH_REGS
+	/* ----------------------------------------------------------
+ 	 * Save the ARMv8.3-PAuth keys as they are not banked
+ 	 * by exception level
+	 * ----------------------------------------------------------
+	 */
+	add	x19, sp, #CTX_PAUTH_REGS_OFFSET
+
+	mrs	x20, APIAKeyLo_EL1	/* x21:x20 = APIAKey */
+	mrs	x21, APIAKeyHi_EL1
+	mrs	x22, APIBKeyLo_EL1	/* x23:x22 = APIBKey */
+	mrs	x23, APIBKeyHi_EL1
+	mrs	x24, APDAKeyLo_EL1	/* x25:x24 = APDAKey */
+	mrs	x25, APDAKeyHi_EL1
+	mrs	x26, APDBKeyLo_EL1	/* x27:x26 = APDBKey */
+	mrs	x27, APDBKeyHi_EL1
+	mrs	x28, APGAKeyLo_EL1	/* x29:x28 = APGAKey */
+	mrs	x29, APGAKeyHi_EL1
+
+	stp	x20, x21, [x19, #CTX_PACIAKEY_LO]
+	stp	x22, x23, [x19, #CTX_PACIBKEY_LO]
+	stp	x24, x25, [x19, #CTX_PACDAKEY_LO]
+	stp	x26, x27, [x19, #CTX_PACDBKEY_LO]
+	stp	x28, x29, [x19, #CTX_PACGAKEY_LO]
+#endif /* CTX_INCLUDE_PAUTH_REGS */
+
+	ret
+endfunc save_gp_pmcr_pauth_regs
+
+/* ------------------------------------------------------------------
+ * This function restores ARMv8.3-PAuth (if enabled) and all general
+ * purpose registers except x30 from the CPU context.
+ * x30 register must be explicitly restored by the caller.
+ * ------------------------------------------------------------------
  */
-func restore_gp_registers
+func restore_gp_pmcr_pauth_regs
+#if CTX_INCLUDE_PAUTH_REGS
+ 	/* Restore the ARMv8.3 PAuth keys */
+	add	x10, sp, #CTX_PAUTH_REGS_OFFSET
+
+	ldp	x0, x1, [x10, #CTX_PACIAKEY_LO]	/* x1:x0 = APIAKey */
+	ldp	x2, x3, [x10, #CTX_PACIBKEY_LO]	/* x3:x2 = APIBKey */
+	ldp	x4, x5, [x10, #CTX_PACDAKEY_LO]	/* x5:x4 = APDAKey */
+	ldp	x6, x7, [x10, #CTX_PACDBKEY_LO]	/* x7:x6 = APDBKey */
+	ldp	x8, x9, [x10, #CTX_PACGAKEY_LO]	/* x9:x8 = APGAKey */
+
+	msr	APIAKeyLo_EL1, x0
+	msr	APIAKeyHi_EL1, x1
+	msr	APIBKeyLo_EL1, x2
+	msr	APIBKeyHi_EL1, x3
+	msr	APDAKeyLo_EL1, x4
+	msr	APDAKeyHi_EL1, x5
+	msr	APDBKeyLo_EL1, x6
+	msr	APDBKeyHi_EL1, x7
+	msr	APGAKeyLo_EL1, x8
+	msr	APGAKeyHi_EL1, x9
+#endif /* CTX_INCLUDE_PAUTH_REGS */
+
+	/* ----------------------------------------------------------
+	 * Restore PMCR_EL0 when returning to Non-secure state if
+	 * Secure Cycle Counter is not disabled in MDCR_EL3 when
+	 * ARMv8.5-PMU is implemented.
+	 * ----------------------------------------------------------
+	 */
+	mrs	x0, scr_el3
+	tst	x0, #SCR_NS_BIT
+	beq	2f
+
+	/* ----------------------------------------------------------
+	 * Back to Non-secure state.
+	 * Check if earlier initialization MDCR_EL3.SCCD to 1 failed,
+	 * meaning that ARMv8-PMU is not implemented and PMCR_EL0
+	 * should be restored from non-secure context.
+	 * ----------------------------------------------------------
+	 */
+	mrs	x0, mdcr_el3
+	tst	x0, #MDCR_SCCD_BIT
+	bne	2f
+	ldr	x0, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
+	msr	pmcr_el0, x0
+2:
 	ldp	x0, x1, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X0]
 	ldp	x2, x3, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X2]
 	ldp	x4, x5, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X4]
@@ -504,49 +468,28 @@
 	msr	sp_el0, x28
 	ldp	x28, x29, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_X28]
 	ret
-endfunc restore_gp_registers
+endfunc restore_gp_pmcr_pauth_regs
 
-/* -----------------------------------------------------
- * Restore general purpose registers (including x30), and exit EL3 via ERET to
- * a lower exception level.
- * -----------------------------------------------------
- */
-func restore_gp_registers_eret
-	bl	restore_gp_registers
-	ldr	x30, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_LR]
-
-#if IMAGE_BL31 && RAS_EXTENSION
-	/*
-	 * Issue Error Synchronization Barrier to synchronize SErrors before
-	 * exiting EL3. We're running with EAs unmasked, so any synchronized
-	 * errors would be taken immediately; therefore no need to inspect
-	 * DISR_EL1 register.
-	 */
-	esb
-#endif
-	eret
-endfunc	restore_gp_registers_eret
-
-/* -----------------------------------------------------
- * This routine assumes that the SP_EL3 is pointing to
- * a valid context structure from where the gp regs and
- * other special registers can be retrieved.
- * -----------------------------------------------------
+/* ------------------------------------------------------------------
+ * This routine assumes that the SP_EL3 is pointing to a valid
+ * context structure from where the gp regs and other special
+ * registers can be retrieved.
+ * ------------------------------------------------------------------
  */
 func el3_exit
-	/* -----------------------------------------------------
-	 * Save the current SP_EL0 i.e. the EL3 runtime stack
-	 * which will be used for handling the next SMC. Then
-	 * switch to SP_EL3
-	 * -----------------------------------------------------
+	/* ----------------------------------------------------------
+	 * Save the current SP_EL0 i.e. the EL3 runtime stack which
+	 * will be used for handling the next SMC.
+	 * Then switch to SP_EL3.
+	 * ----------------------------------------------------------
 	 */
 	mov	x17, sp
-	msr	spsel, #1
+	msr	spsel, #MODE_SP_ELX
 	str	x17, [sp, #CTX_EL3STATE_OFFSET + CTX_RUNTIME_SP]
 
-	/* -----------------------------------------------------
+	/* ----------------------------------------------------------
 	 * Restore SPSR_EL3, ELR_EL3 and SCR_EL3 prior to ERET
-	 * -----------------------------------------------------
+	 * ----------------------------------------------------------
 	 */
 	ldr	x18, [sp, #CTX_EL3STATE_OFFSET + CTX_SCR_EL3]
 	ldp	x16, x17, [sp, #CTX_EL3STATE_OFFSET + CTX_SPSR_EL3]
@@ -554,43 +497,35 @@
 	msr	spsr_el3, x16
 	msr	elr_el3, x17
 
-	/* -----------------------------------------------------
-	 * Restore PMCR_EL0 when returning to Non-secure state
-	 * if Secure Cycle Counter is not disabled in MDCR_EL3
-	 * when ARMv8.5-PMU is implemented
-	 * -----------------------------------------------------
-	 */
-	tst	x18, #SCR_NS_BIT
-	beq	2f
-
-	/* -----------------------------------------------------
-	 * Back to Non-secure state.
-	 * Check if earlier initialization MDCR_EL3.SCCD to 1
-	 * failed, meaning that ARMv8-PMU is not implemented and
-	 * PMCR_EL0 should be restored from non-secure context.
-	 * -----------------------------------------------------
-	 */
-	mrs	x17, mdcr_el3
-	tst	x17, #MDCR_SCCD_BIT
-	bne	2f
-	ldr	x17, [sp, #CTX_EL3STATE_OFFSET + CTX_PMCR_EL0]
-	msr	pmcr_el0, x17
-2:
-
 #if IMAGE_BL31 && DYNAMIC_WORKAROUND_CVE_2018_3639
-	/* Restore mitigation state as it was on entry to EL3 */
+	/* ----------------------------------------------------------
+	 * Restore mitigation state as it was on entry to EL3
+	 * ----------------------------------------------------------
+	 */
 	ldr	x17, [sp, #CTX_CVE_2018_3639_OFFSET + CTX_CVE_2018_3639_DISABLE]
-	cmp	x17, xzr
-	beq	1f
+	cbz	x17, 1f
 	blr	x17
 1:
 #endif
+	/* ----------------------------------------------------------
+	 * Restore general purpose (including x30), PMCR_EL0 and
+	 * ARMv8.3-PAuth registers.
+	 * Exit EL3 via ERET to a lower exception level.
+ 	 * ----------------------------------------------------------
+ 	 */
+	bl	restore_gp_pmcr_pauth_regs
+	ldr	x30, [sp, #CTX_GPREGS_OFFSET + CTX_GPREG_LR]
 
-#if CTX_INCLUDE_PAUTH_REGS
-	/* Restore ARMv8.3-PAuth registers */
-	bl	pauth_context_restore
+#if IMAGE_BL31 && RAS_EXTENSION
+	/* ----------------------------------------------------------
+	 * Issue Error Synchronization Barrier to synchronize SErrors
+	 * before exiting EL3. We're running with EAs unmasked, so
+	 * any synchronized errors would be taken immediately;
+	 * therefore no need to inspect DISR_EL1 register.
+ 	 * ----------------------------------------------------------
+	 */
+	esb
 #endif
+	eret
 
-	/* Restore saved general purpose registers and return */
-	b	restore_gp_registers_eret
 endfunc el3_exit