[PATCH v2 1/4] mm: support promotion-only NUMA hinting scans

[Gregory Price <gourry@gourry.net>]

From: "Gregory Price (Meta)" <gourry@gourry.net>

folio_can_map_prot_numa() derives the eligible memory from the global
balancing mode. Combined mode needs the scanner to distinguish placement
scans from promotion-only scans.

Add MM_CP_PROT_NUMA_PROMO_ONLY and let change_prot_numa() callers request
a promotion-only protection walk. Carry the choice with each walk so the
PTE and PMD paths use the same decision.

Initially derive the value from the normal balancing mode, preserving the
existing behavior for the following fixes.

Fixes: c574bbe91703 ("NUMA balancing: optimize page placement for memory tiering system")
Cc: stable@vger.kernel.org
Assisted-by: LLM
Signed-off-by: Gregory Price (Meta) <gourry@gourry.net>
---
 include/linux/mm.h  |  4 +++-
 kernel/sched/fair.c |  7 ++++++-
 mm/huge_memory.c    |  3 ++-
 mm/internal.h       |  5 +++--
 mm/mempolicy.c      | 22 +++++++++++++---------
 mm/mprotect.c       |  4 +++-
 6 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/include/linux/mm.h b/include/linux/mm.h
index 969594074fd2..2d1b59a27629 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -3408,6 +3408,8 @@ int get_cmdline(struct task_struct *task, char *buffer, int buflen);
 #define  MM_CP_UFFD_RWP_RESOLVE            (1UL << 5) /* resolve rwp */
 #define  MM_CP_UFFD_RWP_ALL                (MM_CP_UFFD_RWP | \
 					    MM_CP_UFFD_RWP_RESOLVE)
+/* Whether a MM_CP_PROT_NUMA change is for promotion only */
+#define  MM_CP_PROT_NUMA_PROMO_ONLY        (1UL << 6)
 
 bool can_change_pte_writable(struct vm_area_struct *vma, unsigned long addr,
 			     pte_t pte);
@@ -4736,7 +4738,7 @@ void vma_set_file(struct vm_area_struct *vma, struct file *file);
 
 #ifdef CONFIG_NUMA_BALANCING
 unsigned long change_prot_numa(struct vm_area_struct *vma,
-			unsigned long start, unsigned long end);
+			unsigned long start, unsigned long end, bool promo_only);
 #endif
 
 struct vm_area_struct *find_extend_vma_locked(struct mm_struct *,
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 8dff37059faf..81359b414947 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -4129,8 +4129,10 @@ static void task_numa_work(struct callback_head *work)
 	unsigned long nr_pte_updates = 0;
 	long pages, virtpages;
 	struct vma_iterator vmi;
+	unsigned int numab_mode = READ_ONCE(sysctl_numa_balancing_mode);
 	bool vma_pids_skipped;
 	bool vma_pids_forced = false;
+	bool promo_only;
 
 	WARN_ON_ONCE(p != container_of(work, struct task_struct, numa_work));
 
@@ -4304,11 +4306,14 @@ static void task_numa_work(struct callback_head *work)
 			continue;
 		}
 
