diff --git a/lib/el3_runtime/aarch64/context.S b/lib/el3_runtime/aarch64/context.S
index 143da95..620ec16 100644
--- a/lib/el3_runtime/aarch64/context.S
+++ b/lib/el3_runtime/aarch64/context.S
@@ -9,7 +9,6 @@
 #include <context.h>
 
 	.global	el1_sysregs_context_save
-	.global el1_sysregs_context_save_post_ops
 	.global	el1_sysregs_context_restore
 #if CTX_INCLUDE_FPREGS
 	.global	fpregs_context_save
@@ -112,36 +111,6 @@
 /* -----------------------------------------------------
  * The following function strictly follows the AArch64
  * PCS to use x9-x17 (temporary caller-saved registers)
- * to do post operations after saving the EL1 system
- * register context.
- * -----------------------------------------------------
- */
-func el1_sysregs_context_save_post_ops
-#if ENABLE_SPE_FOR_LOWER_ELS
-	/* Detect if SPE is implemented */
-	mrs	x9, id_aa64dfr0_el1
-	ubfx	x9, x9, #ID_AA64DFR0_PMS_SHIFT, #ID_AA64DFR0_PMS_LENGTH
-	cmp	x9, #0x1
-	b.ne	1f
-
-	/*
-	 * Before switching from normal world to secure world
-	 * the profiling buffers need to be drained out to memory.  This is
-	 * required to avoid an invalid memory access when TTBR is switched
-	 * for entry to SEL1.
-	 */
-	.arch	armv8.2-a+profile
-	psb	csync
-	dsb	nsh
-	.arch	armv8-a
-1:
-#endif
-	ret
-endfunc el1_sysregs_context_save_post_ops
-
-/* -----------------------------------------------------
- * The following function strictly follows the AArch64
- * PCS to use x9-x17 (temporary caller-saved registers)
  * to restore EL1 system register context.  It assumes
  * that 'x0' is pointing to a 'el1_sys_regs' structure
  * from where the register context will be restored
diff --git a/lib/el3_runtime/aarch64/context_mgmt.c b/lib/el3_runtime/aarch64/context_mgmt.c
index 479acc9..8f1523f 100644
--- a/lib/el3_runtime/aarch64/context_mgmt.c
+++ b/lib/el3_runtime/aarch64/context_mgmt.c
@@ -15,6 +15,7 @@
 #include <platform_def.h>
 #include <pubsub_events.h>
 #include <smcc_helpers.h>
+#include <spe.h>
 #include <string.h>
 #include <utils.h>
 
@@ -216,6 +217,9 @@
 static void enable_extensions_nonsecure(int el2_unused)
 {
 #if IMAGE_BL31
+#if ENABLE_SPE_FOR_LOWER_ELS
+	spe_enable(el2_unused);
+#endif
 #endif
 }
 
@@ -354,13 +358,6 @@
 			 * relying on hw. Some fields are architecturally
 			 * UNKNOWN on reset.
 			 *
-			 * 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.
-			 *
 			 * 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.
@@ -397,22 +394,6 @@
 					| MDCR_EL2_HPME_BIT | MDCR_EL2_TPM_BIT
 					| MDCR_EL2_TPMCR_BIT));
 
-#if ENABLE_SPE_FOR_LOWER_ELS
-			uint64_t id_aa64dfr0_el1;
-
-			/* Detect if SPE is implemented */
-			id_aa64dfr0_el1 = read_id_aa64dfr0_el1() >>
-				ID_AA64DFR0_PMS_SHIFT;
-			if ((id_aa64dfr0_el1 & ID_AA64DFR0_PMS_MASK) == 1) {
-				/*
-				 * Make sure traps to EL2 are not generated if
-				 * EL2 is implemented but not used.
-				 */
-				mdcr_el2 &= ~MDCR_EL2_TPMS;
-				mdcr_el2 |= MDCR_EL2_E2PB(MDCR_EL2_E2PB_EL1);
-			}
-#endif
-
 			write_mdcr_el2(mdcr_el2);
 
 			/*
@@ -454,7 +435,6 @@
 	assert(ctx);
 
 	el1_sysregs_context_save(get_sysregs_ctx(ctx));
-	el1_sysregs_context_save_post_ops();
 
 #if IMAGE_BL31
 	if (security_state == SECURE)
diff --git a/lib/extensions/spe/spe.c b/lib/extensions/spe/spe.c
new file mode 100644
index 0000000..3b297f2
--- /dev/null
+++ b/lib/extensions/spe/spe.c
@@ -0,0 +1,85 @@
+/*
+ * Copyright (c) 2017, ARM Limited and Contributors. All rights reserved.
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#include <arch.h>
+#include <arch_helpers.h>
+#include <pubsub.h>
+
+/*
+ * The assembler does not yet understand the psb csync mnemonic
+ * so use the equivalent hint instruction.
+ */
+#define psb_csync()	asm volatile("hint #17")
+
+void spe_enable(int el2_unused)
+{
+	uint64_t features;
+
+	features = read_id_aa64dfr0_el1() >> ID_AA64DFR0_PMS_SHIFT;
+	if ((features & ID_AA64DFR0_PMS_MASK) == 1) {
+		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
+		 * S-EL1 generate trap exceptions to EL3.
+		 */
+		v = read_mdcr_el3();
+		v |= MDCR_NSPB(MDCR_NSPB_EL1);
+		write_mdcr_el3(v);
+	}
+}
+
+void spe_disable(void)
+{
+	uint64_t features;
+
+	features = read_id_aa64dfr0_el1() >> ID_AA64DFR0_PMS_SHIFT;
+	if ((features & ID_AA64DFR0_PMS_MASK) == 1) {
+		uint64_t v;
+
+		/* Drain buffered data */
+		psb_csync();
+		dsbnsh();
+
+		/* Disable profiling buffer */
+		v = read_pmblimitr_el1();
+		v &= ~(1ULL << 0);
+		write_pmblimitr_el1(v);
+		isb();
+	}
+}
+
+static void *spe_drain_buffers_hook(const void *arg)
+{
+	uint64_t features;
+
+	features = read_id_aa64dfr0_el1() >> ID_AA64DFR0_PMS_SHIFT;
+	if ((features & ID_AA64DFR0_PMS_MASK) == 1) {
+		/* Drain buffered data */
+		psb_csync();
+		dsbnsh();
+	}
+
+	return 0;
+}
+
+SUBSCRIBE_TO_EVENT(cm_entering_secure_world, spe_drain_buffers_hook);