fix(el3_runtime): fix SVE and AMU extension enablement flags

If SVE are enabled for both Non-secure and Secure world along with AMU
extension, then it causes the TAM_BIT in CPTR_EL3 to be set upon exit
from bl31. This restricts access to the AMU register set in normal
world. This fix maintains consistency in both TAM_BIT and CPTR_EZ_BIT
by saving and restoring CPTR_EL3 register from EL3 context.

Signed-off-by: Arunachalam Ganapathy <arunachalam.ganapathy@arm.com>
Change-Id: Id76ce1d27ee48bed65eb32392036377716aff087
diff --git a/include/lib/extensions/amu.h b/include/lib/extensions/amu.h
index 3a70e4f..3a254c9 100644
--- a/include/lib/extensions/amu.h
+++ b/include/lib/extensions/amu.h
@@ -13,6 +13,7 @@
 #include <lib/cassert.h>
 #include <lib/utils_def.h>
 
+#include <context.h>
 #include <platform_def.h>
 
 /* All group 0 counters */
@@ -80,7 +81,11 @@
 };
 
 unsigned int amu_get_version(void);
+#if __aarch64__
+void amu_enable(bool el2_unused, cpu_context_t *ctx);
+#else
 void amu_enable(bool el2_unused);
+#endif
 
 /* Group 0 configuration helpers */
 uint64_t amu_group0_cnt_read(unsigned int idx);
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index d610fd4..40e7ddf 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -903,16 +903,11 @@
 
 #if IMAGE_BL31
 	/* ----------------------------------------------------------
-	 * Restore CPTR_EL3, ZCR_EL3 for SVE support.
-	 * If SVE is not supported - skip the restoration.
+	 * Restore CPTR_EL3.
 	 * ZCR is only restored if SVE is supported and enabled.
 	 * Synchronization is required before zcr_el3 is addressed.
 	 * ----------------------------------------------------------
 	 */
-	mrs	x17, id_aa64pfr0_el1
-	ubfx	x17, x17, ID_AA64PFR0_SVE_SHIFT, ID_AA64PFR0_SVE_LENGTH
-	cbz	x17, sve_not_enabled
-
 	ldp	x19, x20, [sp, #CTX_EL3STATE_OFFSET + CTX_CPTR_EL3]
 	msr	cptr_el3, x19
 
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 7a25151..7c6f953 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -25,6 +25,7 @@
 #include <lib/extensions/twed.h>
 #include <lib/utils.h>
 
+static void enable_extensions_secure(cpu_context_t *ctx);
 
 /*******************************************************************************
  * Context management library initialisation routine. This library is used by
@@ -178,19 +179,13 @@
 	 *  indicated by the interrupt routing model for BL31.
 	 */
 	scr_el3 |= get_scr_el3_from_routing_model(security_state);
-
-#if ENABLE_SVE_FOR_NS
-	if (security_state == NON_SECURE) {
-		sve_enable(ctx);
-	}
 #endif
-#if ENABLE_SVE_FOR_SWD
+
+	/* Save the initialized value of CPTR_EL3 register */
+	write_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3, read_cptr_el3());
 	if (security_state == SECURE) {
-		sve_enable(ctx);
+		enable_extensions_secure(ctx);
 	}
-#endif
-
-#endif
 
 	/*
 	 * SCR_EL3.HCE: Enable HVC instructions if next execution state is
@@ -335,7 +330,7 @@
  * When EL2 is implemented but unused `el2_unused` is non-zero, otherwise
  * it is zero.
  ******************************************************************************/
-static void enable_extensions_nonsecure(bool el2_unused)
+static void enable_extensions_nonsecure(bool el2_unused, cpu_context_t *ctx)
 {
 #if IMAGE_BL31
 #if ENABLE_SPE_FOR_LOWER_ELS
@@ -343,9 +338,13 @@
 #endif
 
 #if ENABLE_AMU
-	amu_enable(el2_unused);
+	amu_enable(el2_unused, ctx);
 #endif
 
+#if ENABLE_SVE_FOR_NS
+	sve_enable(ctx);
+#endif
+
 #if ENABLE_MPAM_FOR_LOWER_ELS
 	mpam_enable(el2_unused);
 #endif
@@ -353,6 +352,18 @@
 }
 
 /*******************************************************************************
+ * Enable architecture extensions on first entry to Secure world.
+ ******************************************************************************/
+static void enable_extensions_secure(cpu_context_t *ctx)
+{
+#if IMAGE_BL31
+#if ENABLE_SVE_FOR_SWD
+	sve_enable(ctx);
+#endif
+#endif
+}
+
+/*******************************************************************************
  * The following function initializes the cpu_context for a CPU specified by
  * its `cpu_idx` for first use, and sets the initial entrypoint state as
  * specified by the entry_point_info structure.
@@ -586,7 +597,7 @@
 			write_cnthp_ctl_el2(CNTHP_CTL_RESET_VAL &
 						~(CNTHP_CTL_ENABLE_BIT));
 		}
-		enable_extensions_nonsecure(el2_unused);
+		enable_extensions_nonsecure(el2_unused, ctx);
 	}
 
 	cm_el1_sysregs_context_restore(security_state);
diff --git a/lib/extensions/amu/aarch64/amu.c b/lib/extensions/amu/aarch64/amu.c
index 24c3737..295c0d5 100644
--- a/lib/extensions/amu/aarch64/amu.c
+++ b/lib/extensions/amu/aarch64/amu.c
@@ -46,7 +46,7 @@
  * Enable counters. This function is meant to be invoked
  * by the context management library before exiting from EL3.
  */
-void amu_enable(bool el2_unused)
+void amu_enable(bool el2_unused, cpu_context_t *ctx)
 {
 	uint64_t v;
 	unsigned int amu_version = amu_get_version();
@@ -88,12 +88,13 @@
 	}
 
 	/*
-	 * CPTR_EL3.TAM: Set to zero so that any accesses to
+	 * 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.
 	 */
-	v = read_cptr_el3();
+	v = read_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3);
 	v &= ~TAM_BIT;
-	write_cptr_el3(v);
+	write_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3, v);
 
 	/* Enable group 0 counters */
 	write_amcntenset0_el0(AMU_GROUP0_COUNTERS_MASK);
diff --git a/lib/extensions/sve/sve.c b/lib/extensions/sve/sve.c
index 7043cc2..2702c30 100644
--- a/lib/extensions/sve/sve.c
+++ b/lib/extensions/sve/sve.c
@@ -27,11 +27,13 @@
 
 void sve_enable(cpu_context_t *context)
 {
+	u_register_t cptr_el3;
+
 	if (!sve_supported()) {
 		return;
 	}
 
-	u_register_t cptr_el3 = read_cptr_el3();
+	cptr_el3 = read_ctx_reg(get_el3state_ctx(context), CTX_CPTR_EL3);
 
 	/* Enable access to SVE functionality for all ELs. */
 	cptr_el3 = (cptr_el3 | CPTR_EZ_BIT) & ~(TFP_BIT);