[PATCH 0/7] sched: core-sched fixes and balancing

[PATCH 0/7] sched: core-sched fixes and balancing

[Peter Zijlstra]

Hi!

Previous 'version' is at [1], but this one is vastly expanded, so I didn't
really bother keeping count.

Anyway, all this started with me breaking core-scheduling [2] and then trying
to fix it all up. Because holidays, this all got delayed, and then TJ had a few
patches fixing core-sched for ext [3] (which, if/when this lands, can be
simplified again).

Anyhow, these patches survive both Aaron's test case [4] and Prateek's [5].


[1] https://lore.kernel.org/r/20260624121327.190063948@infradead.org
[2] https://lore.kernel.org/all/20260611113219.GG187714@noisy.programming.kicks-ass.net/
[3] https://lore.kernel.org/r/20260807210221.232543-1-tj@kernel.org
[4] https://patch.msgid.link/3898f82d-e48a-48d5-a13d-54fa5387a203%40amd.com
[5] https://patch.msgid.link/20260821024335.GA3046103%40bytedance.com

[PATCH 1/7] sched/core: Fix pick_next_task() self recursion

[Peter Zijlstra]

It is possible for another sibling to end up in pick_next_task() when:

  pick_next_task()
    pick_task()
      sched_class::pick_task()

drops the core wide rq->lock. In this case they end up trampling the core wide
task selection state, possibly leading to NULL derefs. Detect this case by
keeping a local copy of core_task_seq, a value that is incremented on
{en,de}queue and schedule.

Since RETRY_TASK is only possible when a higher priority task gets enqueued
during the lock break, this must mean core_task_seq will also be incremented
and is thus completely covered by the seq number mismatch.

