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

1 files changed, 15 insertions, 31 deletions

diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index b7d014a09..d9ff0b6e6 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -496,22 +496,6 @@ void cm_setup_context(cpu_context_t *ctx, const entry_point_info_t *ep)
 }
 
 /*******************************************************************************
- * 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 @@ void cm_manage_extensions_el3(void)
 		spe_init_el3();
 	}
 
+	if (is_feat_amu_supported()) {
+		amu_init_el3();
+	}
+
 	if (is_feat_sme_supported()) {
 		sme_init_el3();
 	}
@@ -553,6 +541,10 @@ void cm_manage_extensions_el3(void)
 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 @@ static void manage_extensions_nonsecure_el2_unused(void)
 		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 @@ void cm_prepare_el3_exit(uint32_t security_state)
 {
 	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 @@ void cm_prepare_el3_exit(uint32_t security_state)
 #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 @@ void cm_prepare_el3_exit(uint32_t security_state)
 
 			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_el2_sysregs_context_restore(uint32_t security_state)
 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