diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index ef48b8600f0f..73a0e232abc3 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -2568,18 +2568,7 @@ __do_set_cpus_allowed(struct task_struct *p, struct affinity_context *ctx)
 	if (running)
 		put_prev_task(rq, p);
 
-	/*
-	 * XXX: ANDROID: we can't use sched_class::set_cpus_allowed() here
-	 * because it doesn't take a struct affinity_context as parameter for
-	 * GKI KMI stability reason -- see b/254812379. To avoid the problem,
-	 * let's hardcode the indirection here and hope for the best. The only
-	 * other potential users of p->set_cpus_allowed() will be in vendor
-	 * modules.
-	 */
-	if (likely(p->sched_class != &dl_sched_class))
-		set_cpus_allowed_common(p, ctx);
-	else
-		set_cpus_allowed_dl(p, ctx);
+	p->sched_class->set_cpus_allowed(p, ctx);
 
 	if (queued)
 		enqueue_task(rq, p, ENQUEUE_RESTORE | ENQUEUE_NOCLOCK);
diff --git a/kernel/sched/deadline.c b/kernel/sched/deadline.c
index 28772e64124b..6b3fc584a3e3 100644
--- a/kernel/sched/deadline.c
+++ b/kernel/sched/deadline.c
@@ -2341,7 +2341,8 @@ static void task_woken_dl(struct rq *rq, struct task_struct *p)
 	}
 }
 
-void set_cpus_allowed_dl(struct task_struct *p, struct affinity_context *ctx)
+static void set_cpus_allowed_dl(struct task_struct *p,
+				struct affinity_context *ctx)
 {
 	struct root_domain *src_rd;
 	struct rq *rq;
@@ -2373,18 +2374,6 @@ void set_cpus_allowed_dl(struct task_struct *p, struct affinity_context *ctx)
 	set_cpus_allowed_common(p, ctx);
 }
 
-static void set_cpus_allowed_dl_cb(struct task_struct *p,
-				   const struct cpumask *new_mask,
-				   u32 flags)
-{
-	struct affinity_context ac = {
-		.new_mask  = new_mask,
-		.flags     = flags,
-	};
-	WARN_ONCE(1, "Unexpected use of dl_sched_class::set_cpus_allowed()");
-	set_cpus_allowed_dl(p, &ac);
-}
-
 /* Assumes rq->lock is held */
 static void rq_online_dl(struct rq *rq)
 {
@@ -2577,7 +2566,7 @@ DEFINE_SCHED_CLASS(dl) = {
 	.pick_task		= pick_task_dl,
 	.select_task_rq		= select_task_rq_dl,
 	.migrate_task_rq	= migrate_task_rq_dl,
-	.set_cpus_allowed       = set_cpus_allowed_dl_cb,
+	.set_cpus_allowed       = set_cpus_allowed_dl,
 	.rq_online              = rq_online_dl,
 	.rq_offline             = rq_offline_dl,
 	.task_woken		= task_woken_dl,
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index ff04806fa877..87b1c7ecec3b 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -11842,7 +11842,7 @@ DEFINE_SCHED_CLASS(fair) = {
 	.rq_offline		= rq_offline_fair,
 
 	.task_dead		= task_dead_fair,
-	.set_cpus_allowed	= set_cpus_allowed_common_cb,
+	.set_cpus_allowed	= set_cpus_allowed_common,
 #endif
 
 	.task_tick		= task_tick_fair,
diff --git a/kernel/sched/idle.c b/kernel/sched/idle.c
index 317d58516948..499a3e286cd0 100644
--- a/kernel/sched/idle.c
+++ b/kernel/sched/idle.c
@@ -516,7 +516,7 @@ DEFINE_SCHED_CLASS(idle) = {
 	.balance		= balance_idle,
 	.pick_task		= pick_task_idle,
 	.select_task_rq		= select_task_rq_idle,
-	.set_cpus_allowed	= set_cpus_allowed_common_cb,
+	.set_cpus_allowed	= set_cpus_allowed_common,
 #endif
 
 	.task_tick		= task_tick_idle,
diff --git a/kernel/sched/rt.c b/kernel/sched/rt.c
index b2d3281416a0..a3002fc257fc 100644
--- a/kernel/sched/rt.c
+++ b/kernel/sched/rt.c
@@ -2635,7 +2635,7 @@ DEFINE_SCHED_CLASS(rt) = {
 	.balance		= balance_rt,
 	.pick_task		= pick_task_rt,
 	.select_task_rq		= select_task_rq_rt,
-	.set_cpus_allowed       = set_cpus_allowed_common_cb,
+	.set_cpus_allowed       = set_cpus_allowed_common,
 	.rq_online              = rq_online_rt,
 	.rq_offline             = rq_offline_rt,
 	.task_woken		= task_woken_rt,
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index ca78455a6027..90c78dd9758f 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -2205,9 +2205,7 @@ struct sched_class {
 
 	void (*task_woken)(struct rq *this_rq, struct task_struct *task);
 
-	void (*set_cpus_allowed)(struct task_struct *p,
-				 const struct cpumask *newmask,
-				 u32 flags);
+	void (*set_cpus_allowed)(struct task_struct *p, struct affinity_context *ctx);
 
 	void (*rq_online)(struct rq *rq);
 	void (*rq_offline)(struct rq *rq);
@@ -2320,20 +2318,8 @@ extern void update_group_capacity(struct sched_domain *sd, int cpu);
 
 extern void trigger_load_balance(struct rq *rq);
 
-extern void set_cpus_allowed_dl(struct task_struct *p, struct affinity_context *ctx);
 extern void set_cpus_allowed_common(struct task_struct *p, struct affinity_context *ctx);
 
-static inline void set_cpus_allowed_common_cb(struct task_struct *p, const struct cpumask *new_mask, u32 flags)
-{
-	struct affinity_context ac = {
-		.new_mask  = new_mask,
-		.flags     = flags,
-	};
-
-	WARN_ONCE(1, "Unexpected use of sched_class::set_cpus_allowed()");
-	set_cpus_allowed_common(p, &ac);
-}
-
 static inline struct task_struct *get_push_task(struct rq *rq)
 {
 	struct task_struct *p = rq->curr;
diff --git a/kernel/sched/stop_task.c b/kernel/sched/stop_task.c
index ad8fa4d37247..f988ebe3febb 100644
--- a/kernel/sched/stop_task.c
+++ b/kernel/sched/stop_task.c
@@ -134,7 +134,7 @@ DEFINE_SCHED_CLASS(stop) = {
 	.balance		= balance_stop,
 	.pick_task		= pick_task_stop,
 	.select_task_rq		= select_task_rq_stop,
-	.set_cpus_allowed	= set_cpus_allowed_common_cb,
+	.set_cpus_allowed	= set_cpus_allowed_common,
 #endif
 
 	.task_tick		= task_tick_stop,