diff --git a/include/lib/extensions/amu.h b/include/lib/extensions/amu.h
index 09d8dee..a396b99 100644
--- a/include/lib/extensions/amu.h
+++ b/include/lib/extensions/amu.h
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2021, Arm Limited and Contributors. All rights reserved.
+ * Copyright (c) 2017-2023, Arm Limited and Contributors. All rights reserved.
  *
  * SPDX-License-Identifier: BSD-3-Clause
  */
@@ -19,9 +19,11 @@
 void amu_enable(cpu_context_t *ctx);
 void amu_init_el3(void);
 void amu_init_el2_unused(void);
+void amu_enable_per_world(per_world_context_t *per_world_ctx);
 #else
 void amu_enable(bool el2_unused);
-#endif
+#endif /* __aarch64__ */
+
 #else
 #if __aarch64__
 void amu_enable(cpu_context_t *ctx)
@@ -33,12 +35,15 @@
 void amu_init_el2_unused(void)
 {
 }
+void amu_enable_per_world(per_world_context_t *per_world_ctx)
+{
+}
 #else
 static inline void amu_enable(bool el2_unused)
 {
 }
-#endif
-#endif
+#endif /*__aarch64__ */
+#endif /* ENABLE_FEAT_AMU */
 
 #if ENABLE_AMU_AUXILIARY_COUNTERS
 /*
diff --git a/include/lib/extensions/sme.h b/include/lib/extensions/sme.h
index dbefdfc..bd7948e 100644
--- a/include/lib/extensions/sme.h
+++ b/include/lib/extensions/sme.h
@@ -21,23 +21,31 @@
 #define SME_SMCR_LEN_MAX	U(0x1FF)
 
 #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_enable(cpu_context_t *context);
 void sme_disable(cpu_context_t *context);
+void sme_enable_per_world(per_world_context_t *per_world_ctx);
+void sme_disable_per_world(per_world_context_t *per_world_ctx);
 #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_enable(cpu_context_t *context)
+{
+}
 static inline void sme_disable(cpu_context_t *context)
 {
 }
+static inline void sme_enable_per_world(per_world_context_t *per_world_ctx)
+{
+}
+static inline void sme_disable_per_world(per_world_context_t *per_world_ctx)
+{
+}
 #endif /* ENABLE_SME_FOR_NS */
 
 #endif /* SME_H */
diff --git a/include/lib/extensions/sve.h b/include/lib/extensions/sve.h
index fc76a16..947c905 100644
--- a/include/lib/extensions/sve.h
+++ b/include/lib/extensions/sve.h
@@ -10,17 +10,17 @@
 #include <context.h>
 
 #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);
+void sve_enable_per_world(per_world_context_t *per_world_ctx);
+void sve_disable_per_world(per_world_context_t *per_world_ctx);
 #else
-static inline void sve_enable(cpu_context_t *context)
+static inline void sve_init_el2_unused(void)
 {
 }
-static inline void sve_init_el2_unused(void)
+static inline void sve_enable_per_world(per_world_context_t *per_world_ctx)
 {
 }
-static inline void sve_disable(cpu_context_t *context)
+static inline void sve_disable_per_world(per_world_context_t *per_world_ctx)
 {
 }
 #endif /* ( ENABLE_SME_FOR_NS | ENABLE_SVE_FOR_NS ) */
diff --git a/include/lib/extensions/sys_reg_trace.h b/include/lib/extensions/sys_reg_trace.h
index beda88a..7004267 100644
--- a/include/lib/extensions/sys_reg_trace.h
+++ b/include/lib/extensions/sys_reg_trace.h
@@ -12,8 +12,8 @@
 #if ENABLE_SYS_REG_TRACE_FOR_NS
 
 #if __aarch64__
-void sys_reg_trace_enable(cpu_context_t *context);
-void sys_reg_trace_disable(cpu_context_t *context);
+void sys_reg_trace_enable_per_world(per_world_context_t *per_world_ctx);
+void sys_reg_trace_disable_per_world(per_world_context_t *per_world_ctx);
 void sys_reg_trace_init_el2_unused(void);
 #else
 void sys_reg_trace_init_el3(void);
@@ -22,10 +22,10 @@
 #else /* !ENABLE_SYS_REG_TRACE_FOR_NS */
 
 #if __aarch64__
-static inline void sys_reg_trace_enable(cpu_context_t *context)
+static inline void sys_reg_trace_enable_per_world(per_world_context_t *per_world_ctx)
 {
 }
-static inline void sys_reg_trace_disable(cpu_context_t *context)
+static inline void sys_reg_trace_disable_per_world(per_world_context_t *per_world_ctx)
 {
 }
 static inline void sys_reg_trace_init_el2_unused(void)