Merge "fix(spmd): perform G0 interrupt acknowledge and deactivation" into integration
diff --git a/docs/getting_started/build-options.rst b/docs/getting_started/build-options.rst
index a5633e9..276daa4 100644
--- a/docs/getting_started/build-options.rst
+++ b/docs/getting_started/build-options.rst
@@ -45,6 +45,12 @@
    compiling TF-A. Its value must be a numeric, and defaults to 0. See also,
    *Armv8 Architecture Extensions* in :ref:`Firmware Design`.
 
+-  ``ARM_BL2_SP_LIST_DTS``: Path to DTS file snippet to override the hardcoded
+   SP nodes in tb_fw_config.
+
+-  ``ARM_SPMC_MANIFEST_DTS`` : path to an alternate manifest file used as the
+   SPMC Core manifest. Valid when ``SPD=spmd`` is selected.
+
 -  ``BL2``: This is an optional build option which specifies the path to BL2
    image for the ``fip`` target. In this case, the BL2 in the TF-A will not be
    built.
@@ -727,6 +733,10 @@
    1 (do save and restore). 0 is the default. An SPD may set this to 1 if it
    wants the timer registers to be saved and restored.
 
+-  ``OPTEE_SP_FW_CONFIG``: DTC build flag to include OP-TEE as SP in
+   tb_fw_config device tree. This flag is defined only when
+   ``ARM_SPMC_MANIFEST_DTS`` manifest file name contains pattern optee_sp.
+
 -  ``OVERRIDE_LIBC``: This option allows platforms to override the default libc
    for the BL image. It can be either 0 (include) or 1 (remove). The default
    value is 0.
@@ -936,6 +946,9 @@
       When ``EL3_EXCEPTION_HANDLING`` is ``1``, ``TSP_NS_INTR_ASYNC_PREEMPT``
       must also be set to ``1``.
 
+-  ``TS_SP_FW_CONFIG``: DTC build flag to include Trusted Services (Crypto and
+   internal-trusted-storage) as SP in tb_fw_config device tree.
+
 -  ``TWED_DELAY``: Numeric value to be set in order to delay the trapping of
    WFE instruction. ``ENABLE_FEAT_TWED`` build option must be enabled to set
    this delay. It can take values in the range (0-15). Default value is ``0``
diff --git a/docs/plat/arm/arm-build-options.rst b/docs/plat/arm/arm-build-options.rst
index e0b9242..3179267 100644
--- a/docs/plat/arm/arm-build-options.rst
+++ b/docs/plat/arm/arm-build-options.rst
@@ -109,19 +109,6 @@
    (```ethosn.bin```). This firmware image will be included in the FIP and
    loaded at runtime.
 
--  ``ARM_SPMC_MANIFEST_DTS`` : path to an alternate manifest file used as the
-   SPMC Core manifest. Valid when ``SPD=spmd`` is selected.
-
--  ``ARM_BL2_SP_LIST_DTS``: Path to DTS file snippet to override the hardcoded
-   SP nodes in tb_fw_config.
-
--  ``OPTEE_SP_FW_CONFIG``: DTC build flag to include OP-TEE as SP in tb_fw_config
-   device tree. This flag is defined only when ``ARM_SPMC_MANIFEST_DTS`` manifest
-   file name contains pattern optee_sp.
-
--  ``TS_SP_FW_CONFIG``: DTC build flag to include Trusted Services (Crypto and
-   internal-trusted-storage) as SP in tb_fw_config device tree.
-
 -  ``ARM_GPT_SUPPORT``: Enable GPT parser to get the entry address and length of
    the various partitions present in the GPT image. This support is available
    only for the BL2 component, and it is disabled by default.
diff --git a/include/lib/extensions/amu.h b/include/lib/extensions/amu.h
index de476e4..80ad68c 100644
--- a/include/lib/extensions/amu.h
+++ b/include/lib/extensions/amu.h
@@ -16,13 +16,21 @@
 
 #if ENABLE_FEAT_AMU
 #if __aarch64__
-void amu_enable(bool el2_unused, cpu_context_t *ctx);
+void amu_enable(cpu_context_t *ctx);
+void amu_init_el3(void);
+void amu_init_el2_unused(void);
 #else
 void amu_enable(bool el2_unused);
 #endif
 #else
 #if __aarch64__
-static inline void amu_enable(bool el2_unused, cpu_context_t *ctx)
+void amu_enable(cpu_context_t *ctx)
+{
+}
+void amu_init_el3(void)
+{
+}
+void amu_init_el2_unused(void)
 {
 }
 #else
diff --git a/include/lib/extensions/brbe.h b/include/lib/extensions/brbe.h
index 9ee2444..194efba 100644
--- a/include/lib/extensions/brbe.h
+++ b/include/lib/extensions/brbe.h
@@ -8,9 +8,9 @@
 #define BRBE_H
 
 #if ENABLE_BRBE_FOR_NS
-void brbe_enable(void);
+void brbe_init_el3(void);
 #else
-static inline void brbe_enable(void)
+static inline void brbe_init_el3(void)
 {
 }
 #endif /* ENABLE_BRBE_FOR_NS */
diff --git a/include/lib/extensions/mpam.h b/include/lib/extensions/mpam.h
index 4327278..e5438ce 100644
--- a/include/lib/extensions/mpam.h
+++ b/include/lib/extensions/mpam.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2018-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -10,11 +10,15 @@
 #include <stdbool.h>
 
 #if ENABLE_MPAM_FOR_LOWER_ELS
-void mpam_enable(bool el2_unused);
+void mpam_init_el3(void);
+void mpam_init_el2_unused(void);
 #else
-static inline void mpam_enable(bool el2_unused)
+static inline void mpam_init_el3(void)
 {
 }
-#endif
+static inline void mpam_init_el2_unused(void)
+{
+}
+#endif /* ENABLE_MPAM_FOR_LOWER_ELS */
 
 #endif /* MPAM_H */
diff --git a/include/lib/extensions/pmuv3.h b/include/lib/extensions/pmuv3.h
index 5d5d055..62fee7b 100644
--- a/include/lib/extensions/pmuv3.h
+++ b/include/lib/extensions/pmuv3.h
@@ -9,7 +9,7 @@
 
 #include <context.h>
 
-void pmuv3_disable_el3(void);
+void pmuv3_init_el3(void);
 
 #ifdef __aarch64__
 void pmuv3_enable(cpu_context_t *ctx);
diff --git a/include/lib/extensions/sme.h b/include/lib/extensions/sme.h
index 0e9c4b9..dbefdfc 100644
--- a/include/lib/extensions/sme.h
+++ b/include/lib/extensions/sme.h
@@ -22,11 +22,19 @@
 
 #if ENABLE_SME_FOR_NS
 void sme_enable(cpu_context_t *context);
+void sme_init_el3(void);
+void sme_init_el2_unused(void);
 void sme_disable(cpu_context_t *context);
 #else
 static inline void sme_enable(cpu_context_t *context)
 {
 }
+static inline void sme_init_el3(void)
+{
+}
+static inline void sme_init_el2_unused(void)
+{
+}
 static inline void sme_disable(cpu_context_t *context)
 {
 }
diff --git a/include/lib/extensions/spe.h b/include/lib/extensions/spe.h
index 02fccae..7b39037 100644
--- a/include/lib/extensions/spe.h
+++ b/include/lib/extensions/spe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2018, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -10,15 +10,19 @@
 #include <stdbool.h>
 
 #if ENABLE_SPE_FOR_NS
-void spe_enable(bool el2_unused);
+void spe_init_el3(void);
+void spe_init_el2_unused(void);
 void spe_disable(void);
 #else
-static inline void spe_enable(bool el2_unused)
+static inline void spe_init_el3(void)
 {
 }
+static inline void spe_init_el2_unused(void)
+{
+}
 static inline void spe_disable(void)
 {
 }
-#endif
+#endif /* ENABLE_SPE_FOR_NS */
 
 #endif /* SPE_H */
diff --git a/include/lib/extensions/sve.h b/include/lib/extensions/sve.h
index 1faed2d..fc76a16 100644
--- a/include/lib/extensions/sve.h
+++ b/include/lib/extensions/sve.h
@@ -11,11 +11,15 @@
 
 #if (ENABLE_SME_FOR_NS || ENABLE_SVE_FOR_NS)
 void sve_enable(cpu_context_t *context);
+void sve_init_el2_unused(void);
 void sve_disable(cpu_context_t *context);
 #else
 static inline void sve_enable(cpu_context_t *context)
 {
 }
+static inline void sve_init_el2_unused(void)
+{
+}
 static inline void sve_disable(cpu_context_t *context)
 {
 }
diff --git a/include/lib/extensions/sys_reg_trace.h b/include/lib/extensions/sys_reg_trace.h
index 5915c55..d9f7f1b 100644
--- a/include/lib/extensions/sys_reg_trace.h
+++ b/include/lib/extensions/sys_reg_trace.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -10,10 +10,12 @@
 #include <context.h>
 
 #if ENABLE_SYS_REG_TRACE_FOR_NS
+
 #if __aarch64__
 void sys_reg_trace_enable(cpu_context_t *context);
+void sys_reg_trace_init_el2_unused(void);
 #else
-void sys_reg_trace_enable(void);
+void sys_reg_trace_init_el3(void);
 #endif /* __aarch64__ */
 
 #else /* !ENABLE_SYS_REG_TRACE_FOR_NS */
@@ -22,11 +24,18 @@
 static inline void sys_reg_trace_enable(cpu_context_t *context)
 {
 }
+static inline void sys_reg_trace_disable(cpu_context_t *context)
+{
+}
+static inline void sys_reg_trace_init_el2_unused(void)
+{
+}
 #else
-static inline void sys_reg_trace_enable(void)
+static inline void sys_reg_trace_init_el3(void)
 {
 }
 #endif /* __aarch64__ */
+
 #endif /* ENABLE_SYS_REG_TRACE_FOR_NS */
 
 #endif /* SYS_REG_TRACE_H */
diff --git a/include/lib/extensions/trbe.h b/include/lib/extensions/trbe.h
index 861a4ad..0bed433 100644
--- a/include/lib/extensions/trbe.h
+++ b/include/lib/extensions/trbe.h
@@ -8,9 +8,13 @@
 #define TRBE_H
 
 #if ENABLE_TRBE_FOR_NS
