Merge "refactor(lib/arch): add macro to update SIMD related trap fields" into integration
diff --git a/lib/arch/include/arch.h b/lib/arch/include/arch.h
index 8d0ed42..c236e15 100644
--- a/lib/arch/include/arch.h
+++ b/lib/arch/include/arch.h
@@ -780,6 +780,10 @@
 #define CPTR_EL2_VHE_SMEN_TRAP_ALL_00	UL(0x0)
 #define CPTR_EL2_VHE_SMEN_NO_TRAP_11	UL(0x3)
 
+#define CPTR_EL2_VHE_SIMD_MASK		(MASK(CPTR_EL2_VHE_FPEN) |	\
+					 MASK(CPTR_EL2_VHE_ZEN) |	\
+					 MASK(CPTR_EL2_VHE_SMEN))
+
 /* Trap all AMU, trace, FPU, SVE, SME accesses */
 #define CPTR_EL2_VHE_INIT		((CPTR_EL2_VHE_ZEN_TRAP_ALL_00 << \
 					  CPTR_EL2_VHE_ZEN_SHIFT)	| \
diff --git a/lib/arch/include/simd.h b/lib/arch/include/simd.h
index 4b7bb41..c3c7a02 100644
--- a/lib/arch/include/simd.h
+++ b/lib/arch/include/simd.h
@@ -218,6 +218,41 @@
 }
 
 /*
+ * Enable SIMD related flags like FPEN, ZEN, SMEN in 'cptr_val' based on SIMD
+ * configuration.
+ */
+#define SIMD_ENABLE_CPTR_FLAGS(simd_cfg, cptr_val)				\
+	do {									\
+		(cptr_val) &= ~(CPTR_EL2_VHE_SIMD_MASK);			\
+										\
+		(cptr_val) |= INPLACE(CPTR_EL2_VHE_FPEN,			\
+				      CPTR_EL2_VHE_FPEN_NO_TRAP_11);		\
+										\
+		if ((simd_cfg)->sve_en) {					\
+			(cptr_val) |= INPLACE(CPTR_EL2_VHE_ZEN,			\
+					      CPTR_EL2_VHE_ZEN_NO_TRAP_11);	\
+		}								\
+										\
+		if ((simd_cfg)->sme_en) {					\
+			(cptr_val) |= INPLACE(CPTR_EL2_VHE_SMEN,		\
+					      CPTR_EL2_VHE_SMEN_NO_TRAP_11);	\
+		}								\
+	} while (false)
+
+/* Disable all SIMD related flags like FPEN, ZEN, SMEN in 'cptr_val' */
+#define SIMD_DISABLE_ALL_CPTR_FLAGS(cptr_val)					\
+	do {									\
+		(cptr_val) &= ~(CPTR_EL2_VHE_SIMD_MASK);			\
+										\
+		(cptr_val) |= INPLACE(CPTR_EL2_VHE_FPEN,			\
+				      CPTR_EL2_VHE_FPEN_TRAP_ALL_00) |		\
+			      INPLACE(CPTR_EL2_VHE_ZEN,				\
+				      CPTR_EL2_VHE_ZEN_TRAP_ALL_00) |		\
+			      INPLACE(CPTR_EL2_VHE_SMEN,			\
+				      CPTR_EL2_VHE_SMEN_TRAP_ALL_00);		\
+	} while (false)
+
+/*
  * RMM support to use SIMD (FPU) at REL2
  */
 #ifdef RMM_FPU_USE_AT_REL2
diff --git a/lib/arch/src/simd.c b/lib/arch/src/simd.c
index 624443b..9ac0db3 100644
--- a/lib/arch/src/simd.c
+++ b/lib/arch/src/simd.c
@@ -336,13 +336,9 @@
 		simd_ctx->sflags &= ~SIMD_SFLAG_SAVED;
 	}
 
-	/*
-	 * Construct the cptr_el2 to be used when this context needs to be
-	 * saved and used by the owner.
-	 */
+	/* Construct the cptr_el2 for this context */
 	simd_ctx->cptr_el2 = CPTR_EL2_VHE_INIT;
