fix(sme): enable SME/SME2 during arch init

This change enables SME/SME2 for nonsecure use at EL2 for TFTF cases
during arch_setup. This removes dependency on testcases to explicitly
call sme_enable or sme2_enable to access SME or SME2 functionality.

This change also adds CPTR_EL2 register in suspend context. CPTR_EL2
register is saved/restored in CPU suspend entry/exit path.

Signed-off-by: Arunachalam Ganapathy <arunachalam.ganapathy@arm.com>
Change-Id: I2c99fd49c48c1a9ff2110747714db858a78d3a32
diff --git a/include/lib/aarch64/arch.h b/include/lib/aarch64/arch.h
index 739cd26..d607029 100644
--- a/include/lib/aarch64/arch.h
+++ b/include/lib/aarch64/arch.h
@@ -575,21 +575,10 @@
 #define EVNTI_SHIFT		U(4)
 #define EVNTI_MASK		U(0xf)
 
-/* CPTR_EL3 definitions */
-#define TCPAC_BIT		(U(1) << 31)
-#define TAM_BIT			(U(1) << 30)
-#define TTA_BIT			(U(1) << 20)
-#define ESM_BIT			(U(1) << 12)
-#define TFP_BIT			(U(1) << 10)
-#define CPTR_EZ_BIT		(U(1) << 8)
-#define CPTR_EL3_RESET_VAL	U(0x0)
-
 /* CPTR_EL2 definitions */
-#define CPTR_EL2_RES1		((ULL(3) << 12) | (ULL(1) << 9) | (ULL(0xff)))
+#define CPTR_EL2_RES1		((ULL(1) << 13) | (ULL(1) << 9) | (ULL(0xff)))
 #define CPTR_EL2_TCPAC_BIT	(ULL(1) << 31)
 #define CPTR_EL2_TAM_BIT	(ULL(1) << 30)
-#define CPTR_EL2_SMEN_MASK	ULL(0x3)
-#define CPTR_EL2_SMEN_SHIFT	U(24)
 #define CPTR_EL2_TTA_BIT	(ULL(1) << 20)
 #define CPTR_EL2_TSM_BIT	(ULL(1) << 12)
 #define CPTR_EL2_TFP_BIT	(ULL(1) << 10)
@@ -981,6 +970,7 @@
 #define SMCR_ELX_LEN_MASK		U(0x1ff)
 #define SMCR_ELX_EZT0_BIT		(U(1) << 30)
 #define SMCR_ELX_FA64_BIT		(U(1) << 31)
+#define SMCR_EL2_RESET_VAL		(SMCR_ELX_EZT0_BIT | SMCR_ELX_FA64_BIT)
 
 /*******************************************************************************
  * Definitions of MAIR encodings for device and normal memory
diff --git a/include/lib/extensions/sme.h b/include/lib/extensions/sme.h
index c89e630..b20a8c0 100644
--- a/include/lib/extensions/sme.h
+++ b/include/lib/extensions/sme.h
@@ -24,13 +24,9 @@
 } smestop_instruction_type_t;
 
 /* SME feature related prototypes. */
-void sme_enable(void);
 void sme_smstart(smestart_instruction_type_t smstart_type);
 void sme_smstop(smestop_instruction_type_t smstop_type);
 
-/* SME2 feature related prototypes. */
-void sme2_enable(void);
-
 /* Assembly function prototypes. */
 uint64_t sme_rdvl_1(void);
 void sme_try_illegal_instruction(void);
