}
 
 #ifdef CONFIG_SMP
+/*
+ * move_queued_task - move a queued task to new rq.
+ *
+ * Returns (locked) new rq. Old rq's lock is released.
+ */
+static struct rq *move_queued_task(struct task_struct *p, int new_cpu)
+{
+       struct rq *rq = task_rq(p);
+
+       lockdep_assert_held(&rq->lock);
+
+       dequeue_task(rq, p, 0);
+       p->on_rq = TASK_ON_RQ_MIGRATING;
+       set_task_cpu(p, new_cpu);
+       raw_spin_unlock(&rq->lock);
+
+       rq = cpu_rq(new_cpu);
+
+       raw_spin_lock(&rq->lock);
+       BUG_ON(task_cpu(p) != new_cpu);
+       p->on_rq = TASK_ON_RQ_QUEUED;
+       enqueue_task(rq, p, 0);
+       check_preempt_curr(rq, p, 0);
+
+       return rq;
+}
+
 void do_set_cpus_allowed(struct task_struct *p, const struct cpumask *new_mask)
 {
        if (p->sched_class && p->sched_class->set_cpus_allowed)
                goto out;
 
        dest_cpu = cpumask_any_and(cpu_active_mask, new_mask);
-       if (task_on_rq_queued(p) || p->state == TASK_WAKING) {
+       if (task_running(rq, p) || p->state == TASK_WAKING) {
                struct migration_arg arg = { p, dest_cpu };
                /* Need help from migration thread: drop lock and wait. */
                task_rq_unlock(rq, p, &flags);
                stop_one_cpu(cpu_of(rq), migration_cpu_stop, &arg);
                tlb_migrate_finish(p->mm);
                return 0;
-       }
+       } else if (task_on_rq_queued(p))
+               rq = move_queued_task(p, dest_cpu);
 out:
        task_rq_unlock(rq, p, &flags);
 
         * If we're not on a rq, the next wake-up will ensure we're
         * placed properly.
         */
-       if (task_on_rq_queued(p)) {
-               dequeue_task(rq, p, 0);
-               p->on_rq = TASK_ON_RQ_MIGRATING;
-               set_task_cpu(p, dest_cpu);
-               raw_spin_unlock(&rq->lock);
-
-               rq = cpu_rq(dest_cpu);
-               raw_spin_lock(&rq->lock);
-               BUG_ON(task_rq(p) != rq);
-               p->on_rq = TASK_ON_RQ_QUEUED;
-               enqueue_task(rq, p, 0);
-               check_preempt_curr(rq, p, 0);
-       }
+       if (task_on_rq_queued(p))
+               rq = move_queued_task(p, dest_cpu);
 done:
        ret = 1;
 fail: