1 files changed, 13 insertions, 2 deletions

diff --git a/arch/x86/oprofile/nmi_int.c b/arch/x86/oprofile/nmi_int.c
index 6890d8498e0b..379e8bd0deea 100644
--- a/arch/x86/oprofile/nmi_int.c
+++ b/arch/x86/oprofile/nmi_int.c
@@ -494,14 +494,19 @@ static int nmi_setup(void)
         if (err)
                 goto fail;
 
+        cpu_notifier_register_begin();
+
+        /* Use get/put_online_cpus() to protect 'nmi_enabled' */
         get_online_cpus();
-        register_cpu_notifier(&oprofile_cpu_nb);
         nmi_enabled = 1;
         /* make nmi_enabled visible to the nmi handler: */
         smp_mb();
         on_each_cpu(nmi_cpu_setup, NULL, 1);
+        __register_cpu_notifier(&oprofile_cpu_nb);
         put_online_cpus();
 
+        cpu_notifier_register_done();
+
         return 0;
 fail:
         free_msrs();
@@ -512,12 +517,18 @@ static void nmi_shutdown(void)
 {
         struct op_msrs *msrs;
 
+        cpu_notifier_register_begin();
+
+        /* Use get/put_online_cpus() to protect 'nmi_enabled' & 'ctr_running' */
         get_online_cpus();
-        unregister_cpu_notifier(&oprofile_cpu_nb);
         on_each_cpu(nmi_cpu_shutdown, NULL, 1);
         nmi_enabled = 0;
         ctr_running = 0;
+        __unregister_cpu_notifier(&oprofile_cpu_nb);
         put_online_cpus();
+
+        cpu_notifier_register_done();
+
         /* make variables visible to the nmi handler: */
         smp_mb();
         unregister_nmi_handler(NMI_LOCAL, "oprofile");