wrmsrl(MSR_KVM_POLL_CONTROL, 1);
 }
 
-void arch_haltpoll_enable(void)
+void arch_haltpoll_enable(unsigned int cpu)
 {
        if (!kvm_para_has_feature(KVM_FEATURE_POLL_CONTROL)) {
-               printk(KERN_ERR "kvm: host does not support poll control\n");
-               printk(KERN_ERR "kvm: host upgrade recommended\n");
+               pr_err_once("kvm: host does not support poll control\n");
+               pr_err_once("kvm: host upgrade recommended\n");
                return;
        }
 
-       preempt_disable();
        /* Enable guest halt poll disables host halt poll */
-       kvm_disable_host_haltpoll(NULL);
-       smp_call_function(kvm_disable_host_haltpoll, NULL, 1);
-       preempt_enable();
+       smp_call_function_single(cpu, kvm_disable_host_haltpoll, NULL, 1);
 }
 EXPORT_SYMBOL_GPL(arch_haltpoll_enable);
 
-void arch_haltpoll_disable(void)
+void arch_haltpoll_disable(unsigned int cpu)
 {
        if (!kvm_para_has_feature(KVM_FEATURE_POLL_CONTROL))
                return;
 
-       preempt_disable();
        /* Enable guest halt poll disables host halt poll */
-       kvm_enable_host_haltpoll(NULL);
-       smp_call_function(kvm_enable_host_haltpoll, NULL, 1);
-       preempt_enable();
+       smp_call_function_single(cpu, kvm_enable_host_haltpoll, NULL, 1);
 }
 EXPORT_SYMBOL_GPL(arch_haltpoll_disable);
 #endif
 
  */
 
 #include <linux/init.h>
+#include <linux/cpu.h>
 #include <linux/cpuidle.h>
 #include <linux/module.h>
 #include <linux/sched/idle.h>
 #include <linux/kvm_para.h>
 #include <linux/cpuidle_haltpoll.h>
 
+static struct cpuidle_device __percpu *haltpoll_cpuidle_devices;
+static enum cpuhp_state haltpoll_hp_state;
+
 static int default_enter_idle(struct cpuidle_device *dev,
                              struct cpuidle_driver *drv, int index)
 {
        .state_count = 2,
 };
 
+static int haltpoll_cpu_online(unsigned int cpu)
+{
+       struct cpuidle_device *dev;
+
+       dev = per_cpu_ptr(haltpoll_cpuidle_devices, cpu);
+       if (!dev->registered) {
+               dev->cpu = cpu;
+               if (cpuidle_register_device(dev)) {
+                       pr_notice("cpuidle_register_device %d failed!\n", cpu);
+                       return -EIO;
+               }
+               arch_haltpoll_enable(cpu);
+       }
+
+       return 0;
+}
+
+static int haltpoll_cpu_offline(unsigned int cpu)
+{
+       struct cpuidle_device *dev;
+
+       dev = per_cpu_ptr(haltpoll_cpuidle_devices, cpu);
+       if (dev->registered) {
+               arch_haltpoll_disable(cpu);
+               cpuidle_unregister_device(dev);
+       }
+
+       return 0;
+}
+
+static void haltpoll_uninit(void)
+{
+       if (haltpoll_hp_state)
+               cpuhp_remove_state(haltpoll_hp_state);
+       cpuidle_unregister_driver(&haltpoll_driver);
+
+       free_percpu(haltpoll_cpuidle_devices);
+       haltpoll_cpuidle_devices = NULL;
+}
+
 static int __init haltpoll_init(void)
 {
        int ret;
        if (!kvm_para_available())
                return 0;
 
-       ret = cpuidle_register(&haltpoll_driver, NULL);
-       if (ret == 0)
-               arch_haltpoll_enable();
+       ret = cpuidle_register_driver(drv);
+       if (ret < 0)
+               return ret;
+
+       haltpoll_cpuidle_devices = alloc_percpu(struct cpuidle_device);
+       if (haltpoll_cpuidle_devices == NULL) {
+               cpuidle_unregister_driver(drv);
+               return -ENOMEM;
+       }
+
+       ret = cpuhp_setup_state(CPUHP_AP_ONLINE_DYN, "cpuidle/haltpoll:online",
+                               haltpoll_cpu_online, haltpoll_cpu_offline);
+       if (ret < 0) {
+               haltpoll_uninit();
+       } else {
+               haltpoll_hp_state = ret;
+               ret = 0;
+       }
 
        return ret;
 }
 
 static void __exit haltpoll_exit(void)
 {
-       arch_haltpoll_disable();
-       cpuidle_unregister(&haltpoll_driver);
+       haltpoll_uninit();
 }
 
 module_init(haltpoll_init);