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.
-	 */
+	/* 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);
+	}
+}
 
-	amcgcr_el0_cg0nc = read_amcgcr_el0_cg0nc();
-	amcntenset0_el0_px = (UINT64_C(1) << (amcgcr_el0_cg0nc)) - 1U;
+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();
 
-	assert(amcgcr_el0_cg0nc <= AMU_AMCGCR_CG0NC_MAX);
-
-	/*
-	 * 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");
+			amcntenset1_el0_px = topology->cores[core_pos].enable;
+		} else {
+			ERROR("AMU: failed to generate AMU topology\n");
+		}
+
+		write_amcntenset1_el0_px(amcntenset1_el0_px);
+	}
+#else /* ENABLE_AMU_AUXILIARY_COUNTERS */
+	if (num_ctr_groups > 0U) {
+		VERBOSE("AMU: auxiliary counters detected but support is disabled\n");
 	}
 #endif /* ENABLE_AMU_AUXILIARY_COUNTERS */
 
-	/*
-	 * Enable the requested counters.
-	 */
-
-	write_amcntenset0_el0_px(amcntenset0_el0_px);
-
-	amcfgr_el0_ncg = read_amcfgr_el0_ncg();
-	if (amcfgr_el0_ncg > 0U) {
-		write_amcntenset1_el0_px(amcntenset1_el0_px);
-
-#if !ENABLE_AMU_AUXILIARY_COUNTERS
-		VERBOSE("AMU: auxiliary counters detected but support is disabled\n");
-#endif
-	}
-
-	/* 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.
 	 */