-void trbe_enable(void);
+void trbe_init_el3(void);
+void trbe_init_el2_unused(void);
 #else
-static inline void trbe_enable(void)
+static inline void trbe_init_el3(void)
+{
+}
+static inline void trbe_init_el2_unused(void)
 {
 }
 #endif /* ENABLE_TRBE_FOR_NS */
diff --git a/include/lib/extensions/trf.h b/include/lib/extensions/trf.h
index 91a9615..1ac7cda 100644
--- a/include/lib/extensions/trf.h
+++ b/include/lib/extensions/trf.h
@@ -8,9 +8,13 @@
 #define TRF_H
 
 #if ENABLE_TRF_FOR_NS
-void trf_enable(void);
+void trf_init_el3(void);
+void trf_init_el2_unused(void);
 #else
-static inline void trf_enable(void)
+static inline void trf_init_el3(void)
+{
+}
+static inline void trf_init_el2_unused(void)
 {
 }
 #endif /* ENABLE_TRF_FOR_NS */
diff --git a/include/lib/psa/psa_manifest/sid.h b/include/lib/psa/psa_manifest/sid.h
index be78bae..7183112 100644
--- a/include/lib/psa/psa_manifest/sid.h
+++ b/include/lib/psa/psa_manifest/sid.h
@@ -8,6 +8,9 @@
 #ifndef PSA_MANIFEST_SID_H
 #define PSA_MANIFEST_SID_H
 
+/******** RSS_SP_CRYPTO ********/
+#define RSS_CRYPTO_HANDLE				(0x40000100U)
+
 /******** RSS_SP_PLATFORM ********/
 #define RSS_PLATFORM_SERVICE_HANDLE			(0x40000105U)
 
diff --git a/include/lib/psa/rss_crypto_defs.h b/include/lib/psa/rss_crypto_defs.h
new file mode 100644
index 0000000..b8c7426
--- /dev/null
+++ b/include/lib/psa/rss_crypto_defs.h
@@ -0,0 +1,58 @@
+/*
+ * Copyright (c) 2023, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+
+#ifndef RSS_CRYPTO_DEFS_H
+#define RSS_CRYPTO_DEFS_H
+
+/* Declares types that encode errors, algorithms, key types, policies, etc. */
+#include "psa/crypto_types.h"
+
+/*
+ * Value identifying export public key function API, used to dispatch the request
+ * to the corresponding API implementation in the Crypto service backend.
+ *
+ */
+#define RSS_CRYPTO_EXPORT_PUBLIC_KEY_SID	(uint16_t)(0x701)
+
+/*
+ * The persistent key identifiers for RSS builtin keys.
+ */
+enum rss_key_id_builtin_t {
+	RSS_BUILTIN_KEY_ID_HOST_S_ROTPK = 0x7FFF816Cu,
+	RSS_BUILTIN_KEY_ID_HOST_NS_ROTPK,
+	RSS_BUILTIN_KEY_ID_HOST_CCA_ROTPK,
+};
+
+/*
+ * This type is used to overcome a limitation within RSS firmware in the number of maximum
+ * IOVECs it can use especially in psa_aead_encrypt and psa_aead_decrypt.
+ */
+#define RSS_CRYPTO_MAX_NONCE_LENGTH (16u)
+struct rss_crypto_aead_pack_input {
+	uint8_t nonce[RSS_CRYPTO_MAX_NONCE_LENGTH];
+	uint32_t nonce_length;
+};
+
+/*
+ * Structure used to pack non-pointer types in a call
+ */
+struct rss_crypto_pack_iovec {
+	psa_key_id_t key_id;	/* Key id */
+	psa_algorithm_t alg;	/* Algorithm */
+	uint32_t op_handle;	/* Frontend context handle associated
+				   to a multipart operation */
+	uint32_t capacity;	/* Key derivation capacity */
+	uint32_t ad_length;	/* Additional Data length for multipart AEAD */
+	uint32_t plaintext_length;	/* Plaintext length for multipart AEAD */
+	struct rss_crypto_aead_pack_input aead_in;	/* Packs AEAD-related inputs */
+	uint16_t function_id;	/* Used to identify the function in the API dispatcher
+				   to the service backend. See rss_crypto_func_sid for
+				   detail */
+	uint16_t step;		/* Key derivation step */
+};
+
+#endif /* RSS_CRYPTO_DEFS_H */
diff --git a/include/lib/psa/rss_platform_api.h b/include/lib/psa/rss_platform_api.h
index 1dd7d05..8f74a51 100644
--- a/include/lib/psa/rss_platform_api.h
+++ b/include/lib/psa/rss_platform_api.h
@@ -11,6 +11,7 @@
 #include <stdint.h>
 
 #include "psa/error.h"
+#include <rss_crypto_defs.h>
 
 #define RSS_PLATFORM_API_ID_NV_READ       (1010)
 #define RSS_PLATFORM_API_ID_NV_INCREMENT  (1011)
@@ -41,4 +42,19 @@
 rss_platform_nv_counter_read(uint32_t counter_id,
 		uint32_t size, uint8_t *val);
 
+/*
+ * Reads the public key or the public part of a key pair in binary format.
+ *
+ * key		Identifier of the key to export.
+ * data		Buffer where the key data is to be written.
+ * data_size	Size of the data buffer in bytes.
+ * data_length	On success, the number of bytes that make up the key data.
+ *
+ * PSA_SUCCESS if the value is read correctly. Otherwise,
+ *	it returns a PSA_ERROR.
+ */
+psa_status_t
+rss_platform_key_read(enum rss_key_id_builtin_t key, uint8_t *data,
+		size_t data_size, size_t *data_length);
+
 #endif /* RSS_PLATFORM_API_H */
diff --git a/lib/el3_runtime/aarch32/context_mgmt.c b/lib/el3_runtime/aarch32/context_mgmt.c
index 6414aaa..b60b8e0 100644
--- a/lib/el3_runtime/aarch32/context_mgmt.c
+++ b/lib/el3_runtime/aarch32/context_mgmt.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2016-2023, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2016-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -142,19 +142,19 @@
 	}
 
 	if (is_feat_sys_reg_trace_supported()) {
-		sys_reg_trace_enable();
+		sys_reg_trace_init_el3();
 	}
 
 	if (is_feat_trf_supported()) {
-		trf_enable();
+		trf_init_el3();
 	}
 
 	/*
 	 * Also applies to PMU < v3. The PMU is only disabled for EL3 and Secure
 	 * state execution. This does not affect lower NS ELs.
 	 */
-	pmuv3_disable_el3();
-#endif
+	pmuv3_init_el3();
+#endif /*  IMAGE_BL32 */
 }
 
 /*******************************************************************************
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 4a6598a..d9ff0b6 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -496,63 +496,42 @@
 }
 
 /*******************************************************************************
- * Enable architecture extensions on first entry to Non-secure world.
- * When EL2 is implemented but unused `el2_unused` is non-zero, otherwise
- * it is zero. This function updates some registers in-place and its contents
- * are being prepared to be moved to cm_manage_extensions_el3 and
- * cm_manage_extensions_nonsecure.
+ * Enable architecture extensions for EL3 execution. This function only updates
+ * registers in-place which are expected to either never change or be
+ * overwritten by el3_exit.
  ******************************************************************************/
