fix(mte): remove CTX_INCLUDE_MTE_REGS usage

commit@0a33adc058080433f73bde73895266068990245c
Deprecated CTX_INCLUDE_MTE_REGS but missed its usage in
context save and restore path.

Change-Id: I30544abdff2cf92ff05d2d4df46ffc6ff10611de
Signed-off-by: Govindraj Raja <govindraj.raja@arm.com>
diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 389c086..176308e 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2013-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2013-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -107,7 +107,12 @@
 #endif /* NS_TIMER_SWITCH */
 
 	/* Save MTE system registers if the build has instructed so */
-#if CTX_INCLUDE_MTE_REGS
+#if ENABLE_FEAT_MTE
+#if ENABLE_FEAT_MTE == 2
+	mrs x8, id_aa64pfr1_el1
+	and x8, x8, #(ID_AA64PFR1_EL1_MTE_MASK << ID_AA64PFR1_EL1_MTE_SHIFT)
+	cbz x8, no_mte_save
+#endif
 	mrs	x15, TFSRE0_EL1
 	mrs	x16, TFSR_EL1
 	stp	x15, x16, [x0, #CTX_TFSRE0_EL1]
@@ -115,7 +120,9 @@
 	mrs	x9, RGSR_EL1
 	mrs	x10, GCR_EL1
 	stp	x9, x10, [x0, #CTX_RGSR_EL1]
-#endif /* CTX_INCLUDE_MTE_REGS */
+
+no_mte_save:
+#endif /* ENABLE_FEAT_MTE */
 
 	ret
 endfunc el1_sysregs_context_save
@@ -206,7 +213,13 @@
 #endif /* NS_TIMER_SWITCH */
 
 	/* Restore MTE system registers if the build has instructed so */
-#if CTX_INCLUDE_MTE_REGS
+#if ENABLE_FEAT_MTE
+#if ENABLE_FEAT_MTE == 2
+	mrs x8, id_aa64pfr1_el1
+	and x8, x8, #(ID_AA64PFR1_EL1_MTE_MASK << ID_AA64PFR1_EL1_MTE_SHIFT)
+	cbz x8, no_mte_restore
+#endif
+
 	ldp	x11, x12, [x0, #CTX_TFSRE0_EL1]
 	msr	TFSRE0_EL1, x11
 	msr	TFSR_EL1, x12
@@ -214,7 +227,9 @@
 	ldp	x13, x14, [x0, #CTX_RGSR_EL1]
 	msr	RGSR_EL1, x13
 	msr	GCR_EL1, x14
-#endif /* CTX_INCLUDE_MTE_REGS */
+
+no_mte_restore:
+#endif /* ENABLE_FEAT_MTE */
 
 	/* No explict ISB required here as ERET covers it */
 	ret
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 780537d..da53194 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -1333,9 +1333,10 @@
 	el2_sysregs_ctx = get_el2_sysregs_ctx(ctx);
 
 	el2_sysregs_context_restore_common(el2_sysregs_ctx);
-#if CTX_INCLUDE_MTE_REGS
-	write_tfsr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_TFSR_EL2));
-#endif
+
+	if (is_feat_mte_supported()) {
+		write_tfsr_el2(read_ctx_reg(el2_sysregs_ctx, CTX_TFSR_EL2));
+	}
 
 #if CTX_INCLUDE_MPAM_REGS
 	if (is_feat_mpam_supported()) {