diff --git a/include/lib/el3_runtime/aarch64/context.h b/include/lib/el3_runtime/aarch64/context.h
index 6f97bed..acf111b 100644
--- a/include/lib/el3_runtime/aarch64/context.h
+++ b/include/lib/el3_runtime/aarch64/context.h
@@ -63,19 +63,20 @@
 #define CTX_ELR_EL3		U(0x20)
 #define CTX_PMCR_EL0		U(0x28)
 #define CTX_IS_IN_EL3		U(0x30)
+#define CTX_MDCR_EL3		U(0x38)
 /* Constants required in supporting nested exception in EL3 */
-#define CTX_SAVED_ELR_EL3	U(0x38)
+#define CTX_SAVED_ELR_EL3	U(0x40)
 /*
  * General purpose flag, to save various EL3 states
  * FFH mode : Used to identify if handling nested exception
  * KFH mode : Used as counter value
  */
-#define CTX_NESTED_EA_FLAG	U(0x40)
+#define CTX_NESTED_EA_FLAG	U(0x48)
 #if FFH_SUPPORT
- #define CTX_SAVED_ESR_EL3	U(0x48)
- #define CTX_SAVED_SPSR_EL3	U(0x50)
- #define CTX_SAVED_GPREG_LR	U(0x58)
- #define CTX_EL3STATE_END	U(0x60) /* Align to the next 16 byte boundary */
+ #define CTX_SAVED_ESR_EL3	U(0x50)
+ #define CTX_SAVED_SPSR_EL3	U(0x58)
+ #define CTX_SAVED_GPREG_LR	U(0x60)
+ #define CTX_EL3STATE_END	U(0x70) /* Align to the next 16 byte boundary */
 #else
  #define CTX_EL3STATE_END	U(0x50) /* Align to the next 16 byte boundary */
 #endif /* FFH_SUPPORT */
diff --git a/include/lib/extensions/brbe.h b/include/lib/extensions/brbe.h
index 194efba..425a037 100644
--- a/include/lib/extensions/brbe.h
+++ b/include/lib/extensions/brbe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2022-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,10 +7,12 @@
 #ifndef BRBE_H
 #define BRBE_H
 
+#include <context.h>
+
 #if ENABLE_BRBE_FOR_NS
-void brbe_init_el3(void);
+void brbe_enable(cpu_context_t *ctx);
 #else
-static inline void brbe_init_el3(void)
+static inline void brbe_enable(cpu_context_t *ctx)
 {
 }
 #endif /* ENABLE_BRBE_FOR_NS */
diff --git a/include/lib/extensions/spe.h b/include/lib/extensions/spe.h
index 7b39037..c6e44f9 100644
--- a/include/lib/extensions/spe.h
+++ b/include/lib/extensions/spe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2023, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2024, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -8,13 +8,14 @@
 #define SPE_H
 
 #include <stdbool.h>
+#include <context.h>
 
 #if ENABLE_SPE_FOR_NS
-void spe_init_el3(void);
+void spe_enable(cpu_context_t *ctx);
 void spe_init_el2_unused(void);
 void spe_disable(void);
 #else
-static inline void spe_init_el3(void)
+static inline void spe_enable(cpu_context_t *ctx)
 {
 }
 static inline void spe_init_el2_unused(void)
diff --git a/include/lib/extensions/trbe.h b/include/lib/extensions/trbe.h
index 0bed433..5db3316 100644
--- a/include/lib/extensions/trbe.h
+++ b/include/lib/extensions/trbe.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,11 +7,13 @@
 #ifndef TRBE_H
 #define TRBE_H
 
+#include <context.h>
+
 #if ENABLE_TRBE_FOR_NS
-void trbe_init_el3(void);
+void trbe_enable(cpu_context_t *ctx);
 void trbe_init_el2_unused(void);
 #else
-static inline void trbe_init_el3(void)
+static inline void trbe_enable(cpu_context_t *ctx)
 {
 }
 static inline void trbe_init_el2_unused(void)
diff --git a/include/lib/extensions/trf.h b/include/lib/extensions/trf.h
index 1ac7cda..f0a946b 100644
--- a/include/lib/extensions/trf.h
+++ b/include/lib/extensions/trf.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, Arm Limited. All rights reserved.
+ * Copyright (c) 2021-2024, Arm Limited. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -7,16 +7,32 @@
 #ifndef TRF_H
 #define TRF_H
 
+#include <context.h>
+
 #if ENABLE_TRF_FOR_NS
-void trf_init_el3(void);
+
+#if __aarch64__
+void trf_enable(cpu_context_t *ctx);
 void trf_init_el2_unused(void);
-#else
-static inline void trf_init_el3(void)
+#else /* !__aarch64 */
+void trf_init_el3(void);
+#endif /* __aarch64__ */
+
+#else /* ENABLE_TRF_FOR_NS=0 */
+
+#if __aarch64__
+static inline void trf_enable(cpu_context_t *ctx)
 {
 }
 static inline void trf_init_el2_unused(void)
 {
 }
+#else /* !__aarch64 */
+static inline void trf_init_el3(void)
+{
+}
+#endif /* __aarch64__*/
+
 #endif /* ENABLE_TRF_FOR_NS */
 
 #endif /* TRF_H */