diff --git a/lib/extensions/sme/aarch64/sme.c b/lib/extensions/sme/aarch64/sme.c
index 632b2a7..b4eaa2f 100644
--- a/lib/extensions/sme/aarch64/sme.c
+++ b/lib/extensions/sme/aarch64/sme.c
@@ -14,32 +14,6 @@
 #include <lib/extensions/sme.h>
 
 /*
- * Function : sme_enable
- * Enable SME for nonsecure use at EL2 for TFTF cases.
- */
-void sme_enable(void)
-{
-	u_register_t reg;
-
-	/*
-	 * Make sure SME accesses don't cause traps by setting appropriate
-	 * fields in CPTR_EL2.
-	 */
-	reg = read_cptr_el2();
-	if ((read_hcr_el2() & HCR_E2H_BIT) == 0U) {
-		/* When HCR_EL2.E2H == 0, clear TSM bit in CPTR_EL2. */
-		reg = reg & ~CPTR_EL2_TSM_BIT;
-	} else {
-		/* When HCR_EL2.E2H == 1, set SMEN bits in CPTR_EL2. */
-		reg = reg | (CPTR_EL2_SMEN_MASK << CPTR_EL2_SMEN_SHIFT);
-	}
-	write_cptr_el2(reg);
-
-	isb();
-
-}
-
-/*
  * Function: sme_smstart
  * This function enables streaming mode and ZA array storage access
  * independently or together based on the type of instruction variant.
diff --git a/lib/extensions/sme/aarch64/sme2.c b/lib/extensions/sme/aarch64/sme2.c
deleted file mode 100644
index 1eeb2ac..0000000
--- a/lib/extensions/sme/aarch64/sme2.c
+++ /dev/null
@@ -1,40 +0,0 @@
-/*
- * Copyright (c) 2023, ARM Limited and Contributors. All rights reserved.
- *
- * SPDX-License-Identifier: BSD-3-Clause
- */
-
-#include <stdbool.h>
-#include <stdio.h>
-
-#include <arch.h>
-#include <arch_features.h>
-#include <arch_helpers.h>
-#include <lib/extensions/sme.h>
-#include <tftf_lib.h>
-
-/*
- * sme2_enable
- * Enable SME2 for nonsecure use at EL2 for TFTF cases.
- */
-void sme2_enable(void)
-{
-	u_register_t reg;
-
-	/* SME2 is an extended version of SME.
-	 * Therefore, SME accesses still must be taken care by setting
-	 * appropriate fields in order to avoid traps in CPTR_EL2.
-	 */
-	sme_enable();
-
-	/*
-	 * Make sure ZT0 register access don't cause traps by setting
-	 * appropriate field in SMCR_EL2 register.
-	 */
-	reg = read_smcr_el2();
-	reg |= SMCR_ELX_EZT0_BIT;
-	write_smcr_el2(reg);
-	isb();
-
-}
-
diff --git a/lib/power_management/suspend/aarch64/asm_tftf_suspend.S b/lib/power_management/suspend/aarch64/asm_tftf_suspend.S
index 09950b5..e715e49 100644
--- a/lib/power_management/suspend/aarch64/asm_tftf_suspend.S
+++ b/lib/power_management/suspend/aarch64/asm_tftf_suspend.S
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2019, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -72,7 +72,7 @@
 	ret
 
 2:	mrs	x1, mair_el2
-	mrs	x2, hcr_el2
+	mrs	x2, cptr_el2
 	mrs	x3, ttbr0_el2
 	mrs	x4, tcr_el2
 	mrs	x5, vbar_el2
