diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index b058334..7bb0f99 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -17,10 +17,6 @@
 	.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_ECV
 	.global	el2_sysregs_context_save_ecv
 	.global	el2_sysregs_context_restore_ecv
@@ -230,208 +226,6 @@
 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	x10, MPAMIDR_EL1
-
-	/*
-	 * The context registers that we intend to save would be part of the
-	 * PE's system register frame only if MPAMIDR_EL1.HAS_HCR == 1.
-	 */
-	tbz	w10, #MPAMIDR_EL1_HAS_HCR_SHIFT, 3f
-
-	/*
-	 * MPAMHCR_EL2, MPAMVPMV_EL2 and MPAMVPM0_EL2 would be present in the
-	 * system register frame if MPAMIDR_EL1.HAS_HCR == 1. Proceed to save
-	 * the context of these registers.
-	 */
-	mrs	x11, MPAMHCR_EL2
-	mrs	x12, MPAMVPM0_EL2
-	stp	x11, x12, [x0, #CTX_MPAMHCR_EL2]
-
-	mrs	x13, MPAMVPMV_EL2
-	str	x13, [x0, #CTX_MPAMVPMV_EL2]
-
-	/*
-	 * MPAMIDR_EL1.VPMR_MAX has to be probed to obtain the maximum supported
-	 * VPMR value. Proceed to save the context of registers from
-	 * MPAMVPM1_EL2 to MPAMVPM<x>_EL2 where x is VPMR_MAX. From MPAM spec,
-	 * VPMR_MAX should not be zero if HAS_HCR == 1.
-	 */
-	ubfx	x10, x10, #MPAMIDR_EL1_VPMR_MAX_SHIFT, \
-		#MPAMIDR_EL1_VPMR_MAX_WIDTH
-
-	/*
-	 * Once VPMR_MAX has been identified, calculate the offset relative to
-	 * PC to jump to so that relevant context can be saved. The offset is
-	 * calculated as (VPMR_POSSIBLE_MAX - VPMR_MAX) * (instruction size for
-	 * saving one VPM register) + (absolute address of label "1").
-	 */
-	mov	w11, #MPAMIDR_EL1_VPMR_MAX_POSSIBLE
-	sub	w10, w11, w10
-
-	/* Calculate the size of one block of MPAMVPM*_EL2 save */
-	adr	x11, 1f
-	adr	x12, 2f
-	sub	x12, x12, x11
-
-	madd	x10, x10, x12, x11
-	br	x10
-
-	/*
-	 * The branch above would land properly on one of the blocks following
-	 * label "1". Make sure that the order of save is retained.
-	 */
-1:
-#if ENABLE_BTI
-	bti	j
-#endif
-	mrs	x10, MPAMVPM7_EL2
-	str	x10, [x0, #CTX_MPAMVPM7_EL2]
-2:
-#if ENABLE_BTI
-	bti	j
-#endif
-	mrs	x11, MPAMVPM6_EL2
-	str	x11, [x0, #CTX_MPAMVPM6_EL2]
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	mrs	x12, MPAMVPM5_EL2
-	str	x12, [x0, #CTX_MPAMVPM5_EL2]
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	mrs	x13, MPAMVPM4_EL2
-	str	x13, [x0, #CTX_MPAMVPM4_EL2]
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	mrs	x14, MPAMVPM3_EL2
-	str	x14, [x0, #CTX_MPAMVPM3_EL2]
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	mrs	x15, MPAMVPM2_EL2
-	str	x15, [x0, #CTX_MPAMVPM2_EL2]
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	mrs	x16, MPAMVPM1_EL2
-	str	x16, [x0, #CTX_MPAMVPM1_EL2]
-
-3:	ret
-endfunc el2_sysregs_context_save_mpam
-
-func el2_sysregs_context_restore_mpam
-	ldr	x10, [x0, #CTX_MPAM2_EL2]
-	msr	MPAM2_EL2, x10
-
-	mrs	x10, MPAMIDR_EL1
-	/*
-	 * The context registers that we intend to restore would be part of the
-	 * PE's system register frame only if MPAMIDR_EL1.HAS_HCR == 1.
-	 */
-	tbz	w10, #MPAMIDR_EL1_HAS_HCR_SHIFT, 3f
-
-	/*
-	 * MPAMHCR_EL2, MPAMVPMV_EL2 and MPAMVPM0_EL2 would be present in the
-	 * system register frame if MPAMIDR_EL1.HAS_HCR == 1. Proceed to restore
-	 * the context of these registers
-	 */
-	ldp	x11, x12, [x0, #CTX_MPAMHCR_EL2]
-	msr	MPAMHCR_EL2, x11
-	msr	MPAMVPM0_EL2, x12
-
-	ldr	x13, [x0, #CTX_MPAMVPMV_EL2]
-	msr	MPAMVPMV_EL2, x13
-
-	/*
-	 * MPAMIDR_EL1.VPMR_MAX has to be probed to obtain the maximum supported
-	 * VPMR value. Proceed to restore the context of registers from
-	 * MPAMVPM1_EL2 to MPAMVPM<x>_EL2 where x is VPMR_MAX. from MPAM spec,
-	 * VPMR_MAX should not be zero if HAS_HCR == 1.
-	 */
-	ubfx	x10, x10, #MPAMIDR_EL1_VPMR_MAX_SHIFT,	\
-		#MPAMIDR_EL1_VPMR_MAX_WIDTH
-
-	/*
-	 * Once VPMR_MAX has been identified, calculate the offset relative to
-	 * PC to jump to so that relevant context can be restored. The offset is
-	 * calculated as (VPMR_POSSIBLE_MAX - VPMR_MAX) * (instruction size for
-	 * restoring one VPM register) + (absolute address of label "1").
-	 */
-	mov	w11, #MPAMIDR_EL1_VPMR_MAX_POSSIBLE
-	sub	w10, w11, w10
-
-	/* Calculate the size of one block of MPAMVPM*_EL2 restore */
-	adr	x11, 1f
-	adr	x12, 2f
-	sub	x12, x12, x11
-
-	madd	x10, x10, x12, x11
-	br	x10
-
-	/*
-	 * The branch above would land properly on one of the blocks following
-	 * label "1". Make sure that the order of restore is retained.
-	 */
-1:
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	ldr	x10, [x0, #CTX_MPAMVPM7_EL2]
-	msr	MPAMVPM7_EL2, x10
-2:
-#if ENABLE_BTI
-	bti	j
-#endif
-	ldr	x11, [x0, #CTX_MPAMVPM6_EL2]
-	msr	MPAMVPM6_EL2, x11
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	ldr	x12, [x0, #CTX_MPAMVPM5_EL2]
-	msr	MPAMVPM5_EL2, x12
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	ldr	x13, [x0, #CTX_MPAMVPM4_EL2]
-	msr	MPAMVPM4_EL2, x13
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	ldr	x14, [x0, #CTX_MPAMVPM3_EL2]
-	msr	MPAMVPM3_EL2, x14
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	ldr	x15, [x0, #CTX_MPAMVPM2_EL2]
-	msr	MPAMVPM2_EL2, x15
-
-#if ENABLE_BTI
-	bti	j
-#endif
-	ldr	x16, [x0, #CTX_MPAMVPM1_EL2]
-	msr	MPAMVPM1_EL2, x16
-
-3:	ret
-endfunc el2_sysregs_context_restore_mpam
-#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
-
 #if ENABLE_FEAT_ECV
 func el2_sysregs_context_save_ecv
 	mrs	x11, CNTPOFF_EL2
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index c5f20ad..eb5d1db 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -498,9 +498,9 @@
 	sve_enable(ctx);
 #endif
 
-#if ENABLE_MPAM_FOR_LOWER_ELS
-	mpam_enable(el2_unused);
-#endif
+	if (is_feat_mpam_supported()) {
+		mpam_enable(el2_unused);
+	}
 
 	if (is_feat_trbe_supported()) {
 		trbe_enable();
@@ -834,6 +834,96 @@
 	write_hfgwtr_el2(read_ctx_reg(ctx, CTX_HFGWTR_EL2));
 }
 
+static void el2_sysregs_context_save_mpam(el2_sysregs_t *ctx)
+{
+	u_register_t mpam_idr = read_mpamidr_el1();
+
+	write_ctx_reg(ctx, CTX_MPAM2_EL2, read_mpam2_el2());
+
+	/*
+	 * The context registers that we intend to save would be part of the
+	 * PE's system register frame only if MPAMIDR_EL1.HAS_HCR == 1.
+	 */
+	if ((mpam_idr & MPAMIDR_HAS_HCR_BIT) == 0U) {
+		return;
+	}
+
+	/*
+	 * MPAMHCR_EL2, MPAMVPMV_EL2 and MPAMVPM0_EL2 are always present if
+	 * MPAMIDR_HAS_HCR_BIT == 1.
+	 */
+	write_ctx_reg(ctx, CTX_MPAMHCR_EL2, read_mpamhcr_el2());
+	write_ctx_reg(ctx, CTX_MPAMVPM0_EL2, read_mpamvpm0_el2());
+	write_ctx_reg(ctx, CTX_MPAMVPMV_EL2, read_mpamvpmv_el2());
+
+	/*
+	 * The number of MPAMVPM registers is implementation defined, their
+	 * number is stored in the MPAMIDR_EL1 register.
+	 */
+	switch ((mpam_idr >> MPAMIDR_EL1_VPMR_MAX_SHIFT) & MPAMIDR_EL1_VPMR_MAX_MASK) {
+	case 7:
+		write_ctx_reg(ctx, CTX_MPAMVPM7_EL2, read_mpamvpm7_el2());
+		__fallthrough;
+	case 6:
+		write_ctx_reg(ctx, CTX_MPAMVPM6_EL2, read_mpamvpm6_el2());
+		__fallthrough;
+	case 5:
+		write_ctx_reg(ctx, CTX_MPAMVPM5_EL2, read_mpamvpm5_el2());
+		__fallthrough;
+	case 4:
+		write_ctx_reg(ctx, CTX_MPAMVPM4_EL2, read_mpamvpm4_el2());
+		__fallthrough;
+	case 3:
+		write_ctx_reg(ctx, CTX_MPAMVPM3_EL2, read_mpamvpm3_el2());
+		__fallthrough;
+	case 2:
+		write_ctx_reg(ctx, CTX_MPAMVPM2_EL2, read_mpamvpm2_el2());
+		__fallthrough;
+	case 1:
+		write_ctx_reg(ctx, CTX_MPAMVPM1_EL2, read_mpamvpm1_el2());
+		break;
+	}
+}
+
+static void el2_sysregs_context_restore_mpam(el2_sysregs_t *ctx)
+{
+	u_register_t mpam_idr = read_mpamidr_el1();
+
+	write_mpam2_el2(read_ctx_reg(ctx, CTX_MPAM2_EL2));
+
+	if ((mpam_idr & MPAMIDR_HAS_HCR_BIT) == 0U) {
+		return;
+	}
+
+	write_mpamhcr_el2(read_ctx_reg(ctx, CTX_MPAMHCR_EL2));
+	write_mpamvpm0_el2(read_ctx_reg(ctx, CTX_MPAMVPM0_EL2));
+	write_mpamvpmv_el2(read_ctx_reg(ctx, CTX_MPAMVPMV_EL2));
+
+	switch ((mpam_idr >> MPAMIDR_EL1_VPMR_MAX_SHIFT) & MPAMIDR_EL1_VPMR_MAX_MASK) {
+	case 7:
+		write_mpamvpm7_el2(read_ctx_reg(ctx, CTX_MPAMVPM7_EL2));
+		__fallthrough;
+	case 6:
+		write_mpamvpm6_el2(read_ctx_reg(ctx, CTX_MPAMVPM6_EL2));
+		__fallthrough;
+	case 5:
+		write_mpamvpm5_el2(read_ctx_reg(ctx, CTX_MPAMVPM5_EL2));
+		__fallthrough;
+	case 4:
+		write_mpamvpm4_el2(read_ctx_reg(ctx, CTX_MPAMVPM4_EL2));
+		__fallthrough;
+	case 3:
+		write_mpamvpm3_el2(read_ctx_reg(ctx, CTX_MPAMVPM3_EL2));
+		__fallthrough;
+	case 2:
+		write_mpamvpm2_el2(read_ctx_reg(ctx, CTX_MPAMVPM2_EL2));
+		__fallthrough;
+	case 1:
+		write_mpamvpm1_el2(read_ctx_reg(ctx, CTX_MPAMVPM1_EL2));
+		break;
+	}
+}
+
 /*******************************************************************************
  * Save EL2 sysreg context
  ******************************************************************************/
@@ -859,9 +949,9 @@
 #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 (is_feat_mpam_supported()) {
+			el2_sysregs_context_save_mpam(el2_sysregs_ctx);
+		}
 
 		if (is_feat_fgt_supported()) {
 			el2_sysregs_context_save_fgt(el2_sysregs_ctx);
@@ -919,9 +1009,9 @@
 #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 (is_feat_mpam_supported()) {
+			el2_sysregs_context_restore_mpam(el2_sysregs_ctx);
+		}
 
 		if (is_feat_fgt_supported()) {
 			el2_sysregs_context_restore_fgt(el2_sysregs_ctx);