+		promo_only = !(numab_mode & NUMA_BALANCING_NORMAL);
+
 		do {
 			start = max(start, vma->vm_start);
 			end = ALIGN(start + (pages << PAGE_SHIFT), HPAGE_SIZE);
 			end = min(end, vma->vm_end);
-			nr_pte_updates = change_prot_numa(vma, start, end);
+			nr_pte_updates = change_prot_numa(vma, start, end,
+							  promo_only);
 
 			/*
 			 * Try to scan sysctl_numa_balancing_size worth of
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 30b7c63b0e35..2f9ada1bbcfc 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -2784,7 +2784,8 @@ int change_huge_pmd(struct mmu_gather *tlb, struct vm_area_struct *vma,
 			goto unlock;
 
 		if (!folio_can_map_prot_numa(pmd_folio(*pmd), vma,
-					     vma_is_single_threaded_private(vma)))
+					     vma_is_single_threaded_private(vma),
+					     cp_flags & MM_CP_PROT_NUMA_PROMO_ONLY))
 			goto unlock;
 	}
 	/*
diff --git a/mm/internal.h b/mm/internal.h
index 0dca33db068f..f3e093168f84 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -1237,11 +1237,12 @@ static inline bool vma_is_single_threaded_private(struct vm_area_struct *vma)
 
 #ifdef CONFIG_NUMA_BALANCING
 bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
-		bool is_private_single_threaded);
+		bool is_private_single_threaded, bool promo_only);
 
 #else
 static inline bool folio_can_map_prot_numa(struct folio *folio,
-		struct vm_area_struct *vma, bool is_private_single_threaded)
+		struct vm_area_struct *vma, bool is_private_single_threaded,
+		bool promo_only)
 {
 	return false;
 }
diff --git a/mm/mempolicy.c b/mm/mempolicy.c
index 2ad0a5f18280..a082ccfa09ec 100644
--- a/mm/mempolicy.c
+++ b/mm/mempolicy.c
@@ -846,6 +846,7 @@ static int queue_folios_hugetlb(pte_t *pte, unsigned long hmask,
  * @folio: The folio whose mapping considered for being made NUMA hintable
  * @vma: The VMA that the folio belongs to.
  * @is_private_single_threaded: Is this a single-threaded private VMA or not
+ * @promo_only: Whether this scan should only collect promotion candidates
  *
  * This function checks to see if the folio actually indicates that
  * we need to make the mapping one which causes a NUMA hinting fault,
@@ -855,7 +856,7 @@ static int queue_folios_hugetlb(pte_t *pte, unsigned long hmask,
  * Return: True if the mapping of the folio needs to be changed, false otherwise.
  */
 bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
-		bool is_private_single_threaded)
+		bool is_private_single_threaded, bool promo_only)
 {
 	int nid;
 
@@ -886,12 +887,8 @@ bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
 	if (is_private_single_threaded && (nid == numa_node_id()))
 		return false;
 
-	/*
-	 * Skip scanning top tier node if normal numa
-	 * balancing is disabled
-	 */
-	if (!(sysctl_numa_balancing_mode & NUMA_BALANCING_NORMAL) &&
-	    node_is_toptier(nid))
+	/* Promotion-only scans do not collect NUMA placement samples */
+	if (promo_only && node_is_toptier(nid))
 		return false;
 
 	if (folio_use_access_time(folio))
@@ -905,19 +902,26 @@ bool folio_can_map_prot_numa(struct folio *folio, struct vm_area_struct *vma,
  * These are later cleared by a NUMA hinting fault. Depending on these
  * faults, pages may be migrated for better NUMA placement.
  *
+ * With @promo_only only folios eligible for promotion are made
+ * hint-faultable, without also sampling for task placement.
+ *
  * This is assuming that NUMA faults are handled using PROT_NONE. If
  * an architecture makes a different choice, it will need further
  * changes to the core.
  */
 unsigned long change_prot_numa(struct vm_area_struct *vma,
-			unsigned long addr, unsigned long end)
+			unsigned long addr, unsigned long end, bool promo_only)
 {
+	unsigned long cp_flags = MM_CP_PROT_NUMA;
 	struct mmu_gather tlb;
 	long nr_updated;
 
+	if (promo_only)
+		cp_flags |= MM_CP_PROT_NUMA_PROMO_ONLY;
+
 	tlb_gather_mmu(&tlb, vma->vm_mm);
 
-	nr_updated = change_protection(&tlb, vma, addr, end, MM_CP_PROT_NUMA);
+	nr_updated = change_protection(&tlb, vma, addr, end, cp_flags);
 	if (nr_updated > 0) {
 		count_vm_numa_events(NUMA_PTE_UPDATES, nr_updated);
 		count_memcg_events_mm(vma->vm_mm, NUMA_PTE_UPDATES, nr_updated);
diff --git a/mm/mprotect.c b/mm/mprotect.c
index 2888ee638d87..4e66e4c7665b 100644
--- a/mm/mprotect.c
+++ b/mm/mprotect.c
@@ -337,6 +337,7 @@ static long change_pte_range(struct mmu_gather *tlb,
 	long pages = 0;
 	bool is_private_single_threaded;
 	bool prot_numa = cp_flags & MM_CP_PROT_NUMA;
+	bool numa_promo_only = cp_flags & MM_CP_PROT_NUMA_PROMO_ONLY;
 	bool uffd_rwp = cp_flags & MM_CP_UFFD_RWP;
 	bool uffd_wp = cp_flags & MM_CP_UFFD_WP;
 	int nr_ptes;
@@ -384,7 +385,8 @@ static long change_pte_range(struct mmu_gather *tlb,
 			 */
 			if (prot_numa &&
 			    !folio_can_map_prot_numa(folio, vma,
-						is_private_single_threaded)) {
+						is_private_single_threaded,
+						numa_promo_only)) {
 
 				/* determine batch to skip */
 				nr_ptes = mprotect_folio_pte_batch(folio,
-- 
2.55.0