diff --git a/lib/extensions/brbe/brbe.c b/lib/extensions/brbe/brbe.c
index 37bd834..dde0266 100644
--- a/lib/extensions/brbe/brbe.c
+++ b/lib/extensions/brbe/brbe.c
@@ -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
  */
@@ -9,9 +9,10 @@
 #include <arch_helpers.h>
 #include <lib/extensions/brbe.h>
 
-void brbe_init_el3(void)
+void brbe_enable(cpu_context_t *ctx)
 {
-	uint64_t val;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.SBRBE = 0b01
@@ -19,8 +20,7 @@
 	 * Allows BRBE usage in non-secure world and prohibited in
 	 * secure world.
 	 */
-	val = read_mdcr_el3();
-	val &= ~(MDCR_SBRBE_MASK << MDCR_SBRBE_SHIFT);
-	val |= (0x1UL << MDCR_SBRBE_SHIFT);
-	write_mdcr_el3(val);
+	mdcr_el3_val &= ~(MDCR_SBRBE_MASK << MDCR_SBRBE_SHIFT);
+	mdcr_el3_val |= (0x1UL << MDCR_SBRBE_SHIFT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
diff --git a/lib/extensions/spe/spe.c b/lib/extensions/spe/spe.c
index d1fb182..c6076fe 100644
--- a/lib/extensions/spe/spe.c
+++ b/lib/extensions/spe/spe.c
@@ -29,9 +29,10 @@
 	__asm__ volatile("hint #17");
 }
 
-void spe_init_el3(void)
+void spe_enable(cpu_context_t *ctx)
 {
-	uint64_t v;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.NSPB (ARM v8.2): SPE enabled in Non-secure state
@@ -46,10 +47,9 @@
 	 * Setting this bit to 1 doesn't have any effect on it when
 	 * FEAT_SPEv1p2 not implemented.
 	 */
-	v = read_mdcr_el3();
-	v |= MDCR_NSPB(MDCR_NSPB_EL1) | MDCR_EnPMSN_BIT;
-	v &= ~(MDCR_NSPBE_BIT);
-	write_mdcr_el3(v);
+	mdcr_el3_val |= MDCR_NSPB(MDCR_NSPB_EL1) | MDCR_EnPMSN_BIT;
+	mdcr_el3_val &= ~(MDCR_NSPBE_BIT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
 
 void spe_init_el2_unused(void)
diff --git a/lib/extensions/trbe/trbe.c b/lib/extensions/trbe/trbe.c
index d4fbdfb..9157734 100644
--- a/lib/extensions/trbe/trbe.c
+++ b/lib/extensions/trbe/trbe.c
@@ -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
  */
@@ -19,9 +19,10 @@
 	__asm__ volatile("hint #18");
 }
 
-void trbe_init_el3(void)
+void trbe_enable(cpu_context_t *ctx)
 {
-	u_register_t val;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.NSTBE = 0b0
@@ -33,10 +34,9 @@
 	 *  NS-EL2, tracing is prohibited in Secure and Realm state (if
 	 *  implemented).
 	 */
-	val = read_mdcr_el3();
-	val |= MDCR_NSTB(MDCR_NSTB_EL1);
-	val &= ~(MDCR_NSTBE_BIT);
-	write_mdcr_el3(val);
+	mdcr_el3_val |= MDCR_NSTB(MDCR_NSTB_EL1);
+	mdcr_el3_val &= ~(MDCR_NSTBE_BIT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
 
 void trbe_init_el2_unused(void)
diff --git a/lib/extensions/trf/aarch64/trf.c b/lib/extensions/trf/aarch64/trf.c
index 83fbf85..d36853a 100644
--- a/lib/extensions/trf/aarch64/trf.c
+++ b/lib/extensions/trf/aarch64/trf.c
@@ -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
  */
@@ -9,9 +9,10 @@
 #include <arch_helpers.h>
 #include <lib/extensions/trf.h>
 
-void trf_init_el3(void)
+void trf_enable(cpu_context_t *ctx)
 {
-	u_register_t val;
+	el3_state_t *state = get_el3state_ctx(ctx);
+	u_register_t mdcr_el3_val = read_ctx_reg(state, CTX_MDCR_EL3);
 
 	/*
 	 * MDCR_EL3.STE = b0
@@ -22,9 +23,8 @@
 	 * Allow access of trace filter control registers from NS-EL2
 	 * and NS-EL1 when NS-EL2 is implemented but not used
 	 */
-	val = read_mdcr_el3();
-	val &= ~(MDCR_STE_BIT | MDCR_TTRF_BIT);
-	write_mdcr_el3(val);
+	mdcr_el3_val &= ~(MDCR_STE_BIT | MDCR_TTRF_BIT);
+	write_ctx_reg(state, CTX_MDCR_EL3, mdcr_el3_val);
 }
 
 void trf_init_el2_unused(void)