[RFC PATCH 3/4] mm: add BPF hook for THP adjustment

[Yafang Shao <laoar.shao@gmail.com>]

We will use the @vma parameter in BPF programs to determine whether THP can
be used. The typical workflow is as follows:

1. Retrieve the mm_struct from the given @vma.
2. Obtain the task_struct associated with the mm_struct
   It depends on CONFIG_MEMCG.
3. Adjust THP behavior dynamically based on task attributes
   E.g., based on the task’s cgroup

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 mm/Makefile   |  3 +++
 mm/bpf.c      | 36 ++++++++++++++++++++++++++++++++++++
 mm/bpf.h      | 21 +++++++++++++++++++++
 mm/internal.h |  3 +++
 4 files changed, 63 insertions(+)
 create mode 100644 mm/bpf.c
 create mode 100644 mm/bpf.h

diff --git a/mm/Makefile b/mm/Makefile
index e7f6bbf8ae5f..97055da04746 100644
--- a/mm/Makefile
+++ b/mm/Makefile
@@ -99,6 +99,9 @@ obj-$(CONFIG_MIGRATION) += migrate.o
 obj-$(CONFIG_NUMA) += memory-tiers.o
 obj-$(CONFIG_DEVICE_MIGRATION) += migrate_device.o
 obj-$(CONFIG_TRANSPARENT_HUGEPAGE) += huge_memory.o khugepaged.o
+ifdef CONFIG_BPF_SYSCALL
+obj-$(CONFIG_TRANSPARENT_HUGEPAGE) += bpf.o
+endif
 obj-$(CONFIG_PAGE_COUNTER) += page_counter.o
 obj-$(CONFIG_MEMCG_V1) += memcontrol-v1.o
 obj-$(CONFIG_MEMCG) += memcontrol.o vmpressure.o
diff --git a/mm/bpf.c b/mm/bpf.c
new file mode 100644
index 000000000000..72eebcdbad56
--- /dev/null
+++ b/mm/bpf.c
@@ -0,0 +1,36 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Author: Yafang Shao <laoar.shao@gmail.com>
+ */
+
+#include <linux/bpf.h>
+#include <linux/mm_types.h>
+
+__bpf_hook_start();
+
+/* Checks if this @vma can use THP. */
+__weak noinline int
+mm_bpf_thp_vma_allowable(struct vm_area_struct *vma)
+{
+	/* At present, fmod_ret exclusively uses 0 to signify that the return
+	 * value remains unchanged.
+	 */
+	return 0;
+}
+
+__bpf_hook_end();
+
+BTF_SET8_START(mm_bpf_fmod_ret_ids)
+BTF_ID_FLAGS(func, mm_bpf_thp_vma_allowable)
+BTF_SET8_END(mm_bpf_fmod_ret_ids)
+
+static const struct btf_kfunc_id_set mm_bpf_fmodret_set = {
+	.owner = THIS_MODULE,
+	.set   = &mm_bpf_fmod_ret_ids,
+};
+
+static int __init bpf_mm_kfunc_init(void)
+{
+	return register_btf_fmodret_id_set(&mm_bpf_fmodret_set);
+}
+late_initcall(bpf_mm_kfunc_init);
diff --git a/mm/bpf.h b/mm/bpf.h
new file mode 100644
index 000000000000..e03a38084b08
--- /dev/null
+++ b/mm/bpf.h
@@ -0,0 +1,21 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later */
+#ifndef __MM_BPF_H
+#define __MM_BPF_H
+
+#define MM_BPF_ALLOWABLE	(1)
+#define MM_BPF_NOT_ALLOWABLE	(-1)
+
+#define MM_BPF_ALLOWABLE_HOOK(func, args...)	{	\
+	int ret = func(args);				\
+							\
+	if (ret == MM_BPF_ALLOWABLE)			\
+		return 1;				\
+	if (ret == MM_BPF_NOT_ALLOWABLE)		\
+		return 0;				\
+}
+
+#ifdef CONFIG_TRANSPARENT_HUGEPAGE
+int mm_bpf_thp_vma_allowable(struct vm_area_struct *vma);
+#endif
+
+#endif
diff --git a/mm/internal.h b/mm/internal.h
index aa698a11dd68..c8bf405fa581 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -21,6 +21,7 @@
 
 /* Internal core VMA manipulation functions. */
 #include "vma.h"
+#include "bpf.h"
 
 struct folio_batch;
 
@@ -1632,6 +1633,7 @@ static inline bool reclaim_pt_is_enabled(unsigned long start, unsigned long end,
  */
 static inline bool hugepage_global_enabled(struct vm_area_struct *vma)
 {
+	MM_BPF_ALLOWABLE_HOOK(mm_bpf_thp_vma_allowable, vma);
 	return transparent_hugepage_flags &
 			((1<<TRANSPARENT_HUGEPAGE_FLAG) |
 			(1<<TRANSPARENT_HUGEPAGE_REQ_MADV_FLAG));
@@ -1639,6 +1641,7 @@ static inline bool hugepage_global_enabled(struct vm_area_struct *vma)
 
 static inline bool hugepage_global_always(struct vm_area_struct *vma)
 {
+	MM_BPF_ALLOWABLE_HOOK(mm_bpf_thp_vma_allowable, vma);
 	return transparent_hugepage_flags &
 			(1<<TRANSPARENT_HUGEPAGE_FLAG);
 }
-- 
2.43.5