@@ -80,6 +80,8 @@
 	stp	x1, x2, [x0, #SUSPEND_CTX_MAIR_OFFSET]
 	stp	x3, x4, [x0, #SUSPEND_CTX_TTBR0_OFFSET]
 	stp	x5, x6, [x0, #SUSPEND_CTX_VBAR_OFFSET]
+	mrs	x1, hcr_el2
+	str	x1, [x0, #SUSPEND_CTX_HCR_OFFSET]
 	ret
 endfunc __tftf_save_arch_context
 
@@ -114,10 +116,13 @@
 	ldp	x3, x4, [x0, #SUSPEND_CTX_TTBR0_OFFSET]
 	ldp	x5, x6, [x0, #SUSPEND_CTX_VBAR_OFFSET]
 	msr	mair_el2, x1
-	msr	hcr_el2, x2
+	msr	cptr_el2, x2
 	msr	ttbr0_el2, x3
 	msr	tcr_el2, x4
 	msr	vbar_el2, x5
+	ldr	x1, [x0, #SUSPEND_CTX_HCR_OFFSET]
+	msr	hcr_el2, x1
+
 	/*
 	 * TLB invalidations need to be completed before enabling MMU
 	 */
diff --git a/lib/power_management/suspend/suspend_private.h b/lib/power_management/suspend/suspend_private.h
index bc2f3a6..2ccf7a3 100644
--- a/lib/power_management/suspend/suspend_private.h
+++ b/lib/power_management/suspend/suspend_private.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2019, Arm Limited. All rights reserved.
+ * Copyright (c) 2018-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -9,20 +9,22 @@
 
 /*
  * Number of system registers we need to save/restore across a CPU suspend:
- * MAIR, CPACR_EL1/HCR_EL2, TTBR0, TCR, VBAR, SCTLR,
+ * EL1: MAIR, CPACR, TTBR0, TCR, VBAR, SCTLR
+ * EL2: MAIR, CPTR, TTBR0, TCR, VBAR, SCTLR, HCR
  * APIAKeyLo_EL1 and APIAKeyHi_EL1 (if enabled).
  */
 #if ENABLE_PAUTH
-#define NR_CTX_REGS 8
+#define NR_CTX_REGS 10
 #else
-#define NR_CTX_REGS 6
+#define NR_CTX_REGS 8
 #endif
 
 /* Offsets of the fields in the context structure. Needed by asm code. */
 #define	SUSPEND_CTX_MAIR_OFFSET		0
 #define	SUSPEND_CTX_TTBR0_OFFSET	16
 #define	SUSPEND_CTX_VBAR_OFFSET		32
-#define	SUSPEND_CTX_APIAKEY_OFFSET	48
+#define	SUSPEND_CTX_HCR_OFFSET		48
+#define	SUSPEND_CTX_APIAKEY_OFFSET	64
 
 #define SUSPEND_CTX_SP_OFFSET (8 * NR_CTX_REGS)
 #define SUSPEND_CTX_SAVE_SYSTEM_CTX_OFFSET (SUSPEND_CTX_SP_OFFSET + 8)
diff --git a/tftf/framework/aarch64/arch.c b/tftf/framework/aarch64/arch.c
index 56369ae..0510678 100644
--- a/tftf/framework/aarch64/arch.c
+++ b/tftf/framework/aarch64/arch.c
@@ -5,6 +5,7 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
+#include <arch_features.h>
 #include <arch_helpers.h>
 
 void tftf_arch_setup(void)
@@ -23,11 +24,20 @@
 		write_hcr_el2(HCR_TGE_BIT);
 
 		/*
-		 * Disable trap of SVE instructions to EL2.
+		 * Disable trap of SVE, SME instructions to EL2.
 		 * The fields of the CPTR_EL2 register reset to an
 		 * architecturally UNKNOWN value.
 		 */
-		write_cptr_el2(CPTR_EL2_RES1);
+		write_cptr_el2(CPTR_EL2_RESET_VAL);
 		isb();
+
+		/*
+		 * Enable access to ZT0 storage when FEAT_SME2 is implemented
+		 * and enable FA64 when FEAT_SME_FA64 is implemented
+		 */
+		if (is_feat_sme_supported()) {
+			write_smcr_el2(SMCR_EL2_RESET_VAL);
+			isb();
+		}
 	}
 }
diff --git a/tftf/framework/framework.mk b/tftf/framework/framework.mk
index f57572c..b25d1d0 100644
--- a/tftf/framework/framework.mk
+++ b/tftf/framework/framework.mk
@@ -87,7 +87,6 @@
 	lib/extensions/pauth/aarch64/pauth.c				\
 	lib/extensions/pauth/aarch64/pauth_helpers.S			\
 	lib/extensions/sme/aarch64/sme.c				\
-	lib/extensions/sme/aarch64/sme2.c				\
 	lib/extensions/sme/aarch64/sme_helpers.S			\
 	lib/extensions/sme/aarch64/sme2_helpers.S			\
 	lib/extensions/sve/aarch64/sve.c				\
diff --git a/tftf/tests/extensions/sme/test_sme.c b/tftf/tests/extensions/sme/test_sme.c
index 64391ab..6d85183 100644
--- a/tftf/tests/extensions/sme/test_sme.c
+++ b/tftf/tests/extensions/sme/test_sme.c
@@ -71,13 +71,11 @@
 	unsigned int requested_vector_len;
 	unsigned int len_max;
 	unsigned int __unused svl_max = 0U;
+	u_register_t saved_smcr;
 
 	/* Skip the test if SME is not supported. */
 	SKIP_TEST_IF_SME_NOT_SUPPORTED();
 
-	/* Enable SME for use at NS EL2. */
-	sme_enable();
-
 	/* Make sure TPIDR2_EL0 is accessible. */
 	write_tpidr2_el0(0);
 	if (read_tpidr2_el0() != 0) {
@@ -97,8 +95,12 @@
 	/* Entering Streaming SVE mode */
 	sme_smstart(SMSTART_SM);
 
+	saved_smcr = read_smcr_el2();
+
 	/* Write SMCR_EL2 with the LEN max to find implemented width. */
 	write_smcr_el2(SME_SMCR_LEN_MAX);
+	isb();
+
 	len_max = (unsigned int)read_smcr_el2();
 	VERBOSE("Maximum SMCR_EL2.LEN value: 0x%x\n", len_max);
 	VERBOSE("Enumerating supported vector lengths...\n");
@@ -108,6 +110,7 @@
 		reg &= ~(SMCR_ELX_LEN_MASK << SMCR_ELX_LEN_SHIFT);
 		reg |= (i << SMCR_ELX_LEN_SHIFT);
 		write_smcr_el2(reg);
+		isb();
 
 		/* Compute current and requested vector lengths in bits. */
 		current_vector_len = ((unsigned int)sme_rdvl_1() * 8U);
@@ -164,6 +167,9 @@
 		sme_try_illegal_instruction();
 	}
 
+	write_smcr_el2(saved_smcr);
+	isb();
+
 	return TEST_RESULT_SUCCESS;
 #endif /* __aarch64__ */
 }
diff --git a/tftf/tests/extensions/sme/test_sme2.c b/tftf/tests/extensions/sme/test_sme2.c
index c16b7c3..e82da08 100644
--- a/tftf/tests/extensions/sme/test_sme2.c
+++ b/tftf/tests/extensions/sme/test_sme2.c
@@ -55,18 +55,16 @@
 	/* Skip the test if SME2 is not supported. */
 	SKIP_TEST_IF_SME2_NOT_SUPPORTED();
 
-	/* Enable SME2 for use at NS EL2. */
-	sme2_enable();
-
 	/*
 	 * FEAT_SME2 adds a 512 BIT architectural register ZT0 to support
 	 * the lookup-table feature.
 	 * System register SMCR_ELx defines a bit SMCR_ELx.EZT0 bit [30] to
-	 * enable/disable access to this register.
+	 * enable/disable access to this register. SMCR_EL2_RESET_VAL enables
+	 * this bit by default.
 	 *
 	 * Instructions to access ZT0 register are being tested to ensure
-	 * SMCR_ELx.EZT0 bit is set at( EL-3 as well as EL-2), so that
-	 * they are not trapped.
+	 * SMCR_EL3.EZT0 bit is set by EL3 firmware so that EL2 access are not
+	 * trapped.
 	 */
 
 	/* Make sure we can acesss SME2 ZT0 storage, PSTATE.ZA = 1*/