feat(cm): handle asymmetry for FEAT_TCR2

With introduction of FEAT_STATE_CHECK_ASYMMETRIC, the asymmetry of cores
can be handled. FEAT_TCR2 is one of the features which can be
asymmetric across cores and the respective support is added here.

Adding a function to handle this asymmetry by re-visting the
feature presence on running core.
There are two possible cases:
 - If the primary core has the feature and secondary does not have it
   then the feature is disabled.
 - If the primary does not have the feature and secondary has it then,
   the feature need to be enabled in secondary cores.

Signed-off-by: Jayanth Dodderi Chidanand <jayanthdodderi.chidanand@arm.com>
Change-Id: I73a70891d52268ddfa4effe40edf04115f5821ca
diff --git a/bl31/bl31.mk b/bl31/bl31.mk
index 7dc71a2..336ad2b 100644
--- a/bl31/bl31.mk
+++ b/bl31/bl31.mk
@@ -111,6 +111,10 @@
 BL31_SOURCES		+=	lib/extensions/fgt/fgt2.c
 endif
 
+ifneq (${ENABLE_FEAT_TCR2},0)
+BL31_SOURCES		+=	lib/extensions/tcr/tcr2.c
+endif
+
 ifeq (${ENABLE_MPMM},1)
 BL31_SOURCES		+=	${MPMM_SOURCES}
 endif
diff --git a/include/lib/extensions/tcr2.h b/include/lib/extensions/tcr2.h
new file mode 100644
index 0000000..08a2b08
--- /dev/null
+++ b/include/lib/extensions/tcr2.h
@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2024, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#ifndef TCR2_H
+#define TCR2_H
+
+#include <context.h>
+
+#if ENABLE_FEAT_TCR2
+void tcr2_enable(cpu_context_t *ctx);
+void tcr2_disable(cpu_context_t *ctx);
+#else
+static inline void tcr2_enable(cpu_context_t *ctx)
+{
+}
+static inline void tcr2_disable(cpu_context_t *ctx)
+{
+}
+#endif /* ENABLE_FEAT_TCR2 */
+
+#endif /* TCR2_H */
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 6f3b51a..218ad11 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -34,6 +34,7 @@
 #include <lib/extensions/spe.h>
 #include <lib/extensions/sve.h>
 #include <lib/extensions/sys_reg_trace.h>
+#include <lib/extensions/tcr2.h>
 #include <lib/extensions/trbe.h>
 #include <lib/extensions/trf.h>
 #include <lib/utils.h>
@@ -1538,28 +1539,37 @@
 *********************************************************************************/
 void cm_handle_asymmetric_features(void)
 {
-#if ENABLE_SPE_FOR_NS == FEAT_STATE_CHECK_ASYMMETRIC
-	cpu_context_t *spe_ctx = cm_get_context(NON_SECURE);
+	cpu_context_t *ctx __maybe_unused = cm_get_context(NON_SECURE);
 
-	assert(spe_ctx != NULL);
+	assert(ctx != NULL);
 
+#if ENABLE_SPE_FOR_NS == FEAT_STATE_CHECK_ASYMMETRIC
 	if (is_feat_spe_supported()) {
-		spe_enable(spe_ctx);
+		spe_enable(ctx);
 	} else {
-		spe_disable(spe_ctx);
+		spe_disable(ctx);
 	}
 #endif
-#if ERRATA_A520_2938996 || ERRATA_X4_2726228
-	cpu_context_t *trbe_ctx = cm_get_context(NON_SECURE);
-
-	assert(trbe_ctx != NULL);
 
+#if ERRATA_A520_2938996 || ERRATA_X4_2726228
 	if (check_if_affected_core() == ERRATA_APPLIES) {
 		if (is_feat_trbe_supported()) {
-			trbe_disable(trbe_ctx);
+			trbe_disable(ctx);
 		}
 	}
 #endif
+
+#if ENABLE_FEAT_TCR2 == FEAT_STATE_CHECK_ASYMMETRIC
+	el3_state_t *el3_state = get_el3state_ctx(ctx);
+	u_register_t spsr = read_ctx_reg(el3_state, CTX_SPSR_EL3);
+
+	if (is_feat_tcr2_supported() && (GET_RW(spsr) == MODE_RW_64)) {
+		tcr2_enable(ctx);
+	} else {
+		tcr2_disable(ctx);
+	}
+#endif
+
 }
 #endif
 
diff --git a/lib/extensions/tcr/tcr2.c b/lib/extensions/tcr/tcr2.c
new file mode 100644
index 0000000..70bc5f8
--- /dev/null
+++ b/lib/extensions/tcr/tcr2.c
@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) 2024, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <arch.h>
+#include <arch_features.h>
+#include <arch_helpers.h>
+#include <lib/extensions/tcr2.h>
+
+void tcr2_enable(cpu_context_t *ctx)
+{
+	u_register_t reg;
+	el3_state_t *state;
+
+	state = get_el3state_ctx(ctx);
+
+	/* Set the TCR2EN bit in SCR_EL3 to enable access to TCR2_EL1,
+	 * and TCR2_EL2 registers .
+	 */
+
+	reg = read_ctx_reg(state, CTX_SCR_EL3);
+	reg |= SCR_TCR2EN_BIT;
+	write_ctx_reg(state, CTX_SCR_EL3, reg);
+}
+
+void tcr2_disable(cpu_context_t *ctx)
+{
+	u_register_t reg;
+	el3_state_t *state;
+
+	state = get_el3state_ctx(ctx);
+
+	/* Clear the TCR2EN bit in SCR_EL3 to disable access to TCR2_EL1,
+	 * and TCR2_EL2 registers .
+	 */
+
+	reg = read_ctx_reg(state, CTX_SCR_EL3);
+	reg &= ~SCR_TCR2EN_BIT;
+	write_ctx_reg(state, CTX_SCR_EL3, reg);
+}