refactor(amu): separate the EL2 and EL3 enablement code

Combining the EL2 and EL3 enablement code necessitates that it must be
called at el3_exit, which is the only place with enough context to make
the decision of what needs to be set.
Decouple them to allow them to be called from elsewhere. Also take
some time to clarify and simplify AMU code.

The sanity check in the context_restore() is now wrong, as the cpu may
turn off on suspend, thus resetting the value of the counter enables.
Remove it.

Finally, this completes the migration to cm_manage_extensions_el3() and
manage_extensions_nonsecure() so manage_extensions_nonsecure_mixed() is
being removed.

Signed-off-by: Boyan Karatotev <boyan.karatotev@arm.com>
Change-Id: I66399132364c32be66017506bb54cbadd8485577
diff --git a/include/lib/extensions/amu.h b/include/lib/extensions/amu.h
index de476e4..80ad68c 100644
--- a/include/lib/extensions/amu.h
+++ b/include/lib/extensions/amu.h
@@ -16,13 +16,21 @@
 
 #if ENABLE_FEAT_AMU
 #if __aarch64__
-void amu_enable(bool el2_unused, cpu_context_t *ctx);
+void amu_enable(cpu_context_t *ctx);
+void amu_init_el3(void);
+void amu_init_el2_unused(void);
 #else
 void amu_enable(bool el2_unused);
 #endif
 #else
 #if __aarch64__
-static inline void amu_enable(bool el2_unused, cpu_context_t *ctx)
+void amu_enable(cpu_context_t *ctx)
+{
+}
+void amu_init_el3(void)
+{
+}
+void amu_init_el2_unused(void)
 {
 }
 #else
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index b7d014a..d9ff0b6 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -496,22 +496,6 @@
 }
 
 /*******************************************************************************
- * Enable architecture extensions on first entry to Non-secure world.
- * When EL2 is implemented but unused `el2_unused` is non-zero, otherwise
- * it is zero. This function updates some registers in-place and its contents
- * are being prepared to be moved to cm_manage_extensions_el3 and
- * cm_manage_extensions_nonsecure.
- ******************************************************************************/
-static void manage_extensions_nonsecure_mixed(bool el2_unused, cpu_context_t *ctx)
-{
-#if IMAGE_BL31
-	if (is_feat_amu_supported()) {
-		amu_enable(el2_unused, ctx);
-	}
-#endif /* IMAGE_BL31 */
-}
-
-/*******************************************************************************
  * Enable architecture extensions for EL3 execution. This function only updates
  * registers in-place which are expected to either never change or be
  * overwritten by el3_exit.
@@ -523,6 +507,10 @@
 		spe_init_el3();
 	}
 
+	if (is_feat_amu_supported()) {
+		amu_init_el3();
+	}
+
 	if (is_feat_sme_supported()) {
 		sme_init_el3();
 	}
@@ -553,6 +541,10 @@
 static void manage_extensions_nonsecure(cpu_context_t *ctx)
 {
 #if IMAGE_BL31
+	if (is_feat_amu_supported()) {
+		amu_enable(ctx);
+	}
+
 	/* Enable SVE and FPU/SIMD */
 	if (is_feat_sve_supported()) {
 		sve_enable(ctx);
@@ -581,6 +573,10 @@
 		spe_init_el2_unused();
 	}
 
+	if (is_feat_amu_supported()) {
+		amu_init_el2_unused();
+	}
+
 	if (is_feat_mpam_supported()) {
 		mpam_init_el2_unused();
 	}
