feat(pmuv3): setup per world MDCR_EL3

MDCR_EL3 register will context switch across all worlds. Thus the pmuv3
init has to be part of context management initialization.

Change-Id: I10ef7a3071c0fc5c11a93d3c9c2a95ec8c6493bf
Signed-off-by: Mateusz Sulimowicz <matsul@google.com>
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index f220d8a..3ad7133 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -573,6 +573,8 @@
 	if (is_feat_trf_supported()) {
 		trf_enable(ctx);
 	}
+
+	pmuv3_enable(ctx);
 #endif /* IMAGE_BL31 */
 
 	/*
@@ -822,8 +824,6 @@
 	if (is_feat_brbe_supported()) {
 		brbe_enable(ctx);
 	}
-
-	pmuv3_enable(ctx);
 #endif /* IMAGE_BL31 */
 }
 
diff --git a/lib/extensions/pmuv3/aarch64/pmuv3.c b/lib/extensions/pmuv3/aarch64/pmuv3.c
index f9e32ca..61d1258 100644
--- a/lib/extensions/pmuv3/aarch64/pmuv3.c
+++ b/lib/extensions/pmuv3/aarch64/pmuv3.c
@@ -21,17 +21,6 @@
 	return mdcr_el2;
 }
 
-void pmuv3_enable(cpu_context_t *ctx)
-{
-#if (CTX_INCLUDE_EL2_REGS && IMAGE_BL31)
-	u_register_t mdcr_el2_val;
-
-	mdcr_el2_val = read_el2_ctx_common(get_el2_sysregs_ctx(ctx), mdcr_el2);
-	mdcr_el2_val = init_mdcr_el2_hpmn(mdcr_el2_val);
-	write_el2_ctx_common(get_el2_sysregs_ctx(ctx), mdcr_el2, mdcr_el2_val);
-#endif /* (CTX_INCLUDE_EL2_REGS && IMAGE_BL31) */
-}
-
 static u_register_t mtpmu_disable_el3(u_register_t mdcr_el3)
 {
 	if (!is_feat_mtpmu_supported()) {
@@ -48,14 +37,20 @@
 	return mdcr_el3;
 }
 
-void pmuv3_init_el3(void)
+void pmuv3_enable(cpu_context_t *ctx)
 {
-	u_register_t mdcr_el3 = read_mdcr_el3();
+#if (CTX_INCLUDE_EL2_REGS && IMAGE_BL31)
+	u_register_t mdcr_el2_val;
+
+	mdcr_el2_val = read_el2_ctx_common(get_el2_sysregs_ctx(ctx), mdcr_el2);
+	mdcr_el2_val = init_mdcr_el2_hpmn(mdcr_el2_val);
+	write_el2_ctx_common(get_el2_sysregs_ctx(ctx), mdcr_el2, mdcr_el2_val);
+#endif /* (CTX_INCLUDE_EL2_REGS && IMAGE_BL31) */
+
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/* ---------------------------------------------------------------------
-	 * Initialise MDCR_EL3, setting all fields rather than relying on hw.
-	 * Some fields are architecturally UNKNOWN on reset.
-	 *
 	 * MDCR_EL3.MPMX: Set to zero to not affect event counters (when
 	 * SPME = 0).
 	 *
@@ -86,11 +81,15 @@
 	 * MDCR_EL3.TPM: Set to zero so that EL0, EL1, and EL2 System register
 	 *  accesses to all Performance Monitors registers do not trap to EL3.
 	 */
-	mdcr_el3 = (mdcr_el3 | MDCR_SCCD_BIT | MDCR_MCCD_BIT) &
+	mdcr_el3_val = (mdcr_el3_val | MDCR_SCCD_BIT | MDCR_MCCD_BIT) &
 		  ~(MDCR_MPMX_BIT | MDCR_SPME_BIT | MDCR_TPM_BIT);
-	mdcr_el3 = mtpmu_disable_el3(mdcr_el3);
-	write_mdcr_el3(mdcr_el3);
+	mdcr_el3_val = mtpmu_disable_el3(mdcr_el3_val);
+
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
+}
 
+void pmuv3_init_el3(void)
+{
 	/* ---------------------------------------------------------------------
 	 * Initialise PMCR_EL0 setting all fields rather than relying
 	 * on hw. Some fields are architecturally UNKNOWN on reset.