feat(sme): add basic SME2 tests

FEAT_SME2 introduces an architectural register ZT0 to support
lookup table feature. This patch ensures that EL3 has
properly enabled the SME2 for use at lower exception levels,
thereby disabling the traps execution at lower exception levels,
when instructions access ZT0 register to EL3.

Signed-off-by: Jayanth Dodderi Chidanand <jayanthdodderi.chidanand@arm.com>
Change-Id: I46d51184b74c1e82c88344530601f2a3c1aee8ea
diff --git a/include/common/test_helpers.h b/include/common/test_helpers.h
index 04b8625..fdc5401 100644
--- a/include/common/test_helpers.h
+++ b/include/common/test_helpers.h
@@ -350,6 +350,14 @@
 		}								\
 	} while (false)
 
+#define SKIP_TEST_IF_SME2_NOT_SUPPORTED()					\
+	do {									\
+		if(!is_feat_sme2_supported()) {					\
+			tftf_testcase_printf("FEAT_SME2 not supported\n");	\
+			return TEST_RESULT_SKIPPED;				\
+		}								\
+	} while (false)
+
 /* Helper macro to verify if system suspend API is supported */
 #define is_psci_sys_susp_supported()	\
 		(tftf_get_psci_feature_info(SMC_PSCI_SYSTEM_SUSPEND)		\
diff --git a/include/lib/aarch64/arch.h b/include/lib/aarch64/arch.h
index 2f4b0c9..1cbe3c3 100644
--- a/include/lib/aarch64/arch.h
+++ b/include/lib/aarch64/arch.h
@@ -332,6 +332,7 @@
 #define ID_AA64PFR1_EL1_SME_MASK		ULL(0xf)
 #define ID_AA64PFR1_EL1_SME_NOT_SUPPORTED	ULL(0x0)
 #define ID_AA64PFR1_EL1_SME_SUPPORTED		ULL(0x1)
+#define ID_AA64PFR1_EL1_SME2_SUPPORTED		ULL(0x2)
 
 /* ID_PFR1_EL1 definitions */
 #define ID_PFR1_VIRTEXT_SHIFT	U(12)
@@ -887,6 +888,7 @@
 /* SMCR_ELx definitions */
 #define SMCR_ELX_LEN_SHIFT		U(0)
 #define SMCR_ELX_LEN_MASK		U(0x1ff)
+#define SMCR_ELX_EZT0_BIT		(U(1) << 30)
 #define SMCR_ELX_FA64_BIT		(U(1) << 31)
 
 /*******************************************************************************
diff --git a/include/lib/aarch64/arch_features.h b/include/lib/aarch64/arch_features.h
index bb1d156..203d140 100644
--- a/include/lib/aarch64/arch_features.h
+++ b/include/lib/aarch64/arch_features.h
@@ -232,4 +232,12 @@
 	return (features & ID_AA64SMFR0_EL1_FA64_BIT) != 0U;
 }
 
+static inline bool is_feat_sme2_supported(void)
+{
+	uint64_t features;
+
+	features = read_id_aa64pfr1_el1() >> ID_AA64PFR1_EL1_SME_SHIFT;
+	return (features & ID_AA64PFR1_EL1_SME_MASK) >= ID_AA64PFR1_EL1_SME2_SUPPORTED;
+}
+
 #endif /* ARCH_FEATURES_H */
diff --git a/include/lib/extensions/sme.h b/include/lib/extensions/sme.h
index f443cea..c89e630 100644
--- a/include/lib/extensions/sme.h
+++ b/include/lib/extensions/sme.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -28,10 +28,15 @@
 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);
 void sme_vector_to_ZA(const uint64_t *input_vector);
 void sme_ZA_to_vector(const uint64_t *output_vector);
+void sme2_load_zt0_instruction(const uint64_t *inputbuf);
+void sme2_store_zt0_instruction(const uint64_t *outputbuf);
 
 #endif /* SME_H */