-static void manage_extensions_nonsecure_mixed(bool el2_unused, cpu_context_t *ctx)
-{
 #if IMAGE_BL31
+void cm_manage_extensions_el3(void)
+{
 	if (is_feat_spe_supported()) {
-		spe_enable(el2_unused);
+		spe_init_el3();
 	}
 
 	if (is_feat_amu_supported()) {
-		amu_enable(el2_unused, ctx);
-	}
-
-	/* Enable SVE and FPU/SIMD */
-	if (is_feat_sve_supported()) {
-		sve_enable(ctx);
+		amu_init_el3();
 	}
 
 	if (is_feat_sme_supported()) {
-		sme_enable(ctx);
+		sme_init_el3();
 	}
 
 	if (is_feat_mpam_supported()) {
-		mpam_enable(el2_unused);
+		mpam_init_el3();
 	}
 
 	if (is_feat_trbe_supported()) {
-		trbe_enable();
+		trbe_init_el3();
 	}
 
 	if (is_feat_brbe_supported()) {
-		brbe_enable();
-	}
-
-	if (is_feat_sys_reg_trace_supported()) {
-		sys_reg_trace_enable(ctx);
+		brbe_init_el3();
 	}
 
 	if (is_feat_trf_supported()) {
-		trf_enable();
+		trf_init_el3();
 	}
-#endif
-}
 
-/*******************************************************************************
- * Enable architecture extensions for EL3 execution. This function only updates
- * registers in-place which are expected to either never change or be
- * overwritten by el3_exit.
- ******************************************************************************/
-#if IMAGE_BL31
-void cm_manage_extensions_el3(void)
-{
-	pmuv3_disable_el3();
+	pmuv3_init_el3();
 }
 #endif /* IMAGE_BL31 */
 
@@ -562,6 +541,23 @@
 static void manage_extensions_nonsecure(cpu_context_t *ctx)
 {
 #if IMAGE_BL31
+	if (is_feat_amu_supported()) {
+		amu_enable(ctx);
+	}
+
+	/* Enable SVE and FPU/SIMD */
+	if (is_feat_sve_supported()) {
+		sve_enable(ctx);
+	}
+
+	if (is_feat_sme_supported()) {
+		sme_enable(ctx);
+	}
+
+	if (is_feat_sys_reg_trace_supported()) {
+		sys_reg_trace_enable(ctx);
+	}
+
 	pmuv3_enable(ctx);
 #endif /* IMAGE_BL31 */
 }
@@ -573,7 +569,39 @@
 static void manage_extensions_nonsecure_el2_unused(void)
 {
 #if IMAGE_BL31
+	if (is_feat_spe_supported()) {
+		spe_init_el2_unused();
+	}
+
+	if (is_feat_amu_supported()) {
+		amu_init_el2_unused();
+	}
+
+	if (is_feat_mpam_supported()) {
+		mpam_init_el2_unused();
+	}
+
+	if (is_feat_trbe_supported()) {
+		trbe_init_el2_unused();
+	}
+
+	if (is_feat_sys_reg_trace_supported()) {
+		sys_reg_trace_init_el2_unused();
+	}
+
+	if (is_feat_trf_supported()) {
+		trf_init_el2_unused();
+	}
+
 	pmuv3_init_el2_unused();
+
+	if (is_feat_sve_supported()) {
+		sve_init_el2_unused();
+	}
+
+	if (is_feat_sme_supported()) {
+		sme_init_el2_unused();
+	}
 #endif /* IMAGE_BL31 */
 }
 
@@ -606,6 +634,7 @@
 		 * Enable SME, SVE, FPU/SIMD in secure context, secure manager
 		 * must ensure SME, SVE, and FPU/SIMD context properly managed.
 		 */
+			sme_init_el3();
 			sme_enable(ctx);
 		} else {
 		/*
@@ -656,7 +685,6 @@
 {
 	u_register_t sctlr_elx, scr_el3, mdcr_el2;
 	cpu_context_t *ctx = cm_get_context(security_state);
-	bool el2_unused = false;
 	uint64_t hcr_el2 = 0U;
 
 	assert(ctx != NULL);
@@ -694,8 +722,6 @@
 #endif
 			write_sctlr_el2(sctlr_elx);
 		} else if (el2_implemented != EL_IMPL_NONE) {
-			el2_unused = true;
-
 			/*
 			 * EL2 present but unused, need to disable safely.
 			 * SCTLR_EL2 can be ignored in this case.
@@ -719,24 +745,8 @@
 			 * Initialise CPTR_EL2 setting all fields rather than
 			 * relying on the hw. All fields have architecturally
 			 * UNKNOWN reset values.
-			 *
-			 * CPTR_EL2.TCPAC: Set to zero so that Non-secure EL1
-			 *  accesses to the CPACR_EL1 or CPACR from both
-			 *  Execution states do not trap to EL2.
-			 *
-			 * CPTR_EL2.TTA: Set to zero so that Non-secure System
-			 *  register accesses to the trace registers from both
-			 *  Execution states do not trap to EL2.
-			 *  If PE trace unit System registers are not implemented
-			 *  then this bit is reserved, and must be set to zero.
-			 *
-			 * CPTR_EL2.TFP: Set to zero so that Non-secure accesses
-			 *  to SIMD and floating-point functionality from both
-			 *  Execution states do not trap to EL2.
 			 */
-			write_cptr_el2(CPTR_EL2_RESET_VAL &
-					~(CPTR_EL2_TCPAC_BIT | CPTR_EL2_TTA_BIT
-					| CPTR_EL2_TFP_BIT));
+			write_cptr_el2(CPTR_EL2_RESET_VAL);
 
 			/*
 			 * Initialise CNTHCTL_EL2. All fields are
@@ -787,16 +797,6 @@
 			 * relying on hw. Some fields are architecturally
 			 * UNKNOWN on reset.
 			 *
-			 * MDCR_EL2.TTRF: Set to zero so that access to Trace
-			 *  Filter Control register TRFCR_EL1 at EL1 is not
-			 *  trapped to EL2. This bit is RES0 in versions of
-			 *  the architecture earlier than ARMv8.4.
-			 *
-			 * MDCR_EL2.TPMS: Set to zero so that accesses to
-			 *  Statistical Profiling control registers from EL1
-			 *  do not trap to EL2. This bit is RES0 when SPE is
-			 *  not implemented.
-			 *
 			 * MDCR_EL2.TDRA: Set to zero so that Non-secure EL0 and
 			 *  EL1 System register accesses to the Debug ROM
 			 *  registers are not trapped to EL2.
@@ -810,16 +810,10 @@
 			 *
 			 * MDCR_EL2.TDE: Set to zero so that debug exceptions
 			 *  are not routed to EL2.
-			 *
-			 * MDCR_EL2.E2TB: Set to zero so that the trace Buffer
-			 *  owning exception level is NS-EL1 and, tracing is
-			 *  prohibited at NS-EL2. These bits are RES0 when
-			 *  FEAT_TRBE is not implemented.
 			 */
-			mdcr_el2 = ((MDCR_EL2_RESET_VAL) & ~(MDCR_EL2_TTRF |
-				     MDCR_EL2_TDRA_BIT | MDCR_EL2_TDOSA_BIT |
-				     MDCR_EL2_TDA_BIT | MDCR_EL2_TDE_BIT |
-				     MDCR_EL2_E2TB(MDCR_EL2_E2TB_EL1)));
+			mdcr_el2 = ((MDCR_EL2_RESET_VAL) &
+				   ~(MDCR_EL2_TDRA_BIT | MDCR_EL2_TDOSA_BIT |
+				     MDCR_EL2_TDA_BIT | MDCR_EL2_TDE_BIT));
 
 			write_mdcr_el2(mdcr_el2);
 
@@ -844,7 +838,6 @@
 
 			manage_extensions_nonsecure_el2_unused();
 		}
-		manage_extensions_nonsecure_mixed(el2_unused, ctx);
 	}
 
 	cm_el1_sysregs_context_restore(security_state);
