}
 EXPORT_SYMBOL(unregister_reboot_notifier);
 
+/* Add backwards compatibility for stable trees. */
+#ifndef PF_NO_SETAFFINITY
+#define PF_NO_SETAFFINITY              PF_THREAD_BOUND
+#endif
+
+static void migrate_to_reboot_cpu(void)
+{
+       /* The boot cpu is always logical cpu 0 */
+       int cpu = 0;
+
+       cpu_hotplug_disable();
+
+       /* Make certain the cpu I'm about to reboot on is online */
+       if (!cpu_online(cpu))
+               cpu = cpumask_first(cpu_online_mask);
+
+       /* Prevent races with other tasks migrating this task */
+       current->flags |= PF_NO_SETAFFINITY;
+
+       /* Make certain I only run on the appropriate processor */
+       set_cpus_allowed_ptr(current, cpumask_of(cpu));
+}
+
 /**
  *     kernel_restart - reboot the system
  *     @cmd: pointer to buffer containing command to execute for restart
 void kernel_restart(char *cmd)
 {
        kernel_restart_prepare(cmd);
-       disable_nonboot_cpus();
+       migrate_to_reboot_cpu();
        syscore_shutdown();
        if (!cmd)
                printk(KERN_EMERG "Restarting system.\n");
 void kernel_halt(void)
 {
        kernel_shutdown_prepare(SYSTEM_HALT);
-       disable_nonboot_cpus();
+       migrate_to_reboot_cpu();
        syscore_shutdown();
        printk(KERN_EMERG "System halted.\n");
        kmsg_dump(KMSG_DUMP_HALT);
        kernel_shutdown_prepare(SYSTEM_POWER_OFF);
        if (pm_power_off_prepare)
                pm_power_off_prepare();
-       disable_nonboot_cpus();
+       migrate_to_reboot_cpu();
        syscore_shutdown();
        printk(KERN_EMERG "Power down.\n");
        kmsg_dump(KMSG_DUMP_POWEROFF);