diff --git a/lib/extensions/sme/aarch64/sme2.c b/lib/extensions/sme/aarch64/sme2.c
new file mode 100644
index 0000000..1eeb2ac
--- /dev/null
+++ b/lib/extensions/sme/aarch64/sme2.c
@@ -0,0 +1,40 @@
+/*
+ * 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/extensions/sme/aarch64/sme2_helpers.S b/lib/extensions/sme/aarch64/sme2_helpers.S
new file mode 100644
index 0000000..e5925e0
--- /dev/null
+++ b/lib/extensions/sme/aarch64/sme2_helpers.S
@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) 2023, ARM Limited and Contributors. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <asm_macros.S>
+#include <arch.h>
+
+	.arch	armv8-a+sve
+	.globl  sme2_load_zt0_instruction
+	.globl  sme2_store_zt0_instruction
+
+/*
+ * TODO: Due to the limitation with toolchain, SME2 intrinsics, still not
+ * being supported, manually instructions are encoded using the opcodes.
+ * Further, when the toolchain supports the requirements, these macros could
+ * be refactored.
+ */
+.macro _check_general_reg_number nr
+	.if ((\nr) < 0) || ((\nr) > 30)
+		.error "Bad register number \nr."
+	.endif
+.endm
+
+/*
+ * LDR (ZT0)  : Load ZT0 register with 64byte data.
+ * Instruction: LDR ZT0, [<Xn|SP>]
+ *
+ * LDR ZT0, nx
+ * Opcode bit field:
+ * nx    : 64-bit name of the general-purpose base register
+ */
+.macro _ldr_zt nx
+	_check_general_reg_number \nx
+	.inst	0xe11f8000 | (((\nx) & 0x1f) << 5)
+.endm
+
+/*
+ * STR (ZT0)  : Store the 64-byte ZT0 register to the memory address
+ *              provided in the 64-bit base register or stack pointer.
+ * Instruction: STR ZT0, [<Xn|SP>]
+ *
+ * STR ZT0, nx
+ * Opcode bit field:
+ * nx    : 64-bit name of the general-purpose base register
+ */
+.macro _str_zt nx
+	.inst	0xe13f8000 | (((\nx) & 0x1f) << 5)
+.endm
+
+/*
+ * void sme2_load_zt0_instruction;
+ *
+ * This function loads data from input buffer pointed
+ * to by X0 register into the 512 bits ZT0 register.
+ */
+func sme2_load_zt0_instruction
+	_ldr_zt	0
+	ret
+endfunc sme2_load_zt0_instruction
+
+/*
+ * void sme2_store_zt0_instruction;
+ *
+ * This function stores data from the SME2 ZT0 register
+ * into the memory section pointed by the x0 register.
+ * It copies 512bits of data to the inout memory buffer.
+ */
+func sme2_store_zt0_instruction
+	_str_zt	0
+	ret
+endfunc sme2_store_zt0_instruction
diff --git a/tftf/framework/framework.mk b/tftf/framework/framework.mk
index 6f4022d..a2d6926 100644
--- a/tftf/framework/framework.mk
+++ b/tftf/framework/framework.mk
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2018-2022, Arm Limited. All rights reserved.
+# Copyright (c) 2018-2023, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -84,7 +84,9 @@
 	lib/extensions/pauth/aarch64/pauth.c				\
 	lib/extensions/pauth/aarch64/pauth_helpers.S			\
 	lib/extensions/sme/aarch64/sme.c				\
-	lib/extensions/sme/aarch64/sme_helpers.S
+	lib/extensions/sme/aarch64/sme2.c				\
+	lib/extensions/sme/aarch64/sme_helpers.S			\
+	lib/extensions/sme/aarch64/sme2_helpers.S
 endif
 
 TFTF_LINKERFILE		:=	tftf/framework/tftf.ld.S