@@ -1149,23 +1142,15 @@
 void cm_prepare_el3_exit_ns(void)
 {
 #if CTX_INCLUDE_EL2_REGS
+#if ENABLE_ASSERTIONS
 	cpu_context_t *ctx = cm_get_context(NON_SECURE);
 	assert(ctx != NULL);
 
 	/* Assert that EL2 is used. */
-#if ENABLE_ASSERTIONS
-	el3_state_t *state = get_el3state_ctx(ctx);
-	u_register_t scr_el3 = read_ctx_reg(state, CTX_SCR_EL3);
-#endif
+	u_register_t scr_el3 = read_ctx_reg(get_el3state_ctx(ctx), CTX_SCR_EL3);
 	assert(((scr_el3 & SCR_HCE_BIT) != 0UL) &&
 			(el_implemented(2U) != EL_IMPL_NONE));
-
-	/*
-	 * Currently some extensions are configured using
-	 * direct register updates. Therefore, do this here
-	 * instead of when setting up context.
-	 */
-	manage_extensions_nonsecure_mixed(0, ctx);
+#endif /* ENABLE_ASSERTIONS */
 
 	/*
 	 * Set the NS bit to be able to access the ICC_SRE_EL2
diff --git a/lib/extensions/amu/aarch64/amu.c b/lib/extensions/amu/aarch64/amu.c
index c650629..e0d3399 100644
--- a/lib/extensions/amu/aarch64/amu.c
+++ b/lib/extensions/amu/aarch64/amu.c
@@ -188,95 +188,66 @@
  * Enable counters. This function is meant to be invoked by the context
  * management library before exiting from EL3.
  */
-void amu_enable(bool el2_unused, cpu_context_t *ctx)
+void amu_enable(cpu_context_t *ctx)
 {
-	uint64_t amcfgr_el0_ncg;		/* Number of counter groups */
-	uint64_t amcgcr_el0_cg0nc;		/* Number of group 0 counters */
-
-	uint64_t amcntenset0_el0_px = 0x0;	/* Group 0 enable mask */
-	uint64_t amcntenset1_el0_px = 0x0;	/* Group 1 enable mask */
-
-	if (el2_unused) {
-		/*
-		 * CPTR_EL2.TAM: Set to zero so any accesses to the Activity
-		 * Monitor registers do not trap to EL2.
-		 */
-		write_cptr_el2_tam(0U);
-	}
-
 	/*
-	 * Retrieve and update the CPTR_EL3 value from the context mentioned
-	 * in 'ctx'. Set CPTR_EL3.TAM to zero so that any accesses to
-	 * the Activity Monitor registers do not trap to EL3.
+	 * Set CPTR_EL3.TAM to zero so that any accesses to the Activity Monitor
+	 * registers do not trap to EL3.
 	 */
 	ctx_write_cptr_el3_tam(ctx, 0U);
 
-	/*
-	 * Retrieve the number of architected counters. All of these counters
-	 * are enabled by default.
-	 */
-
-	amcgcr_el0_cg0nc = read_amcgcr_el0_cg0nc();
-	amcntenset0_el0_px = (UINT64_C(1) << (amcgcr_el0_cg0nc)) - 1U;
+	/* Initialize FEAT_AMUv1p1 features if present. */
+	if (is_feat_amuv1p1_supported()) {
+		/*
+		 * Set SCR_EL3.AMVOFFEN to one so that accesses to virtual
+		 * offset registers at EL2 do not trap to EL3
+		 */
+		ctx_write_scr_el3_amvoffen(ctx, 1U);
+	}
+}
 
-	assert(amcgcr_el0_cg0nc <= AMU_AMCGCR_CG0NC_MAX);
+void amu_init_el3(void)
+{
+	uint64_t group0_impl_ctr = read_amcgcr_el0_cg0nc();
+	uint64_t group0_en_mask = (1 << (group0_impl_ctr)) - 1U;
+	uint64_t num_ctr_groups = read_amcfgr_el0_ncg();
 
-	/*
-	 * The platform may opt to enable specific auxiliary counters. This can
-	 * be done via the common FCONF getter, or via the platform-implemented
-	 * function.
-	 */
+	/* Enable all architected counters by default */
+	write_amcntenset0_el0_px(group0_en_mask);
 
 #if ENABLE_AMU_AUXILIARY_COUNTERS
-	const struct amu_topology *topology;
+	if (num_ctr_groups > 0U) {
+		uint64_t amcntenset1_el0_px = 0x0; /* Group 1 enable mask */
+		const struct amu_topology *topology;
 
+		/*
+		 * The platform may opt to enable specific auxiliary counters.
+		 * This can be done via the common FCONF getter, or via the
+		 * platform-implemented function.
+		 */
 #if ENABLE_AMU_FCONF
-	topology = FCONF_GET_PROPERTY(amu, config, topology);
+		topology = FCONF_GET_PROPERTY(amu, config, topology);
 #else
-	topology = plat_amu_topology();
+		topology = plat_amu_topology();
 #endif /* ENABLE_AMU_FCONF */
 
-	if (topology != NULL) {
-		unsigned int core_pos = plat_my_core_pos();
+		if (topology != NULL) {
+			unsigned int core_pos = plat_my_core_pos();
 
-		amcntenset1_el0_px = topology->cores[core_pos].enable;
-	} else {
-		ERROR("AMU: failed to generate AMU topology\n");
-	}
-#endif /* ENABLE_AMU_AUXILIARY_COUNTERS */
-
-	/*
-	 * Enable the requested counters.
-	 */
-
-	write_amcntenset0_el0_px(amcntenset0_el0_px);
+			amcntenset1_el0_px = topology->cores[core_pos].enable;
+		} else {
+			ERROR("AMU: failed to generate AMU topology\n");
+		}
 
-	amcfgr_el0_ncg = read_amcfgr_el0_ncg();
-	if (amcfgr_el0_ncg > 0U) {
 		write_amcntenset1_el0_px(amcntenset1_el0_px);
-
-#if !ENABLE_AMU_AUXILIARY_COUNTERS
+	}
+#else /* ENABLE_AMU_AUXILIARY_COUNTERS */
+	if (num_ctr_groups > 0U) {
 		VERBOSE("AMU: auxiliary counters detected but support is disabled\n");
-#endif
 	}
+#endif /* ENABLE_AMU_AUXILIARY_COUNTERS */
 
-	/* Initialize FEAT_AMUv1p1 features if present. */
 	if (is_feat_amuv1p1_supported()) {
-		if (el2_unused) {
-			/*
-			 * Make sure virtual offsets are disabled if EL2 not
-			 * used.
-			 */
-			write_hcr_el2_amvoffen(0U);
-		} else {
-			/*
-			 * Virtual offset registers are only accessible from EL3
-			 * and EL2, when clear, this bit traps accesses from EL2
-			 * so we set it to 1 when EL2 is present.
-			 */
-			ctx_write_scr_el3_amvoffen(ctx, 1U);
-		}
-
 #if AMU_RESTRICT_COUNTERS
 		/*
 		 * FEAT_AMUv1p1 adds a register field to restrict access to
@@ -297,6 +268,21 @@
 #endif
 }
 
+void amu_init_el2_unused(void)
+{
+	/*
+	 * CPTR_EL2.TAM: Set to zero so any accesses to the Activity Monitor
+	 *  registers do not trap to EL2.
+	 */
+	write_cptr_el2_tam(0U);
+
+	/* Initialize FEAT_AMUv1p1 features if present. */
+	if (is_feat_amuv1p1_supported()) {
+		/* Make sure virtual offsets are disabled if EL2 not used. */
+		write_hcr_el2_amvoffen(0U);
+	}
+}
+
 /* Read the group 0 counter identified by the given `idx`. */
 static uint64_t amu_group0_cnt_read(unsigned int idx)
 {
@@ -526,10 +512,10 @@
 
 	uint64_t hcr_el2_amvoffen = 0;	/* AMU virtual offsets enabled */
 
-	uint64_t amcfgr_el0_ncg;	/* Number of counter groups */
 	uint64_t amcgcr_el0_cg0nc;	/* Number of group 0 counters */
 
 #if ENABLE_AMU_AUXILIARY_COUNTERS
+	uint64_t amcfgr_el0_ncg;	/* Number of counter groups */
 	uint64_t amcgcr_el0_cg1nc;	/* Number of group 1 counters */
 	uint64_t amcg1idr_el0_voff;	/* Auxiliary counters with virtual offsets */
 #endif
@@ -541,7 +527,6 @@
 	core_pos = plat_my_core_pos();
 	ctx = &amu_ctxs_[core_pos];
 
-	amcfgr_el0_ncg = read_amcfgr_el0_ncg();
 	amcgcr_el0_cg0nc = read_amcgcr_el0_cg0nc();
 
 	if (is_feat_amuv1p1_supported()) {
@@ -549,22 +534,12 @@
 	}
 
 #if ENABLE_AMU_AUXILIARY_COUNTERS
+	amcfgr_el0_ncg = read_amcfgr_el0_ncg();
 	amcgcr_el0_cg1nc = (amcfgr_el0_ncg > 0U) ? read_amcgcr_el0_cg1nc() : 0U;
 	amcg1idr_el0_voff = (hcr_el2_amvoffen != 0U) ? read_amcg1idr_el0_voff() : 0U;
 #endif
 
 	/*
-	 * Sanity check that all counters were disabled when the context was
-	 * previously saved.
-	 */
-
-	assert(read_amcntenset0_el0_px() == 0U);
-
-	if (amcfgr_el0_ncg > 0U) {
-		assert(read_amcntenset1_el0_px() == 0U);
-	}
-
-	/*
 	 * Restore the counter values from the local context.
 	 */
 
diff --git a/lib/extensions/brbe/brbe.c b/lib/extensions/brbe/brbe.c
index 329cf98..37bd834 100644
--- a/lib/extensions/brbe/brbe.c
+++ b/lib/extensions/brbe/brbe.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,8 +7,9 @@
 #include <arch.h>
 #include <arch_features.h>
 #include <arch_helpers.h>
+#include <lib/extensions/brbe.h>
 
-void brbe_enable(void)
+void brbe_init_el3(void)
 {
 	uint64_t val;
 
diff --git a/lib/extensions/mpam/mpam.c b/lib/extensions/mpam/mpam.c
index 62533fc..6462c97 100644
--- a/lib/extensions/mpam/mpam.c
+++ b/lib/extensions/mpam/mpam.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2018-2022, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2018-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -11,7 +11,7 @@
 #include <arch_helpers.h>
 #include <lib/extensions/mpam.h>
 
-void mpam_enable(bool el2_unused)
+void mpam_init_el3(void)
 {
 	/*
 	 * Enable MPAM, and disable trapping to EL3 when lower ELs access their
@@ -19,15 +19,18 @@
 	 */
 	write_mpam3_el3(MPAM3_EL3_MPAMEN_BIT);
 
-	/*
-	 * If EL2 is implemented but unused, disable trapping to EL2 when lower
-	 * ELs access their own MPAM registers.
-	 */
-	if (el2_unused) {
-		write_mpam2_el2(0ULL);
+}
+
+/*
+ * If EL2 is implemented but unused, disable trapping to EL2 when lower ELs
+ * access their own MPAM registers.
+ */
+void mpam_init_el2_unused(void)
+{
+	write_mpam2_el2(0ULL);
 
-		if ((read_mpamidr_el1() & MPAMIDR_HAS_HCR_BIT) != 0U) {
-			write_mpamhcr_el2(0ULL);
-		}
+	if ((read_mpamidr_el1() & MPAMIDR_HAS_HCR_BIT) != 0U) {
+		write_mpamhcr_el2(0ULL);
 	}
+
 }
diff --git a/lib/extensions/pmuv3/aarch32/pmuv3.c b/lib/extensions/pmuv3/aarch32/pmuv3.c
index fe4205e..effb7e0 100644
--- a/lib/extensions/pmuv3/aarch32/pmuv3.c
+++ b/lib/extensions/pmuv3/aarch32/pmuv3.c
@@ -29,7 +29,7 @@
  * Applies to all PMU versions. Name is PMUv3 for compatibility with aarch64 and
  * to not clash with platforms which reuse the PMU name
  */
-void pmuv3_disable_el3(void)
+void pmuv3_init_el3(void)
 {
 	u_register_t sdcr = read_sdcr();
 
diff --git a/lib/extensions/pmuv3/aarch64/pmuv3.c b/lib/extensions/pmuv3/aarch64/pmuv3.c
index f83a5ee..fda71aa 100644
--- a/lib/extensions/pmuv3/aarch64/pmuv3.c
+++ b/lib/extensions/pmuv3/aarch64/pmuv3.c
@@ -48,7 +48,7 @@
 	return mdcr_el3;
 }
 
-void pmuv3_disable_el3(void)
+void pmuv3_init_el3(void)
 {
 	u_register_t mdcr_el3 = read_mdcr_el3();
 
diff --git a/lib/extensions/sme/sme.c b/lib/extensions/sme/sme.c
index 3423dba..d705b64 100644
--- a/lib/extensions/sme/sme.c
+++ b/lib/extensions/sme/sme.c
@@ -17,7 +17,6 @@
 void sme_enable(cpu_context_t *context)
 {
 	u_register_t reg;
-	u_register_t cptr_el3;
 	el3_state_t *state;
 
 	/* Get the context state. */
@@ -32,9 +31,14 @@
 	reg = read_ctx_reg(state, CTX_SCR_EL3);
 	reg |= SCR_ENTP2_BIT;
 	write_ctx_reg(state, CTX_SCR_EL3, reg);
+}
 
-	/* Set CPTR_EL3.ESM bit so we can write SMCR_EL3 without trapping. */
-	cptr_el3 = read_cptr_el3();
+void sme_init_el3(void)
+{
+	u_register_t cptr_el3 = read_cptr_el3();
+	u_register_t smcr_el3;
+
+	/* Set CPTR_EL3.ESM bit so we can access SMCR_EL3 without trapping. */
 	write_cptr_el3(cptr_el3 | ESM_BIT);
 	isb();
 
@@ -43,11 +47,10 @@
 	 * to be the least restrictive, then lower ELs can restrict as needed
 	 * using SMCR_EL2 and SMCR_EL1.
 	 */
-	reg = SMCR_ELX_LEN_MAX;
-
+	smcr_el3 = SMCR_ELX_LEN_MAX;
 	if (read_feat_sme_fa64_id_field() != 0U) {
 		VERBOSE("[SME] FA64 enabled\n");
-		reg |= SMCR_ELX_FA64_BIT;
+		smcr_el3 |= SMCR_ELX_FA64_BIT;
 	}
 
 	/*
@@ -58,15 +61,24 @@
 	 */
 	if (is_feat_sme2_supported()) {
 		VERBOSE("SME2 enabled\n");
-		reg |= SMCR_ELX_EZT0_BIT;
+		smcr_el3 |= SMCR_ELX_EZT0_BIT;
 	}
-	write_smcr_el3(reg);
+	write_smcr_el3(smcr_el3);
 
 	/* Reset CPTR_EL3 value. */
 	write_cptr_el3(cptr_el3);
 	isb();
 }
 
+void sme_init_el2_unused(void)
+{
+	/*
+	 * CPTR_EL2.TCPAC: Set to zero so that Non-secure EL1 accesses to the
+	 *  CPACR_EL1 or CPACR from both Execution states do not trap to EL2.
+	 */
+	write_cptr_el2(read_cptr_el2() & ~CPTR_EL2_TCPAC_BIT);
+}
+
 void sme_disable(cpu_context_t *context)
 {
 	u_register_t reg;
diff --git a/lib/extensions/spe/spe.c b/lib/extensions/spe/spe.c
index b1fe39f..236b102 100644
--- a/lib/extensions/spe/spe.c
+++ b/lib/extensions/spe/spe.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2022, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -21,25 +21,10 @@
 	__asm__ volatile("hint #17");
 }
 
-void spe_enable(bool el2_unused)
+void spe_init_el3(void)
 {
 	uint64_t v;
 
-	if (el2_unused) {
-		/*
-		 * MDCR_EL2.TPMS (ARM v8.2): Do not trap statistical
-		 * profiling controls to EL2.
-		 *
-		 * MDCR_EL2.E2PB (ARM v8.2): SPE enabled in Non-secure
-		 * state. Accesses to profiling buffer controls at
-		 * Non-secure EL1 are not trapped to EL2.
-		 */
-		v = read_mdcr_el2();
-		v &= ~MDCR_EL2_TPMS;
-		v |= MDCR_EL2_E2PB(MDCR_EL2_E2PB_EL1);
-		write_mdcr_el2(v);
-	}
-
 	/*
 	 * MDCR_EL2.NSPB (ARM v8.2): SPE enabled in Non-secure state
 	 * and disabled in secure state. Accesses to SPE registers at
@@ -55,6 +40,24 @@
 	write_mdcr_el3(v);
 }
 
+void spe_init_el2_unused(void)
+{
+	uint64_t v;
+
+	/*
+	 * MDCR_EL2.TPMS (ARM v8.2): Do not trap statistical
+	 * profiling controls to EL2.
+	 *
+	 * MDCR_EL2.E2PB (ARM v8.2): SPE enabled in Non-secure
+	 * state. Accesses to profiling buffer controls at
+	 * Non-secure EL1 are not trapped to EL2.
+	 */
+	v = read_mdcr_el2();
+	v &= ~MDCR_EL2_TPMS;
+	v |= MDCR_EL2_E2PB(MDCR_EL2_E2PB_EL1);
+	write_mdcr_el2(v);
+}
+
 void spe_disable(void)
 {
 	uint64_t v;
diff --git a/lib/extensions/sve/sve.c b/lib/extensions/sve/sve.c
index f551ca7..eb4ac8d 100644
--- a/lib/extensions/sve/sve.c
+++ b/lib/extensions/sve/sve.c
@@ -37,6 +37,16 @@
 		(ZCR_EL3_LEN_MASK & CONVERT_SVE_LENGTH(SVE_VECTOR_LEN)));
 }
 
+void sve_init_el2_unused(void)
+{
+	/*
+	 * CPTR_EL2.TFP: Set to zero so that Non-secure accesses to Advanced
+	 *  SIMD and floating-point functionality from both Execution states do
+	 *  not trap to EL2.
+	 */
+	write_cptr_el2(read_cptr_el2() & ~CPTR_EL2_TFP_BIT);
+}
+
 void sve_disable(cpu_context_t *context)
 {
 	u_register_t reg;
diff --git a/lib/extensions/sys_reg_trace/aarch32/sys_reg_trace.c b/lib/extensions/sys_reg_trace/aarch32/sys_reg_trace.c
index b3f44b7..6da504e 100644
--- a/lib/extensions/sys_reg_trace/aarch32/sys_reg_trace.c
+++ b/lib/extensions/sys_reg_trace/aarch32/sys_reg_trace.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -10,7 +10,7 @@
 #include <arch_helpers.h>
 #include <lib/extensions/sys_reg_trace.h>
 
-void sys_reg_trace_enable(void)
+void sys_reg_trace_init_el3(void)
 {
 	uint32_t val;
 
diff --git a/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c b/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c
index e61cb90..4b57f67 100644
--- a/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c
+++ b/lib/extensions/sys_reg_trace/aarch64/sys_reg_trace.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -24,3 +24,14 @@
 	val &= ~TTA_BIT;
 	write_ctx_reg(get_el3state_ctx(ctx), CTX_CPTR_EL3, val);
 }
+
+void sys_reg_trace_init_el2_unused(void)
+{
+	/*
+	 * CPTR_EL2.TTA: Set to zero so that Non-secure System register accesses
+	 *  to the trace registers from both Execution states do not trap to
+	 *  EL2. If PE trace unit System registers are not implemented then this
+	 *  bit is reserved, and must be set to zero.
+	 */
+	write_cptr_el2(read_cptr_el2() & ~CPTR_EL2_TTA_BIT);
+}
diff --git a/lib/extensions/trbe/trbe.c b/lib/extensions/trbe/trbe.c
index fa139ca..461ea73 100644
--- a/lib/extensions/trbe/trbe.c
+++ b/lib/extensions/trbe/trbe.c
@@ -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
  */
@@ -19,9 +19,9 @@
 	__asm__ volatile("hint #18");
 }
 
-void trbe_enable(void)
+void trbe_init_el3(void)
 {
-	uint64_t val;
+	u_register_t val;
 
 	/*
 	 * MDCR_EL3.NSTB = 0b11
@@ -34,6 +34,17 @@
 	write_mdcr_el3(val);
 }
 
+void trbe_init_el2_unused(void)
+{
+	/*
+	 * MDCR_EL2.E2TB: Set to zero so that the trace Buffer
+	 *  owning exception level is NS-EL1 and, tracing is
+	 *  prohibited at NS-EL2. These bits are RES0 when
+	 *  FEAT_TRBE is not implemented.
+	 */
+	write_mdcr_el2(read_mdcr_el2() & ~MDCR_EL2_E2TB(MDCR_EL2_E2TB_EL1));
+}
+
 static void *trbe_drain_trace_buffers_hook(const void *arg __unused)
 {
 	if (is_feat_trbe_supported()) {
diff --git a/lib/extensions/trf/aarch32/trf.c b/lib/extensions/trf/aarch32/trf.c
index 0c63efa..e13b4db 100644
--- a/lib/extensions/trf/aarch32/trf.c
+++ b/lib/extensions/trf/aarch32/trf.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -10,7 +10,7 @@
 #include <arch_helpers.h>
 #include <lib/extensions/trf.h>
 
-void trf_enable(void)
+void trf_init_el3(void)
 {
 	uint32_t val;
 
diff --git a/lib/extensions/trf/aarch64/trf.c b/lib/extensions/trf/aarch64/trf.c
index 941692b..f681b28 100644
--- a/lib/extensions/trf/aarch64/trf.c
+++ b/lib/extensions/trf/aarch64/trf.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -9,9 +9,9 @@
 #include <arch_helpers.h>
 #include <lib/extensions/trf.h>
 
-void trf_enable(void)
+void trf_init_el3(void)
 {
-	uint64_t val;
+	u_register_t val;
 
 	/*
 	 * MDCR_EL3.TTRF = b0
@@ -22,3 +22,15 @@
 	val &= ~MDCR_TTRF_BIT;
 	write_mdcr_el3(val);
 }
+
+void trf_init_el2_unused(void)
+{
+	/*
+	 * MDCR_EL2.TTRF: Set to zero so that access to Trace
+	 *  Filter Control register TRFCR_EL1 at EL1 is not
+	 *  trapped to EL2. This bit is RES0 in versions of
+	 *  the architecture earlier than ARMv8.4.
+	 *
+	 */
+	write_mdcr_el2(read_mdcr_el2() & ~MDCR_EL2_TTRF);
+}
diff --git a/lib/libc/printf.c b/lib/libc/printf.c
index faccfdf..6931a7e 100644
--- a/lib/libc/printf.c
+++ b/lib/libc/printf.c
@@ -36,7 +36,7 @@
 }
 
 static int unsigned_num_print(unsigned long long int unum, unsigned int radix,
-			      char padc, int padn)
+			      char padc, int padn, bool uppercase)
 {
 	/* Just need enough space to store 64 bit decimal integer */
 	char num_buf[20];
@@ -51,10 +51,13 @@
 
 	do {
 		rem = unum % radix;
-		if (rem < 0xa)
+		if (rem < 0xa) {
 			num_buf[i] = '0' + rem;
-		else
+		} else if (uppercase) {
+			num_buf[i] = 'A' + (rem - 0xa);
+		} else {
 			num_buf[i] = 'a' + (rem - 0xa);
+		}
 		i++;
 		unum /= radix;
 	} while (unum > 0U);
@@ -105,8 +108,10 @@
 	char padc = '\0'; /* Padding character */
 	int padn; /* Number of characters to pad */
 	int count = 0; /* Number of printed characters */
+	bool uppercase; /* Print characters in uppercase */
 
 	while (*fmt != '\0') {
+		uppercase = false;
 		l_count = 0;
 		padn = 0;
 
@@ -129,7 +134,7 @@
 					unum = (unsigned long long int)num;
 
 				count += unsigned_num_print(unum, 10,
-							    padc, padn);
+							    padc, padn, uppercase);
 				break;
 			case 'c':
 				(void)putchar(va_arg(args, int));
@@ -147,12 +152,15 @@
 				}
 
 				count += unsigned_num_print(unum, 16,
-							    padc, padn);
+							    padc, padn, uppercase);
 				break;
+			case 'X':
+				uppercase = true;
+				// fall through
 			case 'x':
 				unum = get_unum_va_args(args, l_count);
 				count += unsigned_num_print(unum, 16,
-							    padc, padn);
+							    padc, padn, uppercase);
 				break;
 			case 'z':
 				if (sizeof(size_t) == 8U)
@@ -167,7 +175,7 @@
 			case 'u':
 				unum = get_unum_va_args(args, l_count);
 				count += unsigned_num_print(unum, 10,
-							    padc, padn);
+							    padc, padn, uppercase);
 				break;
 			case '0':
 				padc = '0';
diff --git a/lib/psa/rss_platform.c b/lib/psa/rss_platform.c
index 359f894..7d90bfc 100644
--- a/lib/psa/rss_platform.c
+++ b/lib/psa/rss_platform.c
@@ -5,10 +5,9 @@
  *
  */
 
-#include <stdint.h>
-
 #include <psa/client.h>
 #include <psa_manifest/sid.h>
+#include <rss_crypto_defs.h>
 #include <rss_platform_api.h>
 
 psa_status_t
@@ -41,3 +40,30 @@
 			RSS_PLATFORM_API_ID_NV_READ,
 			in_vec, 1, out_vec, 1);
 }
+
+psa_status_t
+rss_platform_key_read(enum rss_key_id_builtin_t key, uint8_t *data,
+		size_t data_size, size_t *data_length)
+{
+	psa_status_t status;
+
+	struct rss_crypto_pack_iovec iov = {
+		.function_id = RSS_CRYPTO_EXPORT_PUBLIC_KEY_SID,
+		.key_id = key,
+	};
+
+	psa_invec in_vec[] = {
+		{.base = &iov, .len = sizeof(struct rss_crypto_pack_iovec)},
+	};
+	psa_outvec out_vec[] = {
+		{.base = data, .len = data_size}
+	};
+
+	status = psa_call(RSS_CRYPTO_HANDLE, PSA_IPC_CALL,
+			in_vec, IOVEC_LEN(in_vec),
+			out_vec, IOVEC_LEN(out_vec));
+
+	*data_length = out_vec[0].len;
+
+	return status;
+}
diff --git a/lib/xlat_tables_v2/aarch64/xlat_tables_arch.c b/lib/xlat_tables_v2/aarch64/xlat_tables_arch.c
index c847a9e..bb6a35c 100644
--- a/lib/xlat_tables_v2/aarch64/xlat_tables_arch.c
+++ b/lib/xlat_tables_v2/aarch64/xlat_tables_arch.c
@@ -27,16 +27,14 @@
 	if (size == PAGE_SIZE_4KB) {
 		tgranx = read_id_aa64mmfr0_el0_tgran4_field();
 		/* MSB of TGRAN4 field will be '1' for unsupported feature */
-		return ((tgranx >= ID_AA64MMFR0_EL1_TGRAN4_SUPPORTED) &&
-			(tgranx < 8U));
+		return (tgranx < 8U);
 	} else if (size == PAGE_SIZE_16KB) {
 		tgranx = read_id_aa64mmfr0_el0_tgran16_field();
 		return (tgranx >= ID_AA64MMFR0_EL1_TGRAN16_SUPPORTED);
 	} else if (size == PAGE_SIZE_64KB) {
 		tgranx = read_id_aa64mmfr0_el0_tgran64_field();
 		/* MSB of TGRAN64 field will be '1' for unsupported feature */
-		return ((tgranx >= ID_AA64MMFR0_EL1_TGRAN64_SUPPORTED) &&
-			(tgranx < 8U));
+		return (tgranx < 8U);
 	} else {
 		return false;
 	}
diff --git a/plat/arm/board/tc/include/tc_plat.h b/plat/arm/board/tc/include/tc_plat.h
index 117fbb4..a6b2b0d 100644
--- a/plat/arm/board/tc/include/tc_plat.h
+++ b/plat/arm/board/tc/include/tc_plat.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, ARM Limited and Contributors. All rights reserved.
+ * Copyright (c) 2021-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,6 +7,10 @@
 #ifndef TC_PLAT_H
 #define TC_PLAT_H
 
+#ifdef PLATFORM_TEST_ROTPK
+#include <rss_crypto_defs.h>
+#endif
+
 void tc_bl31_common_platform_setup(void);
 
 #ifdef PLATFORM_TEST_TFM_TESTSUITE
@@ -17,4 +21,13 @@
 int nv_counter_test(void);
 #endif
 
+#ifdef PLATFORM_TEST_ROTPK
+struct key_id_info {
+	enum rss_key_id_builtin_t key_id;
+	const char *key_id_name;
+};
+
+int rotpk_test(void);
+#endif
+
 #endif /* TC_PLAT_H */
diff --git a/plat/arm/board/tc/nv_counter_test.c b/plat/arm/board/tc/nv_counter_test.c
index f9e001e..179ec4b 100644
--- a/plat/arm/board/tc/nv_counter_test.c
+++ b/plat/arm/board/tc/nv_counter_test.c
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, ARM Limited. All rights reserved.
+ * Copyright (c) 2023, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -22,29 +22,29 @@
 
 	status = rss_comms_init(PLAT_RSS_AP_SND_MHU_BASE, PLAT_RSS_AP_RCV_MHU_BASE);
 	if (status != PSA_SUCCESS) {
-		printf("Failed to initialize RSS communication channel\n");
+		printf("Failed to initialize RSS communication channel - psa_status = %d\n", status);
 		return -1;
 	}
 
 	for (id = 0; id < 3; id++) {
 		status = rss_platform_nv_counter_read(id, sizeof(old_val), (uint8_t *)&old_val);
 		if (status != PSA_SUCCESS) {
-			printf("Failed during first id=(%d) rss_platform_nv_counter_read\n",
-				       id);
+			printf("Failed during first id=(%d) rss_platform_nv_counter_read - psa_status = %d\n",
+				       id, status);
 			return -1;
 		}
 
 		status = rss_platform_nv_counter_increment(id);
 		if (status != PSA_SUCCESS) {
-			printf("Failed during id=(%d) rss_platform_nv_counter_increment\n",
-					id);
+			printf("Failed during id=(%d) rss_platform_nv_counter_increment - psa_status = %d\n",
+					id, status);
 			return -1;
 		}
 
 		status = rss_platform_nv_counter_read(id, sizeof(new_val), (uint8_t *)&new_val);
 		if (status != PSA_SUCCESS) {
-			printf("Failed during second id=(%d) rss_platform_nv_counter_read\n",
-					id);
+			printf("Failed during second id=(%d) rss_platform_nv_counter_read - psa_status = %d\n",
+					id, status);
 			return -1;
 		}
 
diff --git a/plat/arm/board/tc/platform.mk b/plat/arm/board/tc/platform.mk
index 5a1d83a..8ca33ca 100644
--- a/plat/arm/board/tc/platform.mk
+++ b/plat/arm/board/tc/platform.mk
@@ -189,31 +189,11 @@
 endif
 
 ifneq (${PLATFORM_TEST},)
-    $(eval $(call add_define,PLATFORM_TESTS))
-
-    ifeq (${PLATFORM_TEST},rss-nv-counters)
-        include drivers/arm/rss/rss_comms.mk
-
-        # Test code.
-        BL31_SOURCES	+=	plat/arm/board/tc/nv_counter_test.c
-
-        # Code under testing.
-        BL31_SOURCES	+=	lib/psa/rss_platform.c \
-				drivers/arm/rss/rss_comms.c \
-				${RSS_COMMS_SOURCES}
-
-        PLAT_INCLUDES	+=	-Iinclude/lib/psa
-
-        $(eval $(call add_define,PLATFORM_TEST_NV_COUNTERS))
-    else ifeq (${PLATFORM_TEST},tfm-testsuite)
-        # Add this include as first, before arm_common.mk. This is necessary
-        # because arm_common.mk builds Mbed TLS, and platform_test.mk can
-        # change the list of Mbed TLS files that are to be compiled
-        # (LIBMBEDTLS_SRCS).
-        include plat/arm/board/tc/platform_test.mk
-    else
-        $(error "Unsupported PLATFORM_TEST value")
-    endif
+    # Add this include as first, before arm_common.mk. This is necessary
+    # because arm_common.mk builds Mbed TLS, and platform_test.mk can
+    # change the list of Mbed TLS files that are to be compiled
+    # (LIBMBEDTLS_SRCS).
+    include plat/arm/board/tc/platform_test.mk
 endif
 
 
diff --git a/plat/arm/board/tc/platform_test.mk b/plat/arm/board/tc/platform_test.mk
index e974855..2fd5ea0 100644
--- a/plat/arm/board/tc/platform_test.mk
+++ b/plat/arm/board/tc/platform_test.mk
@@ -3,7 +3,37 @@
 # SPDX-License-Identifier: BSD-3-Clause
 #
 
-ifeq (${PLATFORM_TEST},tfm-testsuite)
+$(eval $(call add_define,PLATFORM_TESTS))
+
+ifeq (${PLATFORM_TEST},rss-nv-counters)
+    include drivers/arm/rss/rss_comms.mk
+
+    # Test code.
+    BL31_SOURCES	+=	plat/arm/board/tc/nv_counter_test.c
+
+    # Code under testing.
+    BL31_SOURCES	+=	lib/psa/rss_platform.c \
+				drivers/arm/rss/rss_comms.c \
+				${RSS_COMMS_SOURCES}
+
+    PLAT_INCLUDES	+=	-Iinclude/lib/psa
+
+    $(eval $(call add_define,PLATFORM_TEST_NV_COUNTERS))
+else ifeq (${PLATFORM_TEST},rss-rotpk)
+    include drivers/arm/rss/rss_comms.mk
+
+    # Test code.
+    BL31_SOURCES	+=	plat/arm/board/tc/rotpk_test.c
+
+    # Code under testing.
+    BL31_SOURCES	+=	lib/psa/rss_platform.c \
+				drivers/arm/rss/rss_comms.c \
+				${RSS_COMMS_SOURCES}
+
+    PLAT_INCLUDES	+=	-Iinclude/lib/psa
+
+    $(eval $(call add_define,PLATFORM_TEST_ROTPK))
+else ifeq (${PLATFORM_TEST},tfm-testsuite)
 
     # The variables need to be set to compile the platform test:
     ifeq (${TF_M_TESTS_PATH},)
@@ -77,4 +107,6 @@
     $(eval $(call add_define,DELEG_ATTEST_DUMP_TOKEN_AND_KEY))
 
     $(eval $(call add_define,PLATFORM_TEST_TFM_TESTSUITE))
+else
+    $(error "Unsupported PLATFORM_TEST value")
 endif
diff --git a/plat/arm/board/tc/rotpk_test.c b/plat/arm/board/tc/rotpk_test.c
new file mode 100644
index 0000000..ed56c31
--- /dev/null
+++ b/plat/arm/board/tc/rotpk_test.c
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) 2023, Arm Limited. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <stdint.h>
+#include <stdio.h>
+
+#include <drivers/arm/rss_comms.h>
+#include <plat/common/platform.h>
+#include <rss_platform_api.h>
+#include <tc_plat.h>
+
+static void print_hex(const char *key_id_name, size_t key_size, const uint8_t *key_buf)
+{
+	printf("%s = ", key_id_name);
+	for (int i = 0; i < key_size; i++) {
+		printf("%02x", key_buf[i]);
+	}
+	printf("\n\n");
+}
+
+int rotpk_test(void)
+{
+	psa_status_t status;
+	uint8_t key_buf[128];
+	size_t key_size;
+
+	struct key_id_info key_ids[3] = {
+	       {.key_id = RSS_BUILTIN_KEY_ID_HOST_S_ROTPK,  .key_id_name = "Secure-ROTPK"},
+	       {.key_id = RSS_BUILTIN_KEY_ID_HOST_NS_ROTPK,  .key_id_name = "NS-ROTPK"},
+	       {.key_id = RSS_BUILTIN_KEY_ID_HOST_CCA_ROTPK,  .key_id_name = "CCA-ROTPK"}
+	};
+
+	status = rss_comms_init(PLAT_RSS_AP_SND_MHU_BASE, PLAT_RSS_AP_RCV_MHU_BASE);
+	if (status != PSA_SUCCESS) {
+		printf("Failed to initialize RSS communication channel - psa_status = %d\n", status);
+		return -1;
+	}
+
+	for (int i = 0; i < ARRAY_SIZE(key_ids); i++) {
+		status = rss_platform_key_read(key_ids[i].key_id, key_buf,
+			       sizeof(key_buf), &key_size);
+		if (status != PSA_SUCCESS) {
+			printf("Failed to retrieve %s - psa_status = %d\n", key_ids[i].key_id_name, status);
+			return -1;
+		}
+		print_hex(key_ids[i].key_id_name, key_size, key_buf);
+	}
+
+	printf("Passed rotpk_test\n");
+
+	return 0;
+}
diff --git a/plat/arm/board/tc/tc_bl31_setup.c b/plat/arm/board/tc/tc_bl31_setup.c
index 097c75c..ff7809d 100644
--- a/plat/arm/board/tc/tc_bl31_setup.c
+++ b/plat/arm/board/tc/tc_bl31_setup.c
@@ -59,6 +59,8 @@
 
 #ifdef PLATFORM_TEST_NV_COUNTERS
 	tests_failed = nv_counter_test();
+#elif PLATFORM_TEST_ROTPK
+	tests_failed = rotpk_test();
 #elif PLATFORM_TEST_TFM_TESTSUITE
 	tests_failed = run_platform_tests();
 #endif
diff --git a/plat/mediatek/drivers/apusys/devapc/apusys_dapc_v1.h b/plat/mediatek/drivers/apusys/devapc/apusys_dapc_v1.h
index 2f5d47b..1b77942 100644
--- a/plat/mediatek/drivers/apusys/devapc/apusys_dapc_v1.h
+++ b/plat/mediatek/drivers/apusys/devapc/apusys_dapc_v1.h
@@ -155,19 +155,4 @@
 			   FORBIDDEN,     NO_PROTECTION, FORBIDDEN, FORBIDDEN, \
 			   FORBIDDEN,     FORBIDDEN,     FORBIDDEN, FORBIDDEN, \
 			   FORBIDDEN,     FORBIDDEN,     FORBIDDEN, FORBIDDEN)
-
-#define SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT_D3_SEC_RW(domain)			 \
-	APUSYS_APC_AO_ATTR(domain,						 \
-			   NO_PROTECTION, FORBIDDEN,     FORBIDDEN, SEC_RW_ONLY, \
-			   FORBIDDEN,     NO_PROTECTION, FORBIDDEN, FORBIDDEN,   \
-			   FORBIDDEN,     FORBIDDEN,     FORBIDDEN, FORBIDDEN,   \
-			   FORBIDDEN,     FORBIDDEN,     FORBIDDEN, FORBIDDEN)
-
-#define SLAVE_FORBID_EXCEPT_D0_D3_SEC_RW_D5_NO_PROTECT(domain)		       \
-	APUSYS_APC_AO_ATTR(domain,					       \
-			   SEC_RW_ONLY, FORBIDDEN,     FORBIDDEN, SEC_RW_ONLY, \
-			   FORBIDDEN,   NO_PROTECTION, FORBIDDEN, FORBIDDEN,   \
-			   FORBIDDEN,   FORBIDDEN,     FORBIDDEN, FORBIDDEN,   \
-			   FORBIDDEN,   FORBIDDEN,     FORBIDDEN, FORBIDDEN)
-
 #endif /* APUSYS_DAPC_V1_H */
diff --git a/plat/mediatek/drivers/apusys/mt8188/apusys_devapc_def.h b/plat/mediatek/drivers/apusys/mt8188/apusys_devapc_def.h
index e74b022..47a2a94 100644
--- a/plat/mediatek/drivers/apusys/mt8188/apusys_devapc_def.h
+++ b/plat/mediatek/drivers/apusys/mt8188/apusys_devapc_def.h
@@ -29,14 +29,14 @@
 #define SLAVE_AO_BCRM			SLAVE_FORBID_EXCEPT_D5_NO_PROTECT
 #define SLAVE_AO_DAPC_WRAP		SLAVE_FORBID_EXCEPT_D5_NO_PROTECT
 #define SLAVE_AO_DAPC_CON		SLAVE_FORBID_EXCEPT_D0_SEC_RW_D5_NO_PROTECT
-#define SLAVE_RCX_ACX_BULK		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT_D3_SEC_RW
-#define SLAVE_ACX0_BCRM			SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT_D3_SEC_RW
+#define SLAVE_RCX_ACX_BULK		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
+#define SLAVE_ACX0_BCRM			SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
 #define SLAVE_RPCTOP_LITE_ACX0		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
-#define SLAVE_ACX1_BCRM			SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT_D3_SEC_RW
+#define SLAVE_ACX1_BCRM			SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
 #define SLAVE_RPCTOP_LITE_ACX1		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
-#define SLAVE_RCX_TO_ACX0_0		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT_D3_SEC_RW
+#define SLAVE_RCX_TO_ACX0_0		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
 #define SLAVE_RCX_TO_ACX0_1		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
-#define SLAVE_SAE_TO_ACX0_0		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT_D3_SEC_RW
+#define SLAVE_SAE_TO_ACX0_0		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
 #define SLAVE_SAE_TO_ACX0_1		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
 #define SLAVE_RCX_TO_ACX1_0		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
 #define SLAVE_RCX_TO_ACX1_1		SLAVE_FORBID_EXCEPT_D0_D5_NO_PROTECT
@@ -79,7 +79,7 @@
 #define SLAVE_SENSOR_WRAP_ACX1_DLA1	SLAVE_FORBID_EXCEPT_D5_NO_PROTECT
 #define SLAVE_SENSOR_WRAP_ACX1_VPU0	SLAVE_FORBID_EXCEPT_D5_NO_PROTECT
 #define SLAVE_REVISER			SLAVE_FORBID_EXCEPT_D0_SEC_RW
-#define SLAVE_NOC			SLAVE_FORBID_EXCEPT_D0_D3_SEC_RW_D5_NO_PROTECT
+#define SLAVE_NOC			SLAVE_FORBID_EXCEPT_D0_SEC_RW
 #define SLAVE_BCRM			SLAVE_FORBID_EXCEPT_D5_NO_PROTECT
 #define SLAVE_DAPC_WRAP			SLAVE_FORBID_EXCEPT_D5_NO_PROTECT
 #define SLAVE_DAPC_CON			SLAVE_FORBID_EXCEPT_D0_SEC_RW_D5_NO_PROTECT
diff --git a/plat/xilinx/common/include/plat_fdt.h b/plat/xilinx/common/include/plat_fdt.h
new file mode 100644
index 0000000..a1ee1e1
--- /dev/null
+++ b/plat/xilinx/common/include/plat_fdt.h
@@ -0,0 +1,12 @@
+/*
+ * Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+#ifndef PLAT_FDT_H
+#define PLAT_FDT_H
+
+void prepare_dtb(void);
+
+#endif /* PLAT_FDT_H */
diff --git a/plat/xilinx/common/plat_fdt.c b/plat/xilinx/common/plat_fdt.c
new file mode 100644
index 0000000..3d12d51
--- /dev/null
+++ b/plat/xilinx/common/plat_fdt.c
@@ -0,0 +1,50 @@
+/*
+ * Copyright (c) 2023, Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ */
+#include <common/fdt_fixup.h>
+#include <common/fdt_wrappers.h>
+
+#include <plat_fdt.h>
+#include <platform_def.h>
+
+#if (defined(XILINX_OF_BOARD_DTB_ADDR) && !IS_TFA_IN_OCM(BL31_BASE))
+void prepare_dtb(void)
+{
+	void *dtb = (void *)XILINX_OF_BOARD_DTB_ADDR;
+	int ret;
+
+	/* Return if no device tree is detected */
+	if (fdt_check_header(dtb) != 0) {
+		NOTICE("Can't read DT at %p\n", dtb);
+		return;
+	}
+
+	ret = fdt_open_into(dtb, dtb, XILINX_OF_BOARD_DTB_MAX_SIZE);
+	if (ret < 0) {
+		ERROR("Invalid Device Tree at %p: error %d\n", dtb, ret);
+		return;
+	}
+
+	/* Reserve memory used by Trusted Firmware. */
+	if (fdt_add_reserved_memory(dtb, "tf-a", BL31_BASE, BL31_LIMIT - BL31_BASE + 1)) {
+		WARN("Failed to add reserved memory nodes for BL31 to DT.\n");
+		return;
+	}
+
+	ret = fdt_pack(dtb);
+	if (ret < 0) {
+		ERROR("Failed to pack Device Tree at %p: error %d\n", dtb, ret);
+		return;
+	}
+
+	clean_dcache_range((uintptr_t)dtb, fdt_blob_size(dtb));
+	INFO("Changed device tree to advertise PSCI and reserved memories.\n");
+}
+#else
+void prepare_dtb(void)
+{
+}
+#endif
diff --git a/plat/xilinx/versal/bl31_versal_setup.c b/plat/xilinx/versal/bl31_versal_setup.c
index 0ef92a6..6cb529b 100644
--- a/plat/xilinx/versal/bl31_versal_setup.c
+++ b/plat/xilinx/versal/bl31_versal_setup.c
@@ -20,6 +20,7 @@
 #include <plat/common/platform.h>
 #include <plat_arm.h>
 
+#include <plat_fdt.h>
 #include <plat_private.h>
 #include <plat_startup.h>
 #include "pm_api_sys.h"
@@ -191,8 +192,11 @@
 
 	return 0;
 }