XXX words on forward progress go here

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c |   13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6232,14 +6232,14 @@ static struct task_struct *
 pick_next_task(struct rq *rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(rq))
 {
+	bool core_clock_updated = (rq == rq->core);
 	struct task_struct *next, *p, *max;
 	const struct cpumask *smt_mask;
+	int i, cpu, seq, occ = 0;
 	bool fi_before = false;
-	bool core_clock_updated = (rq == rq->core);
+	bool need_sync = false;
 	unsigned long cookie;
-	int i, cpu, occ = 0;
 	struct rq *rq_i;
-	bool need_sync = false;
 
 	if (!sched_core_enabled(rq))
 		return __pick_next_task(rq, rf);
@@ -6314,7 +6314,7 @@ pick_next_task(struct rq *rq, struct rq_
 	 * However, preemptions can cause multiple picks on the same task set.
 	 * 'Fix' this by also increasing @task_seq for every pick.
 	 */
-	rq->core->core_task_seq++;
+	seq = ++rq->core->core_task_seq;
 
 	/*
 	 * Optimize for common case where this CPU has no cookies
@@ -6362,7 +6362,8 @@ pick_next_task(struct rq *rq, struct rq_
 			update_rq_clock(rq_i);
 
 		p = pick_task(rq_i, rf);
-		if (unlikely(p == RETRY_TASK)) {
+		if (unlikely(seq != rq->core->core_task_seq ||
+			     WARN_ON_ONCE(p == RETRY_TASK))) {
 			/* rq lock may have been dropped, clocks invalidated */
 			core_clock_updated = false;
 			if (!(rq->clock_update_flags & RQCF_UPDATED))
@@ -6392,7 +6393,7 @@ pick_next_task(struct rq *rq, struct rq_
 			if (cookie)
 				p = sched_core_find(rq_i, cookie);
 			if (!p)
-				p = idle_sched_class.pick_task(rq_i, rf);
+				p = idle_sched_class.pick_task(rq_i, NULL);
 		}
 
 		rq_i->core_pick = p;

[PATCH 2/7] sched/core: Simplify/fix time updates

[Peter Zijlstra]

Directly test RQCF_UPDATED instead of using convoluted logic to try and
divinate the same.

Notably, the multi-pick loop's pick_task() can, when it ends up balancing, lock
and unlock the calling CPUs RQ and 'lose' the RQCF_UPDATED tag, which then
trips set_next_task().

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c |   41 +++++++++++++++++++++--------------------
 1 file changed, 21 insertions(+), 20 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6208,11 +6208,16 @@ extern void task_vruntime_update(struct
 
 static void queue_core_balance(struct rq *rq);
 
+static void opt_update_rq_clock(struct rq *rq)
+{
+	if (!(rq->clock_update_flags & RQCF_UPDATED))
+		update_rq_clock(rq);
+}
+
 static struct task_struct *
 pick_next_task(struct rq *rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(rq))
 {
-	bool core_clock_updated = (rq == rq->core);
 	struct task_struct *next, *p, *max;
 	const struct cpumask *smt_mask;
 	int i, cpu, seq, occ = 0;
@@ -6269,10 +6274,7 @@ pick_next_task(struct rq *rq, struct rq_
 	/* reset state */
 	rq->core->core_cookie = 0UL;
 	if (rq->core->core_forceidle_count) {
-		if (!core_clock_updated) {
-			update_rq_clock(rq->core);
-			core_clock_updated = true;
-		}
+		opt_update_rq_clock(rq->core);
 		sched_core_account_forceidle(rq);
 		/* reset after accounting force idle */
 		rq->core->core_forceidle_start = 0;
@@ -6299,14 +6301,11 @@ pick_next_task(struct rq *rq, struct rq_
 	 * and there are no cookied tasks running on siblings.
 	 */
 	if (!need_sync) {
+		opt_update_rq_clock(rq);
+
 		next = pick_task(rq, rf);
-		if (unlikely(next == RETRY_TASK)) {
-			/* rq lock may have been dropped, clocks invalidated */
-			core_clock_updated = false;
-			if (!(rq->clock_update_flags & RQCF_UPDATED))
-				update_rq_clock(rq);
+		if (unlikely(next == RETRY_TASK))
 			goto restart;
-		}
 
 		if (!next->core_cookie) {
 			rq->core_pick = NULL;
@@ -6329,6 +6328,7 @@ pick_next_task(struct rq *rq, struct rq_
 	 */
 	max = NULL;
 	for_each_cpu_wrap(i, smt_mask, cpu) {
+		struct rq_flags rf_i = *rf;
 		rq_i = cpu_rq(i);
 
 		/*
@@ -6336,18 +6336,12 @@ pick_next_task(struct rq *rq, struct rq_
 		 * pick_next_task(). If the current cpu is not the core,
 		 * the core may also have been updated above.
 		 */
-		if (i != cpu && (rq_i != rq->core || !core_clock_updated))
-			update_rq_clock(rq_i);
+		opt_update_rq_clock(rq_i);
 
-		p = pick_task(rq_i, rf);
+		p = pick_task(rq_i, &rf_i);
 		if (unlikely(seq != rq->core->core_task_seq ||
-			     WARN_ON_ONCE(p == RETRY_TASK))) {
-			/* rq lock may have been dropped, clocks invalidated */
-			core_clock_updated = false;
-			if (!(rq->clock_update_flags & RQCF_UPDATED))
-				update_rq_clock(rq);
+			     WARN_ON_ONCE(p == RETRY_TASK)))
 			goto restart;
-		}
 
 		rq_i->core_pick = p;
 		rq_i->core_dl_server = rq_i->dl_server;
@@ -6356,6 +6350,13 @@ pick_next_task(struct rq *rq, struct rq_
 			max = p;
 	}
 
+	/*
+	 * The above loop does @cpu first, if any sibling (which comes later)
+	 * does a LOCK+UNLOCK of @rq in order to (try) steal a task, our
+	 * RQCF_UPDATED got lost.
+	 */
+	rq->clock_update_flags |= RQCF_UPDATED;
+
 	cookie = rq->core->core_cookie = max->core_cookie;
 
 	/*

[PATCH 3/7] sched/core: Allow newidle for core-sched

[Peter Zijlstra]

As reported [1], it is possible for newidle, as called from pick_task_fair(),
to migrate a task that was already selected for the SMT sibling.

That is, in case of core scheduling (pick_next_task()'s restart_multi label),
pick_task() is called for each SMT sibling. Suppose SMT0 picks task A, and SMT1
has no tasks and does newidle, it must not be possible (but currently is) to
migrate A to SMT1 and also select A.

Avoids the 'stealing', but doesn't yet enable newidle just yet.

[1] https://patch.msgid.link/20260603095108.GA1684319@bytedance.com

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c  |    3 ++-
 kernel/sched/fair.c  |    6 ++----
 kernel/sched/sched.h |   15 ++++++++++++++-
 3 files changed, 18 insertions(+), 6 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -3049,7 +3049,8 @@ static int affine_move_task(struct rq *r
 		return -EINVAL;
 	}
 
-	if (task_on_cpu(rq, p) || READ_ONCE(p->__state) == TASK_WAKING) {
+	if (task_on_cpu(rq, p) || task_on_core(rq, p) ||
+	    READ_ONCE(p->__state) == TASK_WAKING) {
 		/*
 		 * MIGRATE_ENABLE gets here because 'p == current', but for
 		 * anything else we cannot do is_migration_disabled(), punt
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -10906,7 +10903,8 @@ int can_migrate_task(struct task_struct
 	env->flags &= ~LBF_ALL_PINNED;
 
 	if (task_on_cpu(env->src_rq, p) ||
-	    task_current_donor(env->src_rq, p)) {
+	    task_current_donor(env->src_rq, p) ||
+	    task_on_core(env->src_rq, p)) {
 		schedstat_inc(p->stats.nr_failed_migrations_running);
 		return 0;
 	}
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1573,6 +1573,14 @@ static inline bool task_has_sched_core(s
 	return !!p->core_cookie;
 }
 
+static inline bool task_on_core(struct rq *rq, struct task_struct *p)
+{
+	if (sched_core_disabled())
+		return false;
+
+	return rq->core_pick == p;
+}
+
 #else /* !CONFIG_SCHED_CORE: */
 
 static inline bool sched_core_enabled(struct rq *rq)
@@ -1618,6 +1626,11 @@ static inline bool task_has_sched_core(s
 	return false;
 }
 
+static inline bool task_on_core(struct rq *rq, struct task_struct *p)
+{
+	return false;
+}
+
 #endif /* !CONFIG_SCHED_CORE */
 
 #ifdef CONFIG_RT_GROUP_SCHED
@@ -4148,7 +4161,7 @@ void move_queued_task_locked(struct rq *
 static inline
 bool task_is_pushable(struct rq *rq, struct task_struct *p, int cpu)
 {
-	if (!task_on_cpu(rq, p) &&
+	if (!task_on_cpu(rq, p) && !task_on_core(rq, p) &&
 	    cpumask_test_cpu(cpu, &p->cpus_mask))
 		return true;
 

[PATCH 4/7] sched/rt: Add early exit on balance path

[Peter Zijlstra]

Add/make consistent an early exit on the balance_{rt,dl}() path. When there are
no pushable tasks, there is no point in taking the rq lock and all that jazz.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/deadline.c |   10 +++-------
 kernel/sched/rt.c       |    3 +++
 2 files changed, 6 insertions(+), 7 deletions(-)

--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -3250,17 +3250,13 @@ static void pull_dl_task(struct rq *this
 				   src_rq->dl.earliest_dl.next))
 			continue;
 
+		if (!has_pushable_dl_tasks(src_rq))
+			continue;
+
 		/* Might drop this_rq->lock */
 		push_task = NULL;
 		double_lock_balance(this_rq, src_rq);
 
-		/*
-		 * If there are no more pullable tasks on the
-		 * rq, we're done with it.
-		 */
-		if (src_rq->dl.dl_nr_running <= 1)
-			goto skip;
-
 		p = pick_earliest_pushable_dl_task(src_rq, this_cpu);
 
 		/*
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -2303,6 +2303,9 @@ static void pull_rt_task(struct rq *this
 		    this_rq->rt.highest_prio.curr)
 			continue;
 
+		if (!has_pushable_tasks(src_rq))
+			continue;
+
 		/*
 		 * We can potentially drop this_rq's lock in
 		 * double_lock_balance, and another CPU could

[PATCH 5/7] sched/fair: Reflow pick_task_fair() / newidle

[Peter Zijlstra]

This is all histerical artifacts. Clean up the code to be less convoluted.

Note: this also re-enables balance for core-sched.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/fair.c |   51 +++++++++++++--------------------------------------
 1 file changed, 13 insertions(+), 38 deletions(-)

--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -5942,7 +5942,7 @@ static inline unsigned long cfs_rq_load_
 	return cfs_rq->avg.load_avg;
 }
 
-static int sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
+static void sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(this_rq));
 
 static inline unsigned long task_util(struct task_struct *p)
@@ -10049,11 +10049,17 @@ struct task_struct *pick_task_fair(struc
 	struct cfs_rq *cfs_rq = &rq->cfs;
 	struct sched_entity *se;
 	struct task_struct *p;
-	int new_tasks;
 
 again:
-	if (!cfs_rq->h_nr_queued)
-		goto idle;
+	if (unlikely(!cfs_rq->h_nr_queued)) {
+		rq_modified_begin(rq, &fair_sched_class);
+		sched_balance_newidle(rq, rf);
+		if (rq_modified_above(rq, &fair_sched_class))
+			return RETRY_TASK;
+
+		if (!cfs_rq->h_nr_queued)
+			return NULL;
+	}
 
 	/* Might not have done put_prev_entity() */
 	if (cfs_rq->curr && cfs_rq->curr->on_rq)
@@ -10065,17 +10071,6 @@ struct task_struct *pick_task_fair(struc
 
 	p = task_of(se);
 	return p;
-
-idle:
-	if (sched_core_enabled(rq))
-		return NULL;
-
-	new_tasks = sched_balance_newidle(rq, rf);
-	if (new_tasks < 0)
-		return RETRY_TASK;
-	if (new_tasks > 0)
-		goto again;
-	return NULL;
 }
 
 static struct task_struct *
@@ -14531,13 +14526,8 @@ static inline void nohz_newidle_balance(
 /*
  * sched_balance_newidle is called by schedule() if this_cpu is about to become
  * idle. Attempts to pull tasks from other CPUs.
- *
- * Returns:
- *   < 0 - we released the lock and there are !fair tasks present
- *     0 - failed, no new tasks
- *   > 0 - success, new (fair) tasks present
  */
-static int sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
+static void sched_balance_newidle(struct rq *this_rq, struct rq_flags *rf)
 	__must_hold(__rq_lockp(this_rq))
 {
 	unsigned long next_balance = jiffies + HZ;
@@ -14554,7 +14544,7 @@ static int sched_balance_newidle(struct
 	 * Return 0; the task will be enqueued when switching to idle.
 	 */
 	if (this_rq->ttwu_pending)
-		return 0;
+		return;
 
 	/*
 	 * We must set idle_stamp _before_ calling sched_balance_rq()
@@ -14567,7 +14557,7 @@ static int sched_balance_newidle(struct
 	 * Do not pull tasks towards !active CPUs...
 	 */
 	if (!cpu_active(this_cpu))
-		return 0;
+		return;
 
 	/*
 	 * This is OK, because current is on_cpu, which avoids it being picked
@@ -14596,7 +14586,6 @@ static int sched_balance_newidle(struct
 	t0 = sched_clock_cpu(this_cpu);
 	__sched_balance_update_blocked_averages(this_rq);
 
-	rq_modified_begin(this_rq, &fair_sched_class);
 	raw_spin_rq_unlock(this_rq);
 
 	for_each_domain(this_cpu, sd) {
@@ -14654,18 +14643,6 @@ static int sched_balance_newidle(struct
 	if (curr_cost > this_rq->max_idle_balance_cost)
 		this_rq->max_idle_balance_cost = curr_cost;
 
-	/*
-	 * While browsing the domains, we released the rq lock, a task could
-	 * have been enqueued in the meantime. Since we're not going idle,
-	 * pretend we pulled a task.
-	 */
-	if (this_rq->cfs.h_nr_queued && !pulled_task)
-		pulled_task = 1;
-
-	/* If a higher prio class was modified, restart the pick */
-	if (rq_modified_above(this_rq, &fair_sched_class))
-		pulled_task = -1;
-
 out:
 	/* Move the next balance forward */
 	if (time_after(this_rq->next_balance, next_balance))
@@ -14677,8 +14654,6 @@ static int sched_balance_newidle(struct
 		nohz_newidle_balance(this_rq);
 
 	rq_repin_lock(this_rq, rf);
-
-	return pulled_task;
 }
 
 /*

[PATCH 6/7] sched/fair: Push sched_balance_newidle() unlock down

[Peter Zijlstra]

By pushing down the this_rq unlock, there are more cases where the lock isn't
dropped.

XXX needs numbers

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/fair.c |   19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -13417,6 +13417,7 @@ static int sched_balance_rq(int this_cpu
 			struct sched_domain *sd, enum cpu_idle_type idle,
 			int *continue_balancing)
 {
+	bool lock_rq = false, unlock_rq = idle == CPU_NEWLY_IDLE;
 	int ld_moved, cur_ld_moved, active_balance = 0;
 	struct sched_domain *sd_parent = sd->parent;
 	struct sched_group *group;
@@ -13485,6 +13486,12 @@ static int sched_balance_rq(int this_cpu
 		 */
 		env.loop_max  = min(sysctl_sched_nr_migrate, busiest->nr_running);
 
+		if (unlock_rq) {
+			raw_spin_rq_unlock(this_rq);
+			unlock_rq = false;
+			lock_rq = true;
+		}
+
 more_balance:
 		rq_lock_irqsave(busiest, &rf);
 		update_rq_clock(busiest);
@@ -13613,6 +13620,12 @@ static int sched_balance_rq(int this_cpu
 	if (!need_active_balance(&env))
 		goto out_unbalanced;
 
+	if (unlock_rq) {
+		raw_spin_rq_unlock(this_rq);
+		unlock_rq = false;
+		lock_rq = true;
+	}
+
 	scoped_guard (raw_spin_rq_lock_irqsave, busiest) {
 		/*
 		 * Don't kick the active_load_balance_cpu_stop,
@@ -13706,6 +13719,8 @@ static int sched_balance_rq(int this_cpu
 out:
 	if (need_unlock)
 		atomic_set_release(&sched_balance_running, 0);
+	if (lock_rq)
+		raw_spin_rq_lock(this_rq);
 
 	return ld_moved;
 }
@@ -14582,8 +14597,6 @@ static void sched_balance_newidle(struct
 	t0 = sched_clock_cpu(this_cpu);
 	__sched_balance_update_blocked_averages(this_rq);
 
-	raw_spin_rq_unlock(this_rq);
-
 	for_each_domain(this_cpu, sd) {
 		u64 domain_cost;
 
@@ -14634,8 +14647,6 @@ static void sched_balance_newidle(struct
 			break;
 	}
 
-	raw_spin_rq_lock(this_rq);
-
 	if (curr_cost > this_rq->max_idle_balance_cost)
 		this_rq->max_idle_balance_cost = curr_cost;
 

[PATCH 7/7] sched: Remove sched_class::balance()

[Peter Zijlstra]

Ever since commit 50653216e4ff ("sched: Add support to pick functions to take
rf"), we have the unfortunate situation that both sched_class::balance() and
sched_class::pick_task() have overlapping functionality in that they drop
rq->lock and balance tasks.

Additionally, prev_balance() is only called for a single RQ in the core-sched
case, resulting in 'missed' balance opportunities in this case.

Just like pick_task(), prev_balance() would iterate the classes and stop when
there is a runnable task found. The only difference is that pick_task() starts
from the top class, while prev_balance() would start from prev->class.

This means that if prev was an RT task, we'd get an extra visit to balance_dl()
-- and similary, if prev is fair, we'd visit balance_{dl,rt}() both. However,
this is not a problem, because both need_pull_{dl,rt}_task() will DTRT.

Therefore, move balance_{rt,dl}() into pick_task_{rt,dl}() and remove
sched_class::balance().

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 kernel/sched/core.c      |   23 -----------------------
 kernel/sched/deadline.c  |   10 ++++++----
 kernel/sched/fair.c      |   47 +++++++++++++----------------------------------
 kernel/sched/idle.c      |    7 -------
 kernel/sched/rt.c        |   14 ++++++--------
 kernel/sched/sched.h     |    5 -----
 kernel/sched/stop_task.c |    7 -------
 7 files changed, 25 insertions(+), 88 deletions(-)

--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -6115,25 +6115,6 @@ static inline void schedule_debug(struct
 	schedstat_inc(this_rq()->sched_count);
 }
 
-static void prev_balance(struct rq *rq, struct rq_flags *rf)
-{
-	const struct sched_class *start_class = rq->donor->sched_class;
-	const struct sched_class *class;
-
-	/*
-	 * We must do the balancing pass before put_prev_task(), such
-	 * that when we release the rq->lock the task is in the same
-	 * state as before we took rq->lock.
-	 *
-	 * We can terminate the balance pass as soon as we know there is
-	 * a runnable task of @class priority or higher.
-	 */
-	for_active_class_range(class, start_class, &idle_sched_class) {
-		if (class->balance && class->balance(rq, rf))
-			break;
-	}
-}
-
 /*
  * Pick up the highest-prio task:
  */
@@ -6171,8 +6152,6 @@ __pick_next_task(struct rq *rq, struct r
 	}
 
 restart:
-	prev_balance(rq, rf);
-
 	for_each_active_class(class) {
 		p = class->pick_task(rq, rf);
 		if (unlikely(p == RETRY_TASK))
@@ -6282,8 +6261,6 @@ pick_next_task(struct rq *rq, struct rq_
 		goto out_set_next;
 	}
 
-	prev_balance(rq, rf);
-
 	smt_mask = cpu_smt_mask(cpu);
 
 restart:
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -2705,7 +2705,7 @@ static void check_preempt_equal_dl(struc
 	resched_curr(rq);
 }
 
-static int balance_dl(struct rq *rq, struct rq_flags *rf)
+static void balance_dl(struct rq *rq, struct rq_flags *rf)
 {
 	/*
 	 * Note, rq->donor may change during rq lock drops,
@@ -2724,8 +2724,6 @@ static int balance_dl(struct rq *rq, str
 		pull_dl_task(rq);
 		rq_repin_lock(rq, rf);
 	}
-
-	return sched_stop_runnable(rq) || sched_dl_runnable(rq);
 }
 
 /*
@@ -2820,6 +2818,11 @@ static struct task_struct *__pick_task_d
 	struct dl_rq *dl_rq = &rq->dl;
 	struct task_struct *p;
 
+	rq_modified_begin(rq, &dl_sched_class);
+	balance_dl(rq, rf);
+	if (rq_modified_above(rq, &dl_sched_class))
+		return RETRY_TASK;
+
 again:
 	if (!sched_dl_runnable(rq))
 		return NULL;
@@ -3655,7 +3658,6 @@ DEFINE_SCHED_CLASS(dl) = {
 	.put_prev_task		= put_prev_task_dl,
 	.set_next_task		= set_next_task_dl,
 
-	.balance		= balance_dl,
 	.select_task_rq		= select_task_rq_dl,
 	.migrate_task_rq	= migrate_task_rq_dl,
 	.set_cpus_allowed       = set_cpus_allowed_dl,
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -464,12 +464,6 @@ select_task_rq_idle(struct task_struct *
 	return task_cpu(p); /* IDLE tasks as never migrated */
 }
 
-static int
-balance_idle(struct rq *rq, struct rq_flags *rf)
-{
-	return WARN_ON_ONCE(1);
-}
-
 /*
  * Idle tasks are unconditionally rescheduled:
  */
@@ -581,7 +575,6 @@ DEFINE_SCHED_CLASS(idle) = {
 	.put_prev_task		= put_prev_task_idle,
 	.set_next_task          = set_next_task_idle,
 
-	.balance		= balance_idle,
 	.select_task_rq		= select_task_rq_idle,
 	.set_cpus_allowed	= set_cpus_allowed_common,
 
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -1596,7 +1596,7 @@ static void check_preempt_equal_prio(str
 	resched_curr(rq);
 }
 
-static int balance_rt(struct rq *rq, struct rq_flags *rf)
+static void balance_rt(struct rq *rq, struct rq_flags *rf)
 {
 	/*
 	 * Note, rq->donor may change during rq lock drops,
@@ -1615,8 +1615,6 @@ static int balance_rt(struct rq *rq, str
 		pull_rt_task(rq);
 		rq_repin_lock(rq, rf);
 	}
-
-	return sched_stop_runnable(rq) || sched_dl_runnable(rq) || sched_rt_runnable(rq);
 }
 
 /*
@@ -1715,14 +1713,15 @@ static struct task_struct *_pick_next_ta
 
 static struct task_struct *pick_task_rt(struct rq *rq, struct rq_flags *rf)
 {
-	struct task_struct *p;
+	rq_modified_begin(rq, &rt_sched_class);
+	balance_rt(rq, rf);
+	if (rq_modified_above(rq, &rt_sched_class))
+		return RETRY_TASK;
 
 	if (!sched_rt_runnable(rq))
 		return NULL;
 
-	p = _pick_next_task_rt(rq);
-
-	return p;
+	return _pick_next_task_rt(rq);
 }
 
 static void put_prev_task_rt(struct rq *rq, struct task_struct *p, struct task_struct *next)
@@ -2610,7 +2609,6 @@ DEFINE_SCHED_CLASS(rt) = {
 	.put_prev_task		= put_prev_task_rt,
 	.set_next_task          = set_next_task_rt,
 
-	.balance		= balance_rt,
 	.select_task_rq		= select_task_rq_rt,
 	.set_cpus_allowed       = set_cpus_allowed_common,
 	.rq_online              = rq_online_rt,
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -2658,11 +2658,6 @@ struct sched_class {
 	void (*wakeup_preempt)(struct rq *rq, struct task_struct *p, int flags);
 
 	/*
-	 * schedule/pick_next_task/prev_balance: rq->lock
-	 */
-	int (*balance)(struct rq *rq, struct rq_flags *rf);
-
-	/*
 	 * schedule/pick_next_task: rq->lock
 	 */
 	struct task_struct *(*pick_task)(struct rq *rq, struct rq_flags *rf);
--- a/kernel/sched/stop_task.c
+++ b/kernel/sched/stop_task.c
@@ -15,12 +15,6 @@ select_task_rq_stop(struct task_struct *
 	return task_cpu(p); /* stop tasks as never migrate */
 }
 
-static int
-balance_stop(struct rq *rq, struct rq_flags *rf)
-{
-	return sched_stop_runnable(rq);
-}
-
 static void
 wakeup_preempt_stop(struct rq *rq, struct task_struct *p, int flags)
 {
@@ -107,7 +101,6 @@ DEFINE_SCHED_CLASS(stop) = {
 	.put_prev_task		= put_prev_task_stop,
 	.set_next_task          = set_next_task_stop,
 
-	.balance		= balance_stop,
 	.select_task_rq		= select_task_rq_stop,
 	.set_cpus_allowed	= set_cpus_allowed_common,