-	simd_ctx->cptr_el2 |= INPLACE(CPTR_EL2_VHE_FPEN,
-				      CPTR_EL2_VHE_FPEN_NO_TRAP_11);
+	SIMD_ENABLE_CPTR_FLAGS(simd_cfg, simd_ctx->cptr_el2);
 
 	/* Initialize SVE related fields and config registers */
 	if (simd_cfg->sve_en) {
@@ -351,18 +347,12 @@
 		simd_ctx->el2_regs.zcr_el2 = 0UL;
 		simd_ctx->el2_regs.zcr_el2 |= INPLACE(ZCR_EL2_LEN,
 						      simd_cfg->sve_vq);
-
-		simd_ctx->cptr_el2 |= INPLACE(CPTR_EL2_VHE_ZEN,
-					      CPTR_EL2_VHE_ZEN_NO_TRAP_11);
 	}
 
 	/* Initialize SME related fields */
 	if (simd_cfg->sme_en) {
 		simd_ctx->tflags |= SIMD_TFLAG_SME;
 		simd_ctx->svcr = 0UL;
-
-		simd_ctx->cptr_el2 |= INPLACE(CPTR_EL2_VHE_SMEN,
-					      CPTR_EL2_VHE_SMEN_NO_TRAP_11);
 	}
 
 	simd_ctx->sflags |= SIMD_SFLAG_INIT_DONE;
diff --git a/runtime/core/exit.c b/runtime/core/exit.c
index 1f473b8..c7ead2f 100644
--- a/runtime/core/exit.c
+++ b/runtime/core/exit.c
@@ -338,13 +338,11 @@
 	rec->active_simd_ctx = simd_context_switch(rec->active_simd_ctx,
 						   rec->aux_data.simd_ctx);
 
-	/* Update REC's cptr_el2 */
-	rec->sysregs.cptr_el2 |= INPLACE(CPTR_EL2_VHE_FPEN,
-					 CPTR_EL2_VHE_FPEN_NO_TRAP_11);
-	if (rec->realm_info.simd_cfg.sve_en) {
-		rec->sysregs.cptr_el2 |= INPLACE(CPTR_EL2_VHE_ZEN,
-						 CPTR_EL2_VHE_ZEN_NO_TRAP_11);
-	}
+	/*
+	 * As the REC SIMD context is now restored, enable SIMD flags in REC's
+	 * cptr based on REC's SIMD configuration.
+	 */
+	SIMD_ENABLE_CPTR_FLAGS(&rec->realm_info.simd_cfg, rec->sysregs.cptr_el2);
 
 	/*
 	 * Return 'true' indicating that this exception has been handled and
diff --git a/runtime/core/run.c b/runtime/core/run.c
index 2f6c8ee..6b40c8c 100644
--- a/runtime/core/run.c
+++ b/runtime/core/run.c
@@ -345,15 +345,12 @@
 	if (rec->active_simd_ctx == rec->aux_data.simd_ctx) {
 		(void)simd_context_switch(rec->active_simd_ctx,
 					  &g_ns_simd_ctx[cpuid]);
+
 		/*
-		 * FPU/SVE traps must be enabled in REC's cptr_el2. This enables
-		 * traps again as it was disabled earlier.
+		 * As the REC SIMD context is now saved, disable all SIMD related
+		 * flags in REC's cptr.
 		 */
-		rec->sysregs.cptr_el2 &= ~(MASK(CPTR_EL2_VHE_FPEN) |
-					   MASK(CPTR_EL2_VHE_ZEN));
-		rec->sysregs.cptr_el2 |= INPLACE(CPTR_EL2_VHE_FPEN,
-						 CPTR_EL2_VHE_FPEN_TRAP_ALL_00) |
-			INPLACE(CPTR_EL2_VHE_ZEN, CPTR_EL2_VHE_ZEN_TRAP_ALL_00);
+		SIMD_DISABLE_ALL_CPTR_FLAGS(rec->sysregs.cptr_el2);
 	}
 
 	/* Clear active simd_context */