#endif
 }
 
+#ifdef CONFIG_SMP
+
+static void do_balance_callbacks(struct rq *rq, struct callback_head *head)
+{
+       void (*func)(struct rq *rq);
+       struct callback_head *next;
+
+       lockdep_assert_held(&rq->lock);
+
+       while (head) {
+               func = (void (*)(struct rq *))head->func;
+               next = head->next;
+               head->next = NULL;
+               head = next;
+
+               func(rq);
+       }
+}
+
+static inline struct callback_head *splice_balance_callbacks(struct rq *rq)
+{
+       struct callback_head *head = rq->balance_callback;
+
+       lockdep_assert_held(&rq->lock);
+       if (head)
+               rq->balance_callback = NULL;
+
+       return head;
+}
+
+static void __balance_callbacks(struct rq *rq)
+{
+       do_balance_callbacks(rq, splice_balance_callbacks(rq));
+}
+
+static inline void balance_callbacks(struct rq *rq, struct callback_head *head)
+{
+       unsigned long flags;
+
+       if (unlikely(head)) {
+               raw_spin_lock_irqsave(&rq->lock, flags);
+               do_balance_callbacks(rq, head);
+               raw_spin_unlock_irqrestore(&rq->lock, flags);
+       }
+}
+
+#else
+
+static inline void __balance_callbacks(struct rq *rq)
+{
+}
+
+static inline struct callback_head *splice_balance_callbacks(struct rq *rq)
+{
+       return NULL;
+}
+
+static inline void balance_callbacks(struct rq *rq, struct callback_head *head)
+{
+}
+
+#endif
+
 static inline void
 prepare_lock_switch(struct rq *rq, struct task_struct *next, struct rq_flags *rf)
 {
         * prev into current:
         */
        spin_acquire(&rq->lock.dep_map, 0, 0, _THIS_IP_);
+       __balance_callbacks(rq);
        raw_spin_unlock_irq(&rq->lock);
 }
 
        return rq;
 }
 
-#ifdef CONFIG_SMP
-
-/* rq->lock is NOT held, but preemption is disabled */
-static void __balance_callback(struct rq *rq)
-{
-       struct callback_head *head, *next;
-       void (*func)(struct rq *rq);
-       unsigned long flags;
-
-       raw_spin_lock_irqsave(&rq->lock, flags);
-       head = rq->balance_callback;
-       rq->balance_callback = NULL;
-       while (head) {
-               func = (void (*)(struct rq *))head->func;
-               next = head->next;
-               head->next = NULL;
-               head = next;
-
-               func(rq);
-       }
-       raw_spin_unlock_irqrestore(&rq->lock, flags);
-}
-
-static inline void balance_callback(struct rq *rq)
-{
-       if (unlikely(rq->balance_callback))
-               __balance_callback(rq);
-}
-
-#else
-
-static inline void balance_callback(struct rq *rq)
-{
-}
-
-#endif
-
 /**
  * schedule_tail - first thing a freshly forked thread must call.
  * @prev: the thread we just switched away from.
         */
 
        rq = finish_task_switch(prev);
-       balance_callback(rq);
        preempt_enable();
 
        if (current->set_child_tid)
                rq = context_switch(rq, prev, next, &rf);
        } else {
                rq->clock_update_flags &= ~(RQCF_ACT_SKIP|RQCF_REQ_SKIP);
-               rq_unlock_irq(rq, &rf);
-       }
 
-       balance_callback(rq);
+               rq_unpin_lock(rq, &rf);
+               __balance_callbacks(rq);
+               raw_spin_unlock_irq(&rq->lock);
+       }
 }
 
 void __noreturn do_task_dead(void)
 out_unlock:
        /* Avoid rq from going away on us: */
        preempt_disable();
-       __task_rq_unlock(rq, &rf);
 
-       balance_callback(rq);
+       rq_unpin_lock(rq, &rf);
+       __balance_callbacks(rq);
+       raw_spin_unlock(&rq->lock);
+
        preempt_enable();
 }
 #else
        int retval, oldprio, oldpolicy = -1, queued, running;
        int new_effective_prio, policy = attr->sched_policy;
        const struct sched_class *prev_class;
+       struct callback_head *head;
        struct rq_flags rf;
        int reset_on_fork;
        int queue_flags = DEQUEUE_SAVE | DEQUEUE_MOVE | DEQUEUE_NOCLOCK;
 
        /* Avoid rq from going away on us: */
        preempt_disable();
+       head = splice_balance_callbacks(rq);
        task_rq_unlock(rq, p, &rf);
 
        if (pi) {
        }
 
        /* Run balance callbacks after we've adjusted the PI chain: */
-       balance_callback(rq);
+       balance_callbacks(rq, head);
        preempt_enable();
 
        return 0;