build(mpam): add new build option CTX_INCLUDE_MPAM_REGS

New build option CTX_INCLUDE_MPAM_REGS is added to select
if the firmware needs to save the MPAM EL2 registers during world
switches. This option is currently disabled as MPAM is only
enabled for NS world.

Signed-off-by: Arvind Ram Prakash <arvind.ramprakash@arm.com>
Change-Id: Ie2e5e184cdb65f7e1a98d8fe81590253fd859679
diff --git a/Makefile b/Makefile
index a9c9c4f..e2de2c0 100644
--- a/Makefile
+++ b/Makefile
@@ -1170,6 +1170,7 @@
 	CTX_INCLUDE_AARCH32_REGS \
 	CTX_INCLUDE_FPREGS \
 	CTX_INCLUDE_EL2_REGS \
+	CTX_INCLUDE_MPAM_REGS \
 	DEBUG \
 	DYN_DISABLE_AUTH \
 	EL3_EXCEPTION_HANDLING \
@@ -1320,6 +1321,7 @@
 	CTX_INCLUDE_AARCH32_REGS \
 	CTX_INCLUDE_FPREGS \
 	CTX_INCLUDE_PAUTH_REGS \
+	CTX_INCLUDE_MPAM_REGS \
 	EL3_EXCEPTION_HANDLING \
 	CTX_INCLUDE_MTE_REGS \
 	CTX_INCLUDE_EL2_REGS \
diff --git a/docs/getting_started/build-options.rst b/docs/getting_started/build-options.rst
index 7ca0300..3d409e3 100644
--- a/docs/getting_started/build-options.rst
+++ b/docs/getting_started/build-options.rst
@@ -180,6 +180,11 @@
    registers to be included when saving and restoring the CPU context. Default
    is 0.
 
+-  ``CTX_INCLUDE_MPAM_REGS``: Boolean option that, when set to 1, will cause the
+   Memory System Resource Partitioning and Monitoring (MPAM)
+   registers to be included when saving and restoring the CPU context.
+   Default is '0'.
+
 -  ``CTX_INCLUDE_MTE_REGS``: Numeric value to include Memory Tagging Extension
    registers in cpu context. This must be enabled, if the platform wants to use
    this feature in the Secure world and MTE is enabled at ELX. This flag can
diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index 215156b..cdcfa39 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -207,18 +207,6 @@
 // Only if MTE registers in use
 #define CTX_TFSR_EL2		U(0x100)
 
-#define CTX_MPAM2_EL2		U(0x108)
-#define CTX_MPAMHCR_EL2		U(0x110)
-#define CTX_MPAMVPM0_EL2	U(0x118)
-#define CTX_MPAMVPM1_EL2	U(0x120)
-#define CTX_MPAMVPM2_EL2	U(0x128)
-#define CTX_MPAMVPM3_EL2	U(0x130)
-#define CTX_MPAMVPM4_EL2	U(0x138)
-#define CTX_MPAMVPM5_EL2	U(0x140)
-#define CTX_MPAMVPM6_EL2	U(0x148)
-#define CTX_MPAMVPM7_EL2	U(0x150)
-#define CTX_MPAMVPMV_EL2	U(0x158)
-
 // Starting with Armv8.6
 #define CTX_HDFGRTR_EL2		U(0x160)
 #define CTX_HAFGRTR_EL2		U(0x168)
@@ -338,6 +326,27 @@
 #endif /* CTX_INCLUDE_PAUTH_REGS */
 
 /*******************************************************************************
+ * Registers related to ARMv8.2-MPAM.
+ ******************************************************************************/
+#define CTX_MPAM_REGS_OFFSET	(CTX_PAUTH_REGS_OFFSET + CTX_PAUTH_REGS_END)
+#if CTX_INCLUDE_MPAM_REGS
+#define CTX_MPAM2_EL2		U(0x0)
+#define CTX_MPAMHCR_EL2		U(0x8)
+#define CTX_MPAMVPM0_EL2	U(0x10)
+#define CTX_MPAMVPM1_EL2	U(0x18)
+#define CTX_MPAMVPM2_EL2	U(0x20)
+#define CTX_MPAMVPM3_EL2	U(0x28)
+#define CTX_MPAMVPM4_EL2	U(0x30)
+#define CTX_MPAMVPM5_EL2	U(0x38)
+#define CTX_MPAMVPM6_EL2	U(0x40)
+#define CTX_MPAMVPM7_EL2	U(0x48)
+#define CTX_MPAMVPMV_EL2	U(0x50)
+#define CTX_MPAM_REGS_END	U(0x60)
+#else
+#define CTX_MPAM_REGS_END	U(0x0)
+#endif /* CTX_INCLUDE_MPAM_REGS */
+
+/*******************************************************************************
  * Registers initialised in a per-world context.
  ******************************************************************************/
 #define CTX_CPTR_EL3			U(0x0)
@@ -375,6 +384,9 @@
 #if CTX_INCLUDE_PAUTH_REGS
 # define CTX_PAUTH_REGS_ALL	(CTX_PAUTH_REGS_END >> DWORD_SHIFT)
 #endif
