#include <linux/module.h>
 #include <linux/of.h>
 #include <linux/slab.h>
+#include <linux/topology.h>
 
 #include <asm/cpuidle.h>
 
        return CPU_PM_CPU_IDLE_ENTER(arm_cpuidle_suspend, idx);
 }
 
-static struct cpuidle_driver arm_idle_driver = {
+static struct cpuidle_driver arm_idle_driver __initdata = {
        .name = "arm_idle",
        .owner = THIS_MODULE,
        /*
 static int __init arm_idle_init(void)
 {
        int cpu, ret;
-       struct cpuidle_driver *drv = &arm_idle_driver;
+       struct cpuidle_driver *drv;
        struct cpuidle_device *dev;
 
-       /*
-        * Initialize idle states data, starting at index 1.
-        * This driver is DT only, if no DT idle states are detected (ret == 0)
-        * let the driver initialization fail accordingly since there is no
-        * reason to initialize the idle driver if only wfi is supported.
-        */
-       ret = dt_init_idle_driver(drv, arm_idle_state_match, 1);
-       if (ret <= 0)
-               return ret ? : -ENODEV;
-
-       ret = cpuidle_register_driver(drv);
-       if (ret) {
-               pr_err("Failed to register cpuidle driver\n");
-               return ret;
-       }
-
-       /*
-        * Call arch CPU operations in order to initialize
-        * idle states suspend back-end specific data
-        */
        for_each_possible_cpu(cpu) {
+
+               drv = kmemdup(&arm_idle_driver, sizeof(*drv), GFP_KERNEL);
+               if (!drv) {
+                       ret = -ENOMEM;
+                       goto out_fail;
+               }
+
+               drv->cpumask = (struct cpumask *)cpumask_of(cpu);
+
+               /*
+                * Initialize idle states data, starting at index 1.  This
+                * driver is DT only, if no DT idle states are detected (ret
+                * == 0) let the driver initialization fail accordingly since
+                * there is no reason to initialize the idle driver if only
+                * wfi is supported.
+                */
+               ret = dt_init_idle_driver(drv, arm_idle_state_match, 1);
+               if (ret <= 0) {
+                       ret = ret ? : -ENODEV;
+                       goto out_fail;
+               }
+
+               ret = cpuidle_register_driver(drv);
+               if (ret) {
+                       pr_err("Failed to register cpuidle driver\n");
+                       goto out_fail;
+               }
+
+               /*
+                * Call arch CPU operations in order to initialize
+                * idle states suspend back-end specific data
+                */
                ret = arm_cpuidle_init(cpu);
 
                /*
                dev = per_cpu(cpuidle_devices, cpu);
                cpuidle_unregister_device(dev);
                kfree(dev);
+               drv = cpuidle_get_driver();
+               cpuidle_unregister_driver(drv);
+               kfree(drv);
        }
 
-       cpuidle_unregister_driver(drv);
-
        return ret;
 }
 device_initcall(arm_idle_init);