@@ -689,7 +685,6 @@
 {
 	u_register_t sctlr_elx, scr_el3, mdcr_el2;
 	cpu_context_t *ctx = cm_get_context(security_state);
-	bool el2_unused = false;
 	uint64_t hcr_el2 = 0U;
 
 	assert(ctx != NULL);
@@ -727,8 +722,6 @@
 #endif
 			write_sctlr_el2(sctlr_elx);
 		} else if (el2_implemented != EL_IMPL_NONE) {
-			el2_unused = true;
-
 			/*
 			 * EL2 present but unused, need to disable safely.
 			 * SCTLR_EL2 can be ignored in this case.
@@ -845,7 +838,6 @@
 
 			manage_extensions_nonsecure_el2_unused();
 		}
-		manage_extensions_nonsecure_mixed(el2_unused, ctx);
 	}
 
 	cm_el1_sysregs_context_restore(security_state);
@@ -1150,23 +1142,15 @@
 void cm_prepare_el3_exit_ns(void)
 {
 #if CTX_INCLUDE_EL2_REGS
+#if ENABLE_ASSERTIONS
 	cpu_context_t *ctx = cm_get_context(NON_SECURE);
 	assert(ctx != NULL);
 
 	/* Assert that EL2 is used. */
-#if ENABLE_ASSERTIONS
-	el3_state_t *state = get_el3state_ctx(ctx);
-	u_register_t scr_el3 = read_ctx_reg(state, CTX_SCR_EL3);
-#endif
+	u_register_t scr_el3 = read_ctx_reg(get_el3state_ctx(ctx), CTX_SCR_EL3);
 	assert(((scr_el3 & SCR_HCE_BIT) != 0UL) &&
 			(el_implemented(2U) != EL_IMPL_NONE));
-
-	/*
-	 * Currently some extensions are configured using
-	 * direct register updates. Therefore, do this here
-	 * instead of when setting up context.
-	 */
-	manage_extensions_nonsecure_mixed(0, ctx);
+#endif /* ENABLE_ASSERTIONS */
 
 	/*
 	 * Set the NS bit to be able to access the ICC_SRE_EL2
diff --git a/lib/extensions/amu/aarch64/amu.c b/lib/extensions/amu/aarch64/amu.c
index c650629..e0d3399 100644
--- a/lib/extensions/amu/aarch64/amu.c
+++ b/lib/extensions/amu/aarch64/amu.c
@@ -188,95 +188,66 @@
  * Enable counters. This function is meant to be invoked by the context
  * management library before exiting from EL3.
  */
-void amu_enable(bool el2_unused, cpu_context_t *ctx)
+void amu_enable(cpu_context_t *ctx)
 {
-	uint64_t amcfgr_el0_ncg;		/* Number of counter groups */
-	uint64_t amcgcr_el0_cg0nc;		/* Number of group 0 counters */
-
-	uint64_t amcntenset0_el0_px = 0x0;	/* Group 0 enable mask */
-	uint64_t amcntenset1_el0_px = 0x0;	/* Group 1 enable mask */
-
-	if (el2_unused) {
-		/*
-		 * CPTR_EL2.TAM: Set to zero so any accesses to the Activity
-		 * Monitor registers do not trap to EL2.
-		 */
-		write_cptr_el2_tam(0U);
-	}
-
 	/*
-	 * Retrieve and update the CPTR_EL3 value from the context mentioned
-	 * in 'ctx'. Set CPTR_EL3.TAM to zero so that any accesses to
-	 * the Activity Monitor registers do not trap to EL3.
+	 * Set CPTR_EL3.TAM to zero so that any accesses to the Activity Monitor
+	 * registers do not trap to EL3.
 	 */
 	ctx_write_cptr_el3_tam(ctx, 0U);
 
-	/*
-	 * Retrieve the number of architected counters. All of these counters
-	 * are enabled by default.
-	 */
-
-	amcgcr_el0_cg0nc = read_amcgcr_el0_cg0nc();
-	amcntenset0_el0_px = (UINT64_C(1) << (amcgcr_el0_cg0nc)) - 1U;
+	/* Initialize FEAT_AMUv1p1 features if present. */
+	if (is_feat_amuv1p1_supported()) {
+		/*
+		 * Set SCR_EL3.AMVOFFEN to one so that accesses to virtual
+		 * offset registers at EL2 do not trap to EL3
+		 */
+		ctx_write_scr_el3_amvoffen(ctx, 1U);
+	}
+}
 
-	assert(amcgcr_el0_cg0nc <= AMU_AMCGCR_CG0NC_MAX);
+void amu_init_el3(void)
+{
+	uint64_t group0_impl_ctr = read_amcgcr_el0_cg0nc();
+	uint64_t group0_en_mask = (1 << (group0_impl_ctr)) - 1U;
+	uint64_t num_ctr_groups = read_amcfgr_el0_ncg();
 
-	/*
-	 * The platform may opt to enable specific auxiliary counters. This can
-	 * be done via the common FCONF getter, or via the platform-implemented
-	 * function.
-	 */
+	/* Enable all architected counters by default */
+	write_amcntenset0_el0_px(group0_en_mask);
 
 #if ENABLE_AMU_AUXILIARY_COUNTERS
-	const struct amu_topology *topology;
+	if (num_ctr_groups > 0U) {
+		uint64_t amcntenset1_el0_px = 0x0; /* Group 1 enable mask */
+		const struct amu_topology *topology;
 
+		/*
+		 * The platform may opt to enable specific auxiliary counters.
+		 * This can be done via the common FCONF getter, or via the
+		 * platform-implemented function.
+		 */
 #if ENABLE_AMU_FCONF
-	topology = FCONF_GET_PROPERTY(amu, config, topology);
+		topology = FCONF_GET_PROPERTY(amu, config, topology);
 #else
-	topology = plat_amu_topology();
+		topology = plat_amu_topology();
 #endif /* ENABLE_AMU_FCONF */
 
-	if (topology != NULL) {
-		unsigned int core_pos = plat_my_core_pos();
+		if (topology != NULL) {
+			unsigned int core_pos = plat_my_core_pos();
 
-		amcntenset1_el0_px = topology->cores[core_pos].enable;
-	} else {
-		ERROR("AMU: failed to generate AMU topology\n");
-	}
-#endif /* ENABLE_AMU_AUXILIARY_COUNTERS */
-
-	/*
-	 * Enable the requested counters.
-	 */
-
-	write_amcntenset0_el0_px(amcntenset0_el0_px);
+			amcntenset1_el0_px = topology->cores[core_pos].enable;
+		} else {
+			ERROR("AMU: failed to generate AMU topology\n");
+		}
 
-	amcfgr_el0_ncg = read_amcfgr_el0_ncg();
-	if (amcfgr_el0_ncg > 0U) {
 		write_amcntenset1_el0_px(amcntenset1_el0_px);
-
-#if !ENABLE_AMU_AUXILIARY_COUNTERS
+	}
+#else /* ENABLE_AMU_AUXILIARY_COUNTERS */
+	if (num_ctr_groups > 0U) {
 		VERBOSE("AMU: auxiliary counters detected but support is disabled\n");
-#endif
 	}
+#endif /* ENABLE_AMU_AUXILIARY_COUNTERS */
 
-	/* Initialize FEAT_AMUv1p1 features if present. */
 	if (is_feat_amuv1p1_supported()) {
-		if (el2_unused) {
-			/*
-			 * Make sure virtual offsets are disabled if EL2 not
-			 * used.
-			 */
-			write_hcr_el2_amvoffen(0U);
-		} else {
-			/*
-			 * Virtual offset registers are only accessible from EL3
-			 * and EL2, when clear, this bit traps accesses from EL2
-			 * so we set it to 1 when EL2 is present.
-			 */
-			ctx_write_scr_el3_amvoffen(ctx, 1U);
-		}
-
 #if AMU_RESTRICT_COUNTERS
 		/*
 		 * FEAT_AMUv1p1 adds a register field to restrict access to
@@ -297,6 +268,21 @@
 #endif
 }
 
+void amu_init_el2_unused(void)
+{
+	/*
+	 * CPTR_EL2.TAM: Set to zero so any accesses to the Activity Monitor
+	 *  registers do not trap to EL2.
+	 */
+	write_cptr_el2_tam(0U);
+
+	/* Initialize FEAT_AMUv1p1 features if present. */
+	if (is_feat_amuv1p1_supported()) {
+		/* Make sure virtual offsets are disabled if EL2 not used. */
+		write_hcr_el2_amvoffen(0U);
+	}
+}
+
 /* Read the group 0 counter identified by the given `idx`. */
 static uint64_t amu_group0_cnt_read(unsigned int idx)
 {
@@ -526,10 +512,10 @@
 
 	uint64_t hcr_el2_amvoffen = 0;	/* AMU virtual offsets enabled */
 
-	uint64_t amcfgr_el0_ncg;	/* Number of counter groups */
 	uint64_t amcgcr_el0_cg0nc;	/* Number of group 0 counters */
 
 #if ENABLE_AMU_AUXILIARY_COUNTERS
+	uint64_t amcfgr_el0_ncg;	/* Number of counter groups */
 	uint64_t amcgcr_el0_cg1nc;	/* Number of group 1 counters */
 	uint64_t amcg1idr_el0_voff;	/* Auxiliary counters with virtual offsets */
 #endif
@@ -541,7 +527,6 @@
 	core_pos = plat_my_core_pos();
 	ctx = &amu_ctxs_[core_pos];
 
-	amcfgr_el0_ncg = read_amcfgr_el0_ncg();
 	amcgcr_el0_cg0nc = read_amcgcr_el0_cg0nc();
 
 	if (is_feat_amuv1p1_supported()) {
@@ -549,22 +534,12 @@
 	}
 
 #if ENABLE_AMU_AUXILIARY_COUNTERS
+	amcfgr_el0_ncg = read_amcfgr_el0_ncg();
 	amcgcr_el0_cg1nc = (amcfgr_el0_ncg > 0U) ? read_amcgcr_el0_cg1nc() : 0U;
 	amcg1idr_el0_voff = (hcr_el2_amvoffen != 0U) ? read_amcg1idr_el0_voff() : 0U;
 #endif
 
 	/*
-	 * Sanity check that all counters were disabled when the context was
-	 * previously saved.
-	 */
-
-	assert(read_amcntenset0_el0_px() == 0U);
-
-	if (amcfgr_el0_ncg > 0U) {
-		assert(read_amcntenset1_el0_px() == 0U);
-	}
-
-	/*
 	 * Restore the counter values from the local context.
 	 */