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
diff --git a/Makefile b/Makefile
index 6a2eeca..5edd625 100644
--- a/Makefile
+++ b/Makefile
@@ -1270,7 +1270,7 @@
 	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 @@
 	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 0c1d657..a1fc12b 100644
--- a/bl31/bl31.mk
+++ b/bl31/bl31.mk
@@ -106,7 +106,7 @@
 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 a1ffc39..be22c6e 100644
--- a/common/feat_detect.c
+++ b/common/feat_detect.c
@@ -169,7 +169,7 @@
 	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 cd70a22..654ddc5 100644
--- a/docs/getting_started/build-options.rst
+++ b/docs/getting_started/build-options.rst
@@ -397,7 +397,7 @@
    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 @@
    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 c10102a..001fad5 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 bd41fef..cf8da5e 100644
--- a/include/arch/aarch64/arch_features.h
+++ b/include/arch/aarch64/arch_features.h
@@ -159,7 +159,7 @@
 }
 
 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 ebd0e30..e7e9f58 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 e5438ce..170f919 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 290a93a..758355a 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -350,6 +350,43 @@
 #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 @@
 	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 3a7e50f..98cee16 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -482,6 +482,11 @@
 	}
 #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 @@
 		sme_init_el3();
 	}
 
-	if (is_feat_mpam_supported()) {
-		mpam_init_el3();
-	}
-
 	if (is_feat_trbe_supported()) {
 		trbe_init_el3();
 	}
@@ -621,6 +622,9 @@
 		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 6462c97..875ad9c 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 f90136b..8ade8e5 100644
--- a/make_helpers/arch_features.mk
+++ b/make_helpers/arch_features.mk
@@ -214,7 +214,16 @@
 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 bd56f30..c71c99a 100644
--- a/plat/arm/board/arm_fpga/platform.mk
+++ b/plat/arm/board/arm_fpga/platform.mk
@@ -41,7 +41,7 @@
 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 72f94b0..7a2b5ce 100644
--- a/plat/arm/board/fvp/platform.mk
+++ b/plat/arm/board/fvp/platform.mk
@@ -53,7 +53,7 @@
 	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