ANDROID: sched: support pick_migrate_task in common kernel

__pick_migrate_task() was present in 5.10 and removed in
5.15 in favor of the balance_push concept.

Restore this function so that the ability to pick a task to migrate
across all scheduling classes can be used by a vendor module.

Function Removed In commit 1cf12e08bc
sched/hotplug: Consolidate task migration on CPU unplug

Bug: 205164003
Change-Id: I56ecd96c32b49495132daefbfc106568ab84e528
Signed-off-by: Stephen Dickey <quic_dickey@quicinc.com>

include/linux/sched.h

@@ -298,6 +298,7 @@ extern int __must_check io_schedule_prepare(void);
 extern void io_schedule_finish(int token);
 extern long io_schedule_timeout(long timeout);
 extern void io_schedule(void);
+extern struct task_struct *pick_migrate_task(struct rq *rq);
 
 /**
  * struct prev_cputime - snapshot of system and user cputime

kernel/sched/core.c

@@ -8870,6 +8870,24 @@ void idle_task_exit(void)
 	/* finish_cpu(), as ran on the BP, will clean up the active_mm state */
 }
 
+struct task_struct *pick_migrate_task(struct rq *rq)
+{
+	const struct sched_class *class;
+	struct task_struct *next;
+
+	for_each_class(class) {
+		next = class->pick_next_task(rq);
+		if (next) {
+			next->sched_class->put_prev_task(rq, next);
+			return next;
+		}
+	}
+
+	/* The idle class should always have a runnable task */
+	BUG();
+}
+EXPORT_SYMBOL_GPL(pick_migrate_task);
+
 static int __balance_push_cpu_stop(void *arg)
 {
 	struct task_struct *p = arg;