+#if CTX_INCLUDE_MPAM_REGS
+# define CTX_MPAM_REGS_ALL	(CTX_MPAM_REGS_END >> DWORD_SHIFT)
+#endif
 
 /*
  * AArch64 general purpose register context structure. Usually x0-x18,
@@ -423,6 +435,11 @@
 DEFINE_REG_STRUCT(pauth, CTX_PAUTH_REGS_ALL);
 #endif
 
+/* Registers associated to ARMv8.2 MPAM */
+#if CTX_INCLUDE_MPAM_REGS
+DEFINE_REG_STRUCT(mpam, CTX_MPAM_REGS_ALL);
+#endif
+
 /*
  * Macros to access members of any of the above structures using their
  * offsets
@@ -453,6 +470,9 @@
 #if CTX_INCLUDE_PAUTH_REGS
 	pauth_t pauth_ctx;
 #endif
+#if CTX_INCLUDE_MPAM_REGS
+	mpam_t	mpam_ctx;
+#endif
 } cpu_context_t;
 
 /*
@@ -481,6 +501,9 @@
 #if CTX_INCLUDE_PAUTH_REGS
 # define get_pauth_ctx(h)	(&((cpu_context_t *) h)->pauth_ctx)
 #endif
+#if CTX_INCLUDE_MPAM_REGS
+# define get_mpam_ctx(h)	(&((cpu_context_t *) h)->mpam_ctx)
+#endif
 
 /*
  * Compile time assertions related to the 'cpu_context' structure to
@@ -507,6 +530,10 @@
 CASSERT(CTX_PAUTH_REGS_OFFSET == __builtin_offsetof(cpu_context_t, pauth_ctx),
 	assert_core_context_pauth_offset_mismatch);
 #endif
+#if CTX_INCLUDE_MPAM_REGS
+CASSERT(CTX_MPAM_REGS_OFFSET == __builtin_offsetof(cpu_context_t, mpam_ctx),
+	assert_core_context_mpam_offset_mismatch);
+#endif
 
 /*
  * Helper macro to set the general purpose registers that correspond to
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 05628b5..9ba4d09 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -1048,7 +1048,9 @@
 	write_hfgwtr_el2(read_ctx_reg(ctx, CTX_HFGWTR_EL2));
 }
 
-static void el2_sysregs_context_save_mpam(el2_sysregs_t *ctx)
+#if CTX_INCLUDE_MPAM_REGS
+
+static void el2_sysregs_context_save_mpam(mpam_t *ctx)
 {
 	u_register_t mpam_idr = read_mpamidr_el1();
 
@@ -1099,7 +1101,10 @@
 	}
 }
 
-static void el2_sysregs_context_restore_mpam(el2_sysregs_t *ctx)
+#endif /* CTX_INCLUDE_MPAM_REGS */
+
+#if CTX_INCLUDE_MPAM_REGS
+static void el2_sysregs_context_restore_mpam(mpam_t *ctx)
 {
 	u_register_t mpam_idr = read_mpamidr_el1();
 
@@ -1137,6 +1142,7 @@
 		break;
 	}
 }
+#endif /* CTX_INCLUDE_MPAM_REGS */
 
 /* -----------------------------------------------------
  * The following registers are not added:
@@ -1264,9 +1270,13 @@
 #if CTX_INCLUDE_MTE_REGS
 	write_ctx_reg(el2_sysregs_ctx, CTX_TFSR_EL2, read_tfsr_el2());
 #endif
+
+#if CTX_INCLUDE_MPAM_REGS
 	if (is_feat_mpam_supported()) {
-		el2_sysregs_context_save_mpam(el2_sysregs_ctx);
+		mpam_t *mpam_ctx = get_mpam_ctx(ctx);
+		el2_sysregs_context_save_mpam(mpam_ctx);
 	}
+#endif
 
 	if (is_feat_fgt_supported()) {
 		el2_sysregs_context_save_fgt(el2_sysregs_ctx);
@@ -1337,9 +1347,13 @@
 #if CTX_INCLUDE_MTE_REGS
 	write_tfsr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_TFSR_EL2));
 #endif
+
+#if CTX_INCLUDE_MPAM_REGS
 	if (is_feat_mpam_supported()) {
-		el2_sysregs_context_restore_mpam(el2_sysregs_ctx);
+		mpam_t *mpam_ctx = get_mpam_ctx(ctx);
+		el2_sysregs_context_restore_mpam(mpam_ctx);
 	}
+#endif
 
 	if (is_feat_fgt_supported()) {
 		el2_sysregs_context_restore_fgt(el2_sysregs_ctx);
diff --git a/make_helpers/defaults.mk b/make_helpers/defaults.mk
index 8829fcb..ac9d44b 100644
--- a/make_helpers/defaults.mk
+++ b/make_helpers/defaults.mk
@@ -374,3 +374,7 @@
 # functions must be enabled by platforms if they require it.
 # Disabled by default.
 INIT_UNUSED_NS_EL2		:= 0
+
+# Disable including MPAM EL2 registers in context by default since currently
+# it's only enabled for NS world
+CTX_INCLUDE_MPAM_REGS		:= 0