fix(mpam): refine MPAM initialization and enablement process

Restricts MPAM to only NS world and enables trap to EL3 for access of
MPAM registers from lower ELs of Secure and Realm world.

This patch removes MPAM enablement from global context and adds it to
EL3 State context which enables/disables MPAM during world switches.
Renamed ENABLE_MPAM_FOR_LOWER_ELS to ENABLE_FEAT_MPAM and
removed mpam_init_el3() as RESET behaviour is trapping.

Signed-off-by: Arvind Ram Prakash <arvind.ramprakash@arm.com>
Change-Id: I131f9dba5df236a71959b2d425ee11af7f3c38c4

14 files changed, 87 insertions, 22 deletions

diff --git a/Makefile b/Makefile
index 6a2eeca1a..5edd6259c 100644
--- a/Makefile
+++ b/Makefile
@@ -1270,7 +1270,7 @@ $(eval $(call assert_numerics,\
 	ENABLE_FEAT_GCS \
 	ENABLE_FEAT_VHE \
 	ENABLE_FEAT_MTE_PERM \
-	ENABLE_MPAM_FOR_LOWER_ELS \
+	ENABLE_FEAT_MPAM \
 	ENABLE_RME \
 	ENABLE_SPE_FOR_NS \
 	ENABLE_SYS_REG_TRACE_FOR_NS \
@@ -1323,7 +1323,7 @@ $(eval $(call add_defines,\
 	AMU_RESTRICT_COUNTERS \
 	ENABLE_ASSERTIONS \
 	ENABLE_BTI \
-	ENABLE_MPAM_FOR_LOWER_ELS \
+	ENABLE_FEAT_MPAM \
 	ENABLE_PAUTH \
 	ENABLE_PIE \
 	ENABLE_PMF \
diff --git a/bl31/bl31.mk b/bl31/bl31.mk
index 0c1d65753..a1fc12be8 100644
--- a/bl31/bl31.mk
+++ b/bl31/bl31.mk
@@ -106,7 +106,7 @@ ifneq (${ENABLE_SVE_FOR_NS},0)
 BL31_SOURCES		+=	lib/extensions/sve/sve.c
 endif
 
-ifneq (${ENABLE_MPAM_FOR_LOWER_ELS},0)
+ifneq (${ENABLE_FEAT_MPAM},0)
 BL31_SOURCES		+=	lib/extensions/mpam/mpam.c
 endif
 
diff --git a/common/feat_detect.c b/common/feat_detect.c
index a1ffc39d4..be22c6ed1 100644
--- a/common/feat_detect.c
+++ b/common/feat_detect.c
@@ -169,7 +169,7 @@ void detect_arch_features(void)
 	check_feature(ENABLE_FEAT_DIT, read_feat_dit_id_field(), "DIT", 1, 1);
 	check_feature(ENABLE_FEAT_AMU, read_feat_amu_id_field(),
 		      "AMUv1", 1, 2);
-	check_feature(ENABLE_MPAM_FOR_LOWER_ELS, read_feat_mpam_version(),
+	check_feature(ENABLE_FEAT_MPAM, read_feat_mpam_version(),
 		      "MPAM", 1, 17);
 	check_feature(CTX_INCLUDE_NEVE_REGS, read_feat_nv_id_field(),
 		      "NV2", 2, 2);
diff --git a/docs/getting_started/build-options.rst b/docs/getting_started/build-options.rst
index cd70a2275..654ddc54d 100644
--- a/docs/getting_started/build-options.rst
+++ b/docs/getting_started/build-options.rst
@@ -397,7 +397,7 @@ Common build options
    support in GCC for TF-A. This option is currently only supported for
    AArch64. Default is 0.
 
--  ``ENABLE_MPAM_FOR_LOWER_ELS``: Numeric value to enable lower ELs to use MPAM
+-  ``ENABLE_FEAT_MPAM``: Numeric value to enable lower ELs to use MPAM
    feature. MPAM is an optional Armv8.4 extension that enables various memory
    system components and resources to define partitions; software running at
    various ELs can assign themselves to desired partition to control their
@@ -408,7 +408,9 @@ Common build options
    access their own MPAM registers without trapping into EL3. This option
    doesn't make use of partitioning in EL3, however. Platform initialisation
    code should configure and use partitions in EL3 as required. This option
-   defaults to ``0``.
+   defaults to ``2`` since MPAM is enabled by default for NS world only.
+   The flag is automatically disabled when the target
+   architecture is AArch32.
 
 -  ``ENABLE_MPMM``: Boolean option to enable support for the Maximum Power
    Mitigation Mechanism supported by certain Arm cores, which allows the SoC
diff --git a/include/arch/aarch64/arch.h b/include/arch/aarch64/arch.h
index c10102a56..001fad5c2 100644
--- a/include/arch/aarch64/arch.h
+++ b/include/arch/aarch64/arch.h
@@ -1231,7 +1231,9 @@
 
 /* MPAM register definitions */
 #define MPAM3_EL3_MPAMEN_BIT		(ULL(1) << 63)
+#define MPAM3_EL3_TRAPLOWER_BIT		(ULL(1) << 62)
 #define MPAMHCR_EL2_TRAP_MPAMIDR_EL1	(ULL(1) << 31)
+#define MPAM3_EL3_RESET_VAL		MPAM3_EL3_TRAPLOWER_BIT
 
 #define MPAM2_EL2_TRAPMPAM0EL1		(ULL(1) << 49)
 #define MPAM2_EL2_TRAPMPAM1EL1		(ULL(1) << 48)
diff --git a/include/arch/aarch64/arch_features.h b/include/arch/aarch64/arch_features.h
index bd41fefde..cf8da5e8a 100644
--- a/include/arch/aarch64/arch_features.h
+++ b/include/arch/aarch64/arch_features.h
@@ -159,7 +159,7 @@ static inline unsigned int read_feat_mpam_version(void)
 }
 
 CREATE_FEATURE_FUNCS_VER(feat_mpam, read_feat_mpam_version, 1U,
-			 ENABLE_MPAM_FOR_LOWER_ELS)
+			 ENABLE_FEAT_MPAM)
 
 /* FEAT_HCX: Extended Hypervisor Configuration Register */
 CREATE_FEATURE_FUNCS(feat_hcx, id_aa64mmfr1_el1, ID_AA64MMFR1_EL1_HCX_SHIFT,
diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index ebd0e3055..e7e9f589b 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -63,6 +63,7 @@
 #define CTX_IS_IN_EL3		U(0x30)
 #define CTX_CPTR_EL3		U(0x38)
 #define CTX_ZCR_EL3		U(0x40)
+#define CTX_MPAM3_EL3		U(0x48)
 #define CTX_EL3STATE_END	U(0x50) /* Align to the next 16 byte boundary */
 
 /*******************************************************************************
diff --git a/include/lib/extensions/mpam.h b/include/lib/extensions/mpam.h
index e5438cec6..170f919bb 100644
--- a/include/lib/extensions/mpam.h
+++ b/include/lib/extensions/mpam.h
@@ -9,16 +9,18 @@
 
 #include <stdbool.h>
 
-#if ENABLE_MPAM_FOR_LOWER_ELS
-void mpam_init_el3(void);
+#include <context.h>
+
+#if ENABLE_FEAT_MPAM
+void mpam_enable(cpu_context_t *context);
 void mpam_init_el2_unused(void);
 #else
-static inline void mpam_init_el3(void)
+static inline void mpam_enable(cpu_context_t *context)
 {
 }
 static inline void mpam_init_el2_unused(void)
 {
 }
-#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
+#endif /* ENABLE_FEAT_MPAM */
 
 #endif /* MPAM_H */
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 290a93a8a..758355aae 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -350,6 +350,43 @@ endfunc fpregs_context_restore
 #endif /* ENABLE_FEAT_DIT */
 	.endm /* set_unset_pstate_bits */
 
+/*-------------------------------------------------------------------------
+ * This macro checks the ENABLE_FEAT_MPAM state, performs ID register
+ * check to see if the platform supports MPAM extension and restores MPAM3
+ * register value if it is FEAT_STATE_ENABLED/FEAT_STATE_CHECKED.
+ *
+ * This is particularly more complicated because we can't check
+ * if the platform supports MPAM  by looking for status of a particular bit
+ * in the MDCR_EL3 or CPTR_EL3 register like other extensions.
+ * ------------------------------------------------------------------------
+ */
+
+	.macro	restore_mpam3_el3
+#if ENABLE_FEAT_MPAM
+#if ENABLE_FEAT_MPAM == 2
+
+	mrs x8, id_aa64pfr0_el1
+	lsr x8, x8, #(ID_AA64PFR0_MPAM_SHIFT)
+	and x8, x8, #(ID_AA64PFR0_MPAM_MASK)
+	mrs x7, id_aa64pfr1_el1
+	lsr x7, x7, #(ID_AA64PFR1_MPAM_FRAC_SHIFT)
+	and x7, x7, #(ID_AA64PFR1_MPAM_FRAC_MASK)
+	orr x7, x7, x8
+	cbz x7, no_mpam
+#endif
+	/* -----------------------------------------------------------
+	 * Restore MPAM3_EL3 register as per context state
+	 * Currently we only enable MPAM for NS world and trap to EL3
+	 * for MPAM access in lower ELs of Secure and Realm world
+	 * -----------------------------------------------------------
+	 */
+	ldr	x17, [sp, #CTX_EL3STATE_OFFSET + CTX_MPAM3_EL3]
+	msr	S3_6_C10_C5_0, x17 /* mpam3_el3 */
+
+no_mpam:
+#endif
+	.endm /* restore_mpam3_el3 */
+
 /* ------------------------------------------------------------------
  * The following macro is used to save and restore all the general
  * purpose and ARMv8.3-PAuth (if enabled) registers.
@@ -573,6 +610,9 @@ func el3_exit
 	isb
 	msr	S3_6_C1_C2_0, x20 /* zcr_el3 */
 sve_not_enabled:
+
+	restore_mpam3_el3
+
 #endif /* IMAGE_BL31 */
 
 #if IMAGE_BL31 && DYNAMIC_WORKAROUND_CVE_2018_3639
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 3a7e50f73..98cee1667 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -482,6 +482,11 @@ static void setup_context_common(cpu_context_t *ctx, const entry_point_info_t *e
 	}
 #endif /* (IMAGE_BL31 && defined(SPD_spmd) && SPMD_SPM_AT_SEL2) */
 
+	if (is_feat_mpam_supported()) {
+		write_ctx_reg(get_el3state_ctx(ctx), CTX_MPAM3_EL3, \
+				MPAM3_EL3_RESET_VAL);
+	}
+
 	/*
 	 * Populate EL3 state so that we've the right context
 	 * before doing ERET
@@ -578,10 +583,6 @@ void cm_manage_extensions_el3(void)
 		sme_init_el3();
 	}
 
-	if (is_feat_mpam_supported()) {
-		mpam_init_el3();
-	}
-
 	if (is_feat_trbe_supported()) {
 		trbe_init_el3();
 	}
@@ -621,6 +622,9 @@ static void manage_extensions_nonsecure(cpu_context_t *ctx)
 		sys_reg_trace_enable(ctx);
 	}
 
+	if (is_feat_mpam_supported()) {
+		mpam_enable(ctx);
+	}
 	pmuv3_enable(ctx);
 #endif /* IMAGE_BL31 */
 }
diff --git a/lib/extensions/mpam/mpam.c b/lib/extensions/mpam/mpam.c
index 6462c9748..875ad9c2a 100644
--- a/lib/extensions/mpam/mpam.c
+++ b/lib/extensions/mpam/mpam.c
@@ -11,14 +11,19 @@
 #include <arch_helpers.h>
 #include <lib/extensions/mpam.h>
 
-void mpam_init_el3(void)
+void mpam_enable(cpu_context_t *context)
 {
+	u_register_t mpam3_el3;
+
+	mpam3_el3 = read_ctx_reg(get_el3state_ctx(context), CTX_MPAM3_EL3);
+
 	/*
 	 * Enable MPAM, and disable trapping to EL3 when lower ELs access their
-	 * own MPAM registers.
+	 * own MPAM registers
 	 */
-	write_mpam3_el3(MPAM3_EL3_MPAMEN_BIT);
-
+	mpam3_el3 = (mpam3_el3 | MPAM3_EL3_MPAMEN_BIT) &
+				~(MPAM3_EL3_TRAPLOWER_BIT);
+	write_ctx_reg(get_el3state_ctx(context), CTX_MPAM3_EL3, mpam3_el3);
 }
 
 /*
diff --git a/make_helpers/arch_features.mk b/make_helpers/arch_features.mk
index f90136b78..8ade8e5e0 100644
--- a/make_helpers/arch_features.mk
+++ b/make_helpers/arch_features.mk
@@ -214,7 +214,16 @@ ENABLE_AMU_FCONF			?=	0
 AMU_RESTRICT_COUNTERS			?=	0
 
 # Build option to enable MPAM for lower ELs.
-ENABLE_MPAM_FOR_LOWER_ELS		?=	0
+# Enabling it by default
+ifeq (${ARCH},aarch64)
+	ENABLE_FEAT_MPAM		?=	2
+else ifeq (${ARCH},aarch32)
+	ifdef ENABLE_FEAT_MPAM
+		$(error ENABLE_FEAT_MPAM is not supported for AArch32)
+	else
+		ENABLE_FEAT_MPAM	:=	0
+	endif
+endif
 
 # Include nested virtualization control (Armv8.4-NV) registers in cpu context.
 # This must be set to 1 if architecture implements Nested Virtualization
diff --git a/plat/arm/board/arm_fpga/platform.mk b/plat/arm/board/arm_fpga/platform.mk
index bd56f30db..c71c99a21 100644
--- a/plat/arm/board/arm_fpga/platform.mk
+++ b/plat/arm/board/arm_fpga/platform.mk
@@ -41,7 +41,7 @@ ENABLE_FEAT_CSV2_2		:= 2
 ENABLE_FEAT_ECV			:= 2
 ENABLE_FEAT_FGT			:= 2
 ENABLE_FEAT_HCX			:= 2
-ENABLE_MPAM_FOR_LOWER_ELS	:= 2
+ENABLE_FEAT_MPAM		:= 2
 ENABLE_SYS_REG_TRACE_FOR_NS	:= 2
 ENABLE_TRF_FOR_NS		:= 2
 
diff --git a/plat/arm/board/fvp/platform.mk b/plat/arm/board/fvp/platform.mk
index 72f94b0a1..7a2b5ced2 100644
--- a/plat/arm/board/fvp/platform.mk
+++ b/plat/arm/board/fvp/platform.mk
@@ -53,7 +53,7 @@ ifneq (${SPD}, tspd)
 	ENABLE_FEAT_AMU			:= 2
 	ENABLE_FEAT_AMUv1p1		:= 2
 	ENABLE_FEAT_HCX			:= 2
-	ENABLE_MPAM_FOR_LOWER_ELS	:= 2
+	ENABLE_FEAT_MPAM		:= 2
 	ENABLE_FEAT_RNG			:= 2
 	ENABLE_FEAT_TWED		:= 2
 	ENABLE_FEAT_GCS			:= 2