+
 void bl31_platform_setup(void)
 {
+	prepare_dtb();
+
 	/* Initialize the gic cpu and distributor interfaces */
 	plat_versal_gic_driver_init();
 	plat_versal_gic_init();
@@ -220,6 +224,10 @@
 	plat_arm_interconnect_enter_coherency();
 
 	const mmap_region_t bl_regions[] = {
+#if (defined(XILINX_OF_BOARD_DTB_ADDR) && !IS_TFA_IN_OCM(BL31_BASE))
+		MAP_REGION_FLAT(XILINX_OF_BOARD_DTB_ADDR, XILINX_OF_BOARD_DTB_MAX_SIZE,
+				MT_MEMORY | MT_RW | MT_NS),
+#endif
 		MAP_REGION_FLAT(BL31_BASE, BL31_END - BL31_BASE,
 			MT_MEMORY | MT_RW | MT_SECURE),
 		MAP_REGION_FLAT(BL_CODE_BASE, BL_CODE_END - BL_CODE_BASE,
diff --git a/plat/xilinx/versal/include/platform_def.h b/plat/xilinx/versal/include/platform_def.h
index 6c1d8b6..4c0df4f 100644
--- a/plat/xilinx/versal/include/platform_def.h
+++ b/plat/xilinx/versal/include/platform_def.h
@@ -76,8 +76,29 @@
  ******************************************************************************/
 #define PLAT_PHY_ADDR_SPACE_SIZE	(1ull << 32)
 #define PLAT_VIRT_ADDR_SPACE_SIZE	(1ull << 32)
+
+#define XILINX_OF_BOARD_DTB_MAX_SIZE	U(0x200000)
+
+#define PLAT_OCM_BSE			U(0xFFFE0000)
+#define PLAT_OCM_LIMIT			U(0xFFFFFFFF)
+
+#define IS_TFA_IN_OCM(x)	((x >= PLAT_OCM_BASE) && (x < PLAT_OCM_LIMIT))
+
+#ifndef MAX_MMAP_REGIONS
+#if (defined(XILINX_OF_BOARD_DTB_ADDR) && !IS_TFA_IN_OCM(BL31_BASE))
+#define MAX_MMAP_REGIONS		9
+#else
 #define MAX_MMAP_REGIONS		8
-#define MAX_XLAT_TABLES			5
+#endif
+#endif
+
+#ifndef MAX_XLAT_TABLES
+#if !IS_TFA_IN_OCM(BL31_BASE)
+#define MAX_XLAT_TABLES		9
+#else
+#define MAX_XLAT_TABLES		5
+#endif
+#endif
 
 #define CACHE_WRITEBACK_SHIFT	6
 #define CACHE_WRITEBACK_GRANULE	(1 << CACHE_WRITEBACK_SHIFT)
diff --git a/plat/xilinx/versal/platform.mk b/plat/xilinx/versal/platform.mk
index 0b6aea0..c936220 100644
--- a/plat/xilinx/versal/platform.mk
+++ b/plat/xilinx/versal/platform.mk
@@ -45,12 +45,17 @@
 VERSAL_PLATFORM ?= silicon
 $(eval $(call add_define_val,VERSAL_PLATFORM,VERSAL_PLATFORM_ID_${VERSAL_PLATFORM}))
 
+ifdef XILINX_OF_BOARD_DTB_ADDR
+$(eval $(call add_define,XILINX_OF_BOARD_DTB_ADDR))
+endif
+
 PLAT_INCLUDES		:=	-Iinclude/plat/arm/common/			\
 				-Iplat/xilinx/common/include/			\
 				-Iplat/xilinx/common/ipi_mailbox_service/	\
 				-Iplat/xilinx/versal/include/			\
 				-Iplat/xilinx/versal/pm_service/
 
+include lib/libfdt/libfdt.mk
 # Include GICv3 driver files
 include drivers/arm/gic/v3/gicv3.mk
 include lib/xlat_tables_v2/xlat_tables.mk
@@ -80,6 +85,7 @@
 				lib/cpus/aarch64/cortex_a72.S			\
 				plat/common/plat_psci_common.c			\
 				plat/xilinx/common/ipi.c			\
+				plat/xilinx/common/plat_fdt.c			\
 				plat/xilinx/common/plat_startup.c		\
 				plat/xilinx/common/ipi_mailbox_service/ipi_mailbox_svc.c \
 				plat/xilinx/common/pm_service/pm_ipi.c		\
@@ -93,7 +99,9 @@
 				plat/xilinx/versal/sip_svc_setup.c		\
 				plat/xilinx/versal/versal_gicv3.c		\
 				plat/xilinx/versal/versal_ipi.c			\
-				plat/xilinx/versal/pm_service/pm_client.c
+				plat/xilinx/versal/pm_service/pm_client.c	\
+				common/fdt_fixup.c				\
+				${LIBFDT_SRCS}
 
 ifeq ($(HARDEN_SLS_ALL), 1)
 TF_CFLAGS_aarch64      +=      -mharden-sls=all
diff --git a/plat/xilinx/versal_net/bl31_versal_net_setup.c b/plat/xilinx/versal_net/bl31_versal_net_setup.c
index 79205a3..d2b6c7b 100644
--- a/plat/xilinx/versal_net/bl31_versal_net_setup.c
+++ b/plat/xilinx/versal_net/bl31_versal_net_setup.c
@@ -12,17 +12,15 @@
 #include <bl31/bl31.h>
 #include <common/bl_common.h>
 #include <common/debug.h>
-#include <common/fdt_fixup.h>
-#include <common/fdt_wrappers.h>
 #include <drivers/arm/dcc.h>
 #include <drivers/arm/pl011.h>
 #include <drivers/console.h>
 #include <lib/mmio.h>
 #include <lib/xlat_tables/xlat_tables_v2.h>
-#include <libfdt.h>
 #include <plat/common/platform.h>
 #include <plat_arm.h>
 
+#include <plat_fdt.h>
 #include <plat_private.h>
 #include <plat_startup.h>
 #include <pm_api_sys.h>
@@ -226,6 +224,8 @@
 
 void bl31_platform_setup(void)
 {
+	prepare_dtb();
+
 	/* Initialize the gic cpu and distributor interfaces */
 	plat_versal_net_gic_driver_init();
 	plat_versal_net_gic_init();
@@ -250,6 +250,10 @@
 void bl31_plat_arch_setup(void)
 {
 	const mmap_region_t bl_regions[] = {
+#if (defined(XILINX_OF_BOARD_DTB_ADDR) && !IS_TFA_IN_OCM(BL31_BASE))
+		MAP_REGION_FLAT(XILINX_OF_BOARD_DTB_ADDR, XILINX_OF_BOARD_DTB_MAX_SIZE,
+				MT_MEMORY | MT_RW | MT_NS),
+#endif
 		MAP_REGION_FLAT(BL31_BASE, BL31_END - BL31_BASE,
 			MT_MEMORY | MT_RW | MT_SECURE),
 		MAP_REGION_FLAT(BL_CODE_BASE, BL_CODE_END - BL_CODE_BASE,
diff --git a/plat/xilinx/versal_net/include/platform_def.h b/plat/xilinx/versal_net/include/platform_def.h
index b256b05..872b6ee 100644
--- a/plat/xilinx/versal_net/include/platform_def.h
+++ b/plat/xilinx/versal_net/include/platform_def.h
@@ -84,13 +84,25 @@
 
 #define PLAT_PHY_ADDR_SPACE_SIZE	(1ULL << 32U)
 #define PLAT_VIRT_ADDR_SPACE_SIZE	(1ULL << 32U)
-#if (BL31_LIMIT < PLAT_DDR_LOWMEM_MAX)
-#define MAX_MMAP_REGIONS		U(10)
+
+#define XILINX_OF_BOARD_DTB_MAX_SIZE	U(0x200000)
+
+#define PLAT_OCM_BASE			U(0xBBF00000)
+#define PLAT_OCM_LIMIT			U(0xBC000000)
+
+#define IS_TFA_IN_OCM(x)	((x >= PLAT_OCM_BASE) && (x < PLAT_OCM_LIMIT))
+
+#ifndef MAX_MMAP_REGIONS
+#if (defined(XILINX_OF_BOARD_DTB_ADDR) && !IS_TFA_IN_OCM(BL31_BASE))
+#define MAX_MMAP_REGIONS		9
 #else
-#define MAX_MMAP_REGIONS		U(9)
+#define MAX_MMAP_REGIONS		8
+#endif
 #endif
 
-#define MAX_XLAT_TABLES			U(8)
+#ifndef MAX_XLAT_TABLES
+#define MAX_XLAT_TABLES			U(9)
+#endif
 
 #define CACHE_WRITEBACK_SHIFT	U(6)
 #define CACHE_WRITEBACK_GRANULE	(1 << CACHE_WRITEBACK_SHIFT)
diff --git a/plat/xilinx/versal_net/platform.mk b/plat/xilinx/versal_net/platform.mk
index 398ef85..9c4cfa0 100644
--- a/plat/xilinx/versal_net/platform.mk
+++ b/plat/xilinx/versal_net/platform.mk
@@ -65,6 +65,10 @@
 
 $(eval $(call add_define_val,VERSAL_NET_CONSOLE,VERSAL_NET_CONSOLE_ID_${VERSAL_NET_CONSOLE}))
 
+ifdef XILINX_OF_BOARD_DTB_ADDR
+$(eval $(call add_define,XILINX_OF_BOARD_DTB_ADDR))
+endif
+
 PLAT_INCLUDES		:=	-Iinclude/plat/arm/common/			\
 				-Iplat/xilinx/common/include/			\
 				-Iplat/xilinx/common/ipi_mailbox_service/	\
@@ -101,7 +105,8 @@
 else
 BL31_SOURCES		+=	${PLAT_PATH}/plat_psci.c
 endif
-BL31_SOURCES		+=	plat/xilinx/common/plat_startup.c		\
+BL31_SOURCES		+=	plat/xilinx/common/plat_fdt.c			\
+				plat/xilinx/common/plat_startup.c		\
 				plat/xilinx/common/ipi.c			\
 				plat/xilinx/common/ipi_mailbox_service/ipi_mailbox_svc.c \
 				plat/xilinx/common/versal.c			\