#include <linux/of.h>
 #include <linux/reboot.h>
 #include <linux/slab.h>
+#include <linux/cpu.h>
 
 #include <asm/cputhreads.h>
 #include <asm/firmware.h>
 {
        struct chip *chip = container_of(work, struct chip, throttle);
        unsigned int cpu;
-       cpumask_var_t mask;
+       cpumask_t mask;
 
-       smp_call_function_any(&chip->mask,
+       get_online_cpus();
+       cpumask_and(&mask, &chip->mask, cpu_online_mask);
+       smp_call_function_any(&mask,
                              powernv_cpufreq_throttle_check, NULL, 0);
 
        if (!chip->restore)
-               return;
+               goto out;
 
        chip->restore = false;
-       cpumask_copy(mask, &chip->mask);
-       for_each_cpu_and(cpu, mask, cpu_online_mask) {
-               int index, tcpu;
+       for_each_cpu(cpu, &mask) {
+               int index;
                struct cpufreq_policy policy;
 
                cpufreq_get_policy(&policy, cpu);
                                               policy.cur,
                                               CPUFREQ_RELATION_C, &index);
                powernv_cpufreq_target_index(&policy, index);
-               for_each_cpu(tcpu, policy.cpus)
-                       cpumask_clear_cpu(tcpu, mask);
+               cpumask_andnot(&mask, &mask, policy.cpus);
        }
+out:
+       put_online_cpus();
 }
 
 static char throttle_reason[][30] = {