[PATCH 1/4] sched/numa: Track per-process automatic NUMA balancing mode

@@ -1382,6 +1382,12 @@ struct task_struct {
 	short				pref_node_fork;
 #endif
 #ifdef CONFIG_NUMA_BALANCING
+	/*
+	 * Scheduler snapshot of signal_struct::numa_balancing_enabled. It is
+	 * updated with sched_change for runqueue NUMA accounting; the
+	 * user-visible process mode lives in signal_struct.
+	 */
+	bool				numa_balancing_sched_enabled;
 	int				numa_scan_seq;
 	unsigned int			numa_scan_period;
 	unsigned int			numa_scan_period_max;

@@ -8,6 +8,7 @@
  */
 
 #include <linux/sched.h>
+#include <linux/sched/signal.h>
 
 #define TNF_MIGRATED	0x01
 #define TNF_NO_GROUP	0x02

@@ -30,6 +31,22 @@ extern void task_numa_fault(int last_node, int node, int pages, int flags);
 extern pid_t task_numa_group_id(struct task_struct *p);
 extern void set_numabalancing_state(bool enabled);
 extern void task_numa_free(struct task_struct *p, bool final);
+static inline bool task_numa_sched_snapshot_enabled(struct task_struct *p)
+{
+	return READ_ONCE(p->numa_balancing_sched_enabled);
+}
+
+static inline bool task_numa_process_mode_enabled(struct task_struct *p)
+{
+	return READ_ONCE(p->signal->numa_balancing_enabled);
+}
+
+int task_numa_balancing_set_current(bool enabled);
+
+static inline int task_numa_balancing_get_current(void)
+{
+	return task_numa_process_mode_enabled(current);
+}
 bool should_numa_migrate_memory(struct task_struct *p, struct folio *folio,
 				int src_nid, int dst_cpu);
 #else

@@ -47,6 +64,26 @@ static inline void set_numabalancing_state(bool enabled)
 static inline void task_numa_free(struct task_struct *p, bool final)
 {
 }
+
+static inline bool task_numa_sched_snapshot_enabled(struct task_struct *p)
+{
+	return false;
+}
+
+static inline bool task_numa_process_mode_enabled(struct task_struct *p)
+{
+	return false;
+}
+
+static inline int task_numa_balancing_set_current(bool enabled)
+{
+	return -EINVAL;
+}
+
+static inline int task_numa_balancing_get_current(void)
+{
+	return -EINVAL;
+}
 static inline bool should_numa_migrate_memory(struct task_struct *p,
 				struct folio *folio, int src_nid, int dst_cpu)
 {

@@ -179,6 +179,14 @@ struct signal_struct {
 
 #ifdef CONFIG_SCHED_AUTOGROUP
 	struct autogroup *autogroup;
+#endif
+#ifdef CONFIG_NUMA_BALANCING
+	/*
+	 * Thread-group automatic NUMA balancing mode configured through prctl().
+	 * Scheduler hot paths use task_struct::numa_balancing_sched_enabled as
+	 * their per-task runqueue accounting snapshot.
+	 */
+	bool numa_balancing_enabled;
 #endif
 	/*
 	 * Cumulative resource counters for dead threads in the group,

@@ -40,6 +40,9 @@ static struct signal_struct init_signals = {
 	.cputimer		= {
 		.cputime_atomic	= INIT_CPUTIME_ATOMIC,
 	},
+#endif
+#ifdef CONFIG_NUMA_BALANCING
+	.numa_balancing_enabled = true,
 #endif
 	INIT_CPU_TIMERS(init_signals)
 	.pids = {

@@ -224,6 +227,7 @@ struct task_struct init_task __aligned(L1_CACHE_BYTES) = {
 	.vtime.state	= VTIME_SYS,
 #endif
 #ifdef CONFIG_NUMA_BALANCING
+	.numa_balancing_sched_enabled = true,
 	.numa_preferred_nid = NUMA_NO_NODE,
 	.numa_group	= NULL,
 	.numa_faults	= NULL,

@@ -2496,6 +2496,14 @@ __latent_entropy struct task_struct *copy_process(
 
 	/* No more failure paths after this point. */
 
+#ifdef CONFIG_NUMA_BALANCING
+	p->numa_balancing_sched_enabled =
+		READ_ONCE(current->signal->numa_balancing_enabled);
+	if (!(clone_flags & CLONE_THREAD))
+		p->signal->numa_balancing_enabled =
+			p->numa_balancing_sched_enabled;
+#endif
+
 	/*
 	 * Copy seccomp details explicitly here, in case they were changed
 	 * before holding sighand lock.

@@ -36,6 +36,7 @@
 #include <linux/sched/isolation.h>
 #include <linux/sched/loadavg.h>
 #include <linux/sched/mm.h>
+#include <linux/sched/numa_balancing.h>
 #include <linux/sched/nohz.h>
 #include <linux/sched/rseq_api.h>
 #include <linux/sched/rt.h>

@@ -603,6 +604,8 @@ int task_llc(const struct task_struct *p)
  *				p->se.load, p->rt_priority,
  *				p->dl.dl_{runtime, deadline, period, flags, bw, density}
  *  - sched_setnuma():		p->numa_preferred_nid
+ *  - task_numa_balancing_set_current():	p->signal->numa_balancing_enabled,
+ *					p->numa_balancing_sched_enabled
  *  - sched_move_task():	p->sched_task_group
  *  - uclamp_update_active()	p->uclamp*
  *

@@ -8407,6 +8410,42 @@ void sched_setnuma(struct task_struct *p, int nid)
 	scoped_guard (sched_change, p, DEQUEUE_SAVE)
 		p->numa_preferred_nid = nid;
 }
+
+static void sched_numa_balancing_change_task(struct task_struct *p, bool enabled)
+{
+	guard(task_rq_lock)(p);
+	scoped_guard (sched_change, p, DEQUEUE_SAVE)
+		WRITE_ONCE(p->numa_balancing_sched_enabled, enabled);
+}
+
+int task_numa_balancing_set_current(bool enabled)
+{
+	static DEFINE_MUTEX(task_numa_balancing_mutex);
+	struct task_struct *t;
+	unsigned long flags;
+	bool old_enabled;
+
+	guard(mutex)(&task_numa_balancing_mutex);
+
+	if (WARN_ON_ONCE(!lock_task_sighand(current, &flags)))
+		return -ESRCH;
+
+	old_enabled = current->signal->numa_balancing_enabled;
+	if (old_enabled != enabled)
+		WRITE_ONCE(current->signal->numa_balancing_enabled, enabled);
+
+	unlock_task_sighand(current, &flags);
+
+	if (old_enabled != enabled) {
+		sched_numa_balancing_change_task(current, enabled);
+		read_lock(&tasklist_lock);
+		for_other_threads(current, t)
+			sched_numa_balancing_change_task(t, enabled);
+		read_unlock(&tasklist_lock);
+	}
+
+	return 0;
+}
 #endif /* CONFIG_NUMA_BALANCING */
 
 #ifdef CONFIG_HOTPLUG_CPU

@@ -38,6 +38,7 @@
 #include <linux/sched/cond_resched.h>
 #include <linux/sched/cputime.h>
 #include <linux/sched/isolation.h>
+#include <linux/sched/numa_balancing.h>
 #include <linux/sched/nohz.h>
 #include <linux/sched/prio.h>
 #include <linux/static_call.h>

@@ -1712,6 +1713,7 @@ static int get_pref_llc(struct task_struct *p, struct mm_struct *mm)
 		 * conflict only exists for a short period of time.
 		 */
 		if (static_branch_likely(&sched_numa_balancing) &&
+		    task_numa_sched_snapshot_enabled(p) &&
 		    p->numa_preferred_nid >= 0 &&
 		    cpu_to_node(mm_sched_cpu) != p->numa_preferred_nid)
 			mm_sched_llc = -1;

@@ -1808,6 +1810,9 @@ static void get_scan_cpumasks(cpumask_var_t cpus, struct task_struct *p)
 	if (!static_branch_likely(&sched_numa_balancing))
 		goto out;
 
+	if (!task_numa_sched_snapshot_enabled(p))
+		goto out;
+
 	cpu = READ_ONCE(p->mm->sc_stat.cpu);
 	if (cpu != -1)
 		nid = cpu_to_node(cpu);

@@ -2387,14 +2392,22 @@ static unsigned int task_scan_max(struct task_struct *p)
 
 static void account_numa_enqueue(struct rq *rq, struct task_struct *p)
 {
-	rq->nr_numa_running += (p->numa_preferred_nid != NUMA_NO_NODE);
-	rq->nr_preferred_running += (p->numa_preferred_nid == task_node(p));
+	bool enabled = task_numa_sched_snapshot_enabled(p);
+
+	rq->nr_numa_running += enabled &&
+				 p->numa_preferred_nid != NUMA_NO_NODE;
+	rq->nr_preferred_running += enabled &&
+				     p->numa_preferred_nid == task_node(p);
 }
 
 static void account_numa_dequeue(struct rq *rq, struct task_struct *p)
 {
-	rq->nr_numa_running -= (p->numa_preferred_nid != NUMA_NO_NODE);
-	rq->nr_preferred_running -= (p->numa_preferred_nid == task_node(p));
+	bool enabled = task_numa_sched_snapshot_enabled(p);
+
+	rq->nr_numa_running -= enabled &&
+				 p->numa_preferred_nid != NUMA_NO_NODE;
+	rq->nr_preferred_running -= enabled &&
+				     p->numa_preferred_nid == task_node(p);
 }
 
 /* Shared or private faults. */

@@ -3085,6 +3098,9 @@ static bool task_numa_compare(struct task_numa_env *env,
 			goto unlock;
 	}
 
+	if (!task_numa_sched_snapshot_enabled(cur))
+		goto unlock;
+
 	/* Skip this swap candidate if cannot move to the source cpu. */
 	if (!cpumask_test_cpu(env->src_cpu, cur->cpus_ptr))
 		goto unlock;

@@ -3992,6 +4008,9 @@ void task_numa_fault(int last_cpupid, int mem_node, int pages, int flags)
 	if (!static_branch_likely(&sched_numa_balancing))
 		return;
 
+	if (!task_numa_sched_snapshot_enabled(p))
+		return;
+
 	/* for example, ksmd faulting in a user's mm */
 	if (!p->mm)
 		return;

@@ -4436,6 +4455,9 @@ static void task_tick_numa(struct rq *rq, struct task_struct *curr)
 	if (!curr->mm || (curr->flags & (PF_EXITING | PF_KTHREAD)) || work->next != work)
 		return;
 
+	if (!task_numa_sched_snapshot_enabled(curr))
+		return;
+
 	/*
 	 * Using runtime rather than walltime has the dual advantage that
 	 * we (mostly) drive the selection from busy threads and that the

@@ -4463,6 +4485,9 @@ static void update_scan_period(struct task_struct *p, int new_cpu)
 	if (!static_branch_likely(&sched_numa_balancing))
 		return;
 
+	if (!task_numa_sched_snapshot_enabled(p))
+		return;
+
 	if (!p->mm || !p->numa_faults || (p->flags & PF_EXITING))
 		return;
 

@@ -10475,6 +10500,9 @@ static long migrate_degrades_locality(struct task_struct *p, struct lb_env *env)
 	if (!static_branch_likely(&sched_numa_balancing))
 		return 0;
 
+	if (!task_numa_sched_snapshot_enabled(p))
+		return 0;
+
 	if (!p->numa_faults || !(env->sd->flags & SD_NUMA))
 		return 0;