Expose the process automatic NUMA balancing mode through prctl().
PR_SET_NUMA_BALANCING accepts PR_NUMA_BALANCING_DISABLE to opt the whole
thread group out of future NUMA scanning and PR_NUMA_BALANCING_ENABLE to
allow it to participate again when global NUMA balancing and memory
policy permit it.

PR_GET_NUMA_BALANCING returns the configured process mode, not the
effective state after combining the global sysctl/static key and memory
policy restrictions. This keeps the ABI weak and predictable: a process
can observe and restore the value it configured, while administrators
retain kernel.numa_balancing as the top-level hard-off switch.

Signed-off-by: Li Zhe <lizhe.67@bytedance.com>
---
 include/uapi/linux/prctl.h |  6 ++++++
 kernel/sys.c               | 18 ++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/include/uapi/linux/prctl.h b/include/uapi/linux/prctl.h
index b6ec6f693719..5bc0a2a2a930 100644
--- a/include/uapi/linux/prctl.h
+++ b/include/uapi/linux/prctl.h
@@ -416,4 +416,10 @@ struct prctl_mm_map {
 # define PR_CFI_DISABLE		_BITUL(1)
 # define PR_CFI_LOCK		_BITUL(2)
 
+/* Per-process automatic NUMA balancing control */
+#define PR_SET_NUMA_BALANCING	82
+#define PR_GET_NUMA_BALANCING	83
+# define PR_NUMA_BALANCING_DISABLE	0
+# define PR_NUMA_BALANCING_ENABLE	1
+
 #endif /* _LINUX_PRCTL_H */
diff --git a/kernel/sys.c b/kernel/sys.c
index 35b538ba843c..fb01841c0ce9 100644
--- a/kernel/sys.c
+++ b/kernel/sys.c
@@ -60,6 +60,7 @@
 #include <linux/sched/stat.h>
 #include <linux/sched/mm.h>
 #include <linux/sched/coredump.h>
+#include <linux/sched/numa_balancing.h>
 #include <linux/sched/task.h>
 #include <linux/sched/cputime.h>
 #include <linux/rcupdate.h>
@@ -2907,6 +2908,23 @@ SYSCALL_DEFINE5(prctl, int, option, unsigned long, arg2, unsigned long, arg3,
 		if (arg3 & PR_CFI_LOCK && !(arg3 & PR_CFI_DISABLE))
 			error = arch_prctl_lock_branch_landing_pad_state(me);
 		break;
+	case PR_SET_NUMA_BALANCING:
+		if (arg3 || arg4 || arg5)
+			return -EINVAL;
+		if (arg2 != PR_NUMA_BALANCING_DISABLE &&
+		    arg2 != PR_NUMA_BALANCING_ENABLE)
+			return -EINVAL;
+		error = task_numa_balancing_set_current(arg2 ==
+							 PR_NUMA_BALANCING_ENABLE);
+		break;
+	case PR_GET_NUMA_BALANCING:
+		if (arg2 || arg3 || arg4 || arg5)
+			return -EINVAL;
+		error = task_numa_balancing_get_current();
+		if (error >= 0)
+			error = error ? PR_NUMA_BALANCING_ENABLE :
+				PR_NUMA_BALANCING_DISABLE;
+		break;
 	default:
 		trace_task_prctl_unknown(option, arg2, arg3, arg4, arg5);
 		error = -EINVAL;
-- 
2.20.1