stop_one_cpu(cpu_of(rq), migration_cpu_stop, &arg);
                tlb_migrate_finish(p->mm);
                return 0;
-       } else if (task_on_rq_queued(p))
+       } else if (task_on_rq_queued(p)) {
+               /*
+                * OK, since we're going to drop the lock immediately
+                * afterwards anyway.
+                */
+               lockdep_unpin_lock(&rq->lock);
                rq = move_queued_task(rq, p, dest_cpu);
+               lockdep_pin_lock(&rq->lock);
+       }
 out:
        task_rq_unlock(rq, p, &flags);
 
 static inline
 int select_task_rq(struct task_struct *p, int cpu, int sd_flags, int wake_flags)
 {
+       lockdep_assert_held(&p->pi_lock);
+
        if (p->nr_cpus_allowed > 1)
                cpu = p->sched_class->select_task_rq(p, cpu, sd_flags, wake_flags);
 
 #ifdef CONFIG_SMP
        if (p->sched_class->task_woken) {
                /*
-                * XXX can drop rq->lock; most likely ok.
+                * Our task @p is fully woken up and running; so its safe to
+                * drop the rq->lock, hereafter rq is only used for statistics.
                 */
+               lockdep_unpin_lock(&rq->lock);
                p->sched_class->task_woken(rq, p);
+               lockdep_pin_lock(&rq->lock);
        }
 
        if (rq->idle_stamp) {
 static void
 ttwu_do_activate(struct rq *rq, struct task_struct *p, int wake_flags)
 {
+       lockdep_assert_held(&rq->lock);
+
 #ifdef CONFIG_SMP
        if (p->sched_contributes_to_load)
                rq->nr_uninterruptible--;
                return;
 
        raw_spin_lock_irqsave(&rq->lock, flags);
+       lockdep_pin_lock(&rq->lock);
 
        while (llist) {
                p = llist_entry(llist, struct task_struct, wake_entry);
                ttwu_do_activate(rq, p, 0);
        }
 
+       lockdep_unpin_lock(&rq->lock);
        raw_spin_unlock_irqrestore(&rq->lock, flags);
 }
 
 #endif
 
        raw_spin_lock(&rq->lock);
+       lockdep_pin_lock(&rq->lock);
        ttwu_do_activate(rq, p, 0);
+       lockdep_unpin_lock(&rq->lock);
        raw_spin_unlock(&rq->lock);
 }
 
        lockdep_assert_held(&rq->lock);
 
        if (!raw_spin_trylock(&p->pi_lock)) {
+               /*
+                * This is OK, because current is on_cpu, which avoids it being
+                * picked for load-balance and preemption/IRQs are still
+                * disabled avoiding further scheduler activity on it and we've
+                * not yet picked a replacement task.
+                */
+               lockdep_unpin_lock(&rq->lock);
                raw_spin_unlock(&rq->lock);
                raw_spin_lock(&p->pi_lock);
                raw_spin_lock(&rq->lock);
+               lockdep_pin_lock(&rq->lock);
        }
 
        if (!(p->state & TASK_NORMAL))
         * of the scheduler it's an obvious special-case), so we
         * do an early lockdep release here:
         */
+       lockdep_unpin_lock(&rq->lock);
        spin_release(&rq->lock.dep_map, 1, _THIS_IP_);
 
        context_tracking_task_switch(prev, next);
         */
        smp_mb__before_spinlock();
        raw_spin_lock_irq(&rq->lock);
+       lockdep_pin_lock(&rq->lock);
 
        rq->clock_skip_update <<= 1; /* promote REQ to ACT */
 
 
                rq = context_switch(rq, prev, next); /* unlocks the rq */
                cpu = cpu_of(rq);
-       } else
+       } else {
+               lockdep_unpin_lock(&rq->lock);
                raw_spin_unlock_irq(&rq->lock);
+       }
 
        balance_callback(rq);
 }
                if (rq->nr_running == 1)
                        break;
 
+               /*
+                * Ensure rq->lock covers the entire task selection
+                * until the migration.
+                */
+               lockdep_pin_lock(&rq->lock);
                next = pick_next_task(rq, &fake_task);
                BUG_ON(!next);
                next->sched_class->put_prev_task(rq, next);
                /* Find suitable destination for @next, with force if needed. */
                dest_cpu = select_fallback_rq(dead_rq->cpu, next);
 
+               lockdep_unpin_lock(&rq->lock);
                rq = __migrate_task(rq, next, dest_cpu);
                if (rq != dead_rq) {
                        raw_spin_unlock(&rq->lock);
 
        for (;;) {
                rq = task_rq(p);
                raw_spin_lock(&rq->lock);
-               if (likely(rq == task_rq(p) && !task_on_rq_migrating(p)))
+               if (likely(rq == task_rq(p) && !task_on_rq_migrating(p))) {
+                       lockdep_pin_lock(&rq->lock);
                        return rq;
+               }
                raw_spin_unlock(&rq->lock);
 
                while (unlikely(task_on_rq_migrating(p)))
                 * If we observe the new cpu in task_rq_lock, the acquire will
                 * pair with the WMB to ensure we must then also see migrating.
                 */
-               if (likely(rq == task_rq(p) && !task_on_rq_migrating(p)))
+               if (likely(rq == task_rq(p) && !task_on_rq_migrating(p))) {
+                       lockdep_pin_lock(&rq->lock);
                        return rq;
+               }
                raw_spin_unlock(&rq->lock);
                raw_spin_unlock_irqrestore(&p->pi_lock, *flags);
 
 static inline void __task_rq_unlock(struct rq *rq)
        __releases(rq->lock)
 {
+       lockdep_unpin_lock(&rq->lock);
        raw_spin_unlock(&rq->lock);
 }
 
        __releases(rq->lock)
        __releases(p->pi_lock)
 {
+       lockdep_unpin_lock(&rq->lock);
        raw_spin_unlock(&rq->lock);
        raw_spin_unlock_irqrestore(&p->pi_lock, *flags);
 }