pid_t machine__get_current_tid(struct machine *machine, int cpu)
 {
-       if (cpu < 0 || cpu >= MAX_NR_CPUS || !machine->current_tid)
+       int nr_cpus = min(machine->env->nr_cpus_online, MAX_NR_CPUS);
+
+       if (cpu < 0 || cpu >= nr_cpus || !machine->current_tid)
                return -1;
 
        return machine->current_tid[cpu];
                             pid_t tid)
 {
        struct thread *thread;
+       int nr_cpus = min(machine->env->nr_cpus_online, MAX_NR_CPUS);
 
        if (cpu < 0)
                return -EINVAL;
        if (!machine->current_tid) {
                int i;
 
-               machine->current_tid = calloc(MAX_NR_CPUS, sizeof(pid_t));
+               machine->current_tid = calloc(nr_cpus, sizeof(pid_t));
                if (!machine->current_tid)
                        return -ENOMEM;
-               for (i = 0; i < MAX_NR_CPUS; i++)
+               for (i = 0; i < nr_cpus; i++)
                        machine->current_tid[i] = -1;
        }
 
-       if (cpu >= MAX_NR_CPUS) {
+       if (cpu >= nr_cpus) {
                pr_err("Requested CPU %d too large. ", cpu);
                pr_err("Consider raising MAX_NR_CPUS\n");
                return -EINVAL;