ANDROID: sched: Move scratch_mask to a percpu variable

The recently introduced scratch_mask in struct rq breaks the KMI. Turn
it into per-cpu table instead, similar to how load_balance_mask or
select_idle_mask are allocated in the scheduler.

Bug: 254942195
Bug: 254447891
Bug: 254812379

Signed-off-by: Quentin Perret <qperret@google.com>
Change-Id: I331336dd13eaaf83be0a4029f89becda45c863f8
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 19761724..4a16ab3 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -57,6 +57,9 @@
 DEFINE_PER_CPU_SHARED_ALIGNED(struct rq, runqueues);
 EXPORT_SYMBOL_GPL(runqueues);
 
+/* Scratch cpumask to be temporarily used under rq_lock */
+DEFINE_PER_CPU(cpumask_var_t, scratch_mask);
+
 #ifdef CONFIG_SCHED_DEBUG
 /*
  * Debugging: various feature bits
@@ -2947,18 +2950,20 @@
 static int __set_cpus_allowed_ptr(struct task_struct *p,
 				  struct affinity_context *ctx)
 {
+	struct cpumask *cpus;
 	struct rq_flags rf;
 	struct rq *rq;
 
 	rq = task_rq_lock(p, &rf);
+	cpus = per_cpu(scratch_mask, rq->cpu);
 	/*
 	 * Masking should be skipped if SCA_USER or any of the SCA_MIGRATE_*
 	 * flags are set.
 	 */
 	if (p->user_cpus_ptr &&
 	    !(ctx->flags & (SCA_USER | SCA_MIGRATE_ENABLE | SCA_MIGRATE_DISABLE)) &&
-	    cpumask_and(rq->scratch_mask, ctx->new_mask, p->user_cpus_ptr))
-		ctx->new_mask = rq->scratch_mask;
+	    cpumask_and(cpus, ctx->new_mask, p->user_cpus_ptr))
+		ctx->new_mask = cpus;
 
 	return __set_cpus_allowed_ptr_locked(p, ctx, rq, &rf);
 }
@@ -9622,6 +9627,8 @@
 			cpumask_size(), GFP_KERNEL, cpu_to_node(i));
 		per_cpu(select_idle_mask, i) = (cpumask_var_t)kzalloc_node(
 			cpumask_size(), GFP_KERNEL, cpu_to_node(i));
+		per_cpu(scratch_mask, i) = (cpumask_var_t)kzalloc_node(
+			cpumask_size(), GFP_KERNEL, cpu_to_node(i));
 	}
 #endif /* CONFIG_CPUMASK_OFFSTACK */
 
@@ -9727,7 +9734,6 @@
 
 		rq->core_cookie = 0UL;
 #endif
-		zalloc_cpumask_var_node(&rq->scratch_mask, GFP_KERNEL, cpu_to_node(i));
 	}
 
 	set_load_weight(&init_task, false);
diff --git a/kernel/sched/sched.h b/kernel/sched/sched.h
index 615bde0..7c30836 100644
--- a/kernel/sched/sched.h
+++ b/kernel/sched/sched.h
@@ -1131,9 +1131,6 @@
 	unsigned int		core_forceidle_seq;
 #endif
 
-	/* Scratch cpumask to be temporarily used under rq_lock */
-	cpumask_var_t		scratch_mask;
-
 	ANDROID_VENDOR_DATA_ARRAY(1, 96);
 	ANDROID_OEM_DATA_ARRAY(1, 16);