diff --git a/tftf/tests/extensions/sme/test_sme2.c b/tftf/tests/extensions/sme/test_sme2.c
new file mode 100644
index 0000000..67122f2
--- /dev/null
+++ b/tftf/tests/extensions/sme/test_sme2.c
@@ -0,0 +1,115 @@
+/*
+ * Copyright (c) 2023, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#include <arch_features.h>
+#include <arch_helpers.h>
+#include <lib/extensions/sme.h>
+#include <test_helpers.h>
+#include <tftf_lib.h>
+
+#ifdef __aarch64__
+
+#define SME2_ARRAYSIZE		(512/64)
+#define SME2_INPUT_DATA		(0x1fffffffffffffff)
+
+/* Global buffers */
+static uint64_t sme2_input_buffer[SME2_ARRAYSIZE] = {0};
+static uint64_t sme2_output_buffer[SME2_ARRAYSIZE] = {0};
+
+/*
+ * clear_ZT0: ZERO all bytes of the ZT0 register.
+ *
+ */
+static void clear_ZT0(void)
+{
+	/**
+	 * Due to the lack of support from the toolchain, instruction
+	 * opcodes are used here.
+	 * TODO: Further, once the toolchain adds support for SME features
+	 * this could be replaced with the instruction ZERO {ZT0}.
+	 */
+	asm volatile(".inst 0xc0480001" : : : );
+}
+
+#endif /* __aarch64__ */
+
+/*
+ * test_sme2_support: Test SME2 support when the extension is enabled.
+ *
+ * Execute some SME2 instructions. These should not be trapped to EL3,
+ * as TF-A is responsible for enabling SME2 for Non-secure world.
+ *
+ */
+test_result_t test_sme2_support(void)
+{
+	/* SME2 is an AArch64-only feature.*/
+	SKIP_TEST_IF_AARCH32();
+
+#ifdef __aarch64__
+	/* 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.
+	 *
+	 * 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.
+	 */
+
+	/* Make sure we can acesss SME2 ZT0 storage, PSTATE.ZA = 1*/
+	VERBOSE("Enabling SME ZA storage and ZT0 storage.\n");
+
+	sme_smstart(SMSTART_ZA);
+
+	/*
+	 * LDR (ZT0) : Load ZT0 register.
+	 *             Load the 64-byte ZT0 register from the memory address
+	 *             provided in the 64-bit scalar base register.
+	 */
+	for (int i = 0; i < SME2_ARRAYSIZE; i++) {
+		sme2_input_buffer[i] = SME2_INPUT_DATA;
+	}
+	sme2_load_zt0_instruction(sme2_input_buffer);
+
+	/*
+	 * STR (ZT0) : Store ZT0 register.
+	 *             Store the 64-byte ZT0 register to the memory address
+	 *             provided in the 64-bit scalar base register
+	 */
+
+	sme2_store_zt0_instruction(sme2_output_buffer);
+
+	/**
+	 * compare the input and output buffer to verify the operations of
+	 * LDR and STR instructions with ZT0 register.
+	 */
+	for (int i = 0; i < SME2_ARRAYSIZE; i++) {
+		if (sme2_input_buffer[i] != sme2_output_buffer[i]) {
+			return TEST_RESULT_FAIL;
+		}
+	}
+
+	/* ZER0 (ZT0) */
+	clear_ZT0();
+
+	/* Finally disable the acesss to SME2 ZT0 storage, PSTATE.ZA = 0*/
+	VERBOSE("Disabling SME ZA storage and ZT0 storage.\n");
+
+	sme_smstop(SMSTOP_ZA);
+
+	return TEST_RESULT_SUCCESS;
+#endif /* __aarch64__ */
+}
diff --git a/tftf/tests/tests-cpu-extensions.mk b/tftf/tests/tests-cpu-extensions.mk
index 3b37673..f838b4b 100644
--- a/tftf/tests/tests-cpu-extensions.mk
+++ b/tftf/tests/tests-cpu-extensions.mk
@@ -1,5 +1,5 @@
 #
-# Copyright (c) 2018-2022, Arm Limited. All rights reserved.
+# Copyright (c) 2018-2023, Arm Limited. All rights reserved.
 #
 # SPDX-License-Identifier: BSD-3-Clause
 #
@@ -15,7 +15,8 @@
 	extensions/pauth/test_pauth.c					\
 	extensions/sve/sve_operations.S					\
 	extensions/sme/test_sme.c					\
-	extensions/spe/test_spe.c                                       \
+	extensions/sme/test_sme2.c					\
+	extensions/spe/test_spe.c					\
 	extensions/sve/test_sve.c					\
 	extensions/sys_reg_trace/test_sys_reg_trace.c			\
 	extensions/trbe/test_trbe.c					\
diff --git a/tftf/tests/tests-cpu-extensions.xml b/tftf/tests/tests-cpu-extensions.xml
index 44ed522..3935911 100644
--- a/tftf/tests/tests-cpu-extensions.xml
+++ b/tftf/tests/tests-cpu-extensions.xml
@@ -1,7 +1,7 @@
 <?xml version="1.0" encoding="utf-8"?>
 
 <!--
-  Copyright (c) 2018-2022, Arm Limited. All rights reserved.
+  Copyright (c) 2018-2023, Arm Limited. All rights reserved.
 
   SPDX-License-Identifier: BSD-3-Clause
 -->
@@ -25,6 +25,7 @@
     <testcase name="Use trace filter control Registers" function="test_trf_enabled" />
     <testcase name="Use trace system Registers" function="test_sys_reg_trace_enabled" />
     <testcase name="SME support" function="test_sme_support" />
+    <testcase name="SME2 support" function="test_sme2_support" />
     <testcase name="SPE support" function="test_spe_support" />
     <testcase name="AFP support" function="test_afp_support" />
     <testcase name="Test wfit instruction" function="test_wfit_instruction" />