[PATCH 4/9] mm/mglru: invalidate empty-walk skip on page fault and migration

empty_map skips an mm on a node for up to K generations after an empty walk.
Notify MGLRU when a page of the mm appears: set the node's bitmap
bit and clear its empty_map bit, so the next aging pass walks it again.

Major paths that map a page into the mm:
- mm/memory.c: do_anonymous_page()/finish_fault() (anon and file/COW faults),
  wp_page_copy() (anon COW), do_swap_page() (swap-in).
- mm/huge_memory.c: __do_huge_pmd_anonymous_page() (THP anon faults).
- mm/migrate.c: remove_migration_pte() (folio destination node, incl.
  NUMA-balancing migration).

The transitions are rare (the bitmap bit is usually already set and the
empty-marking bit usually clear), so each is guarded by test_bit() to avoid a
locked RMW bouncing the mm cache line on every fault; the periodic re-scan
covers any path not listed here.

Signed-off-by: Baoquan He <[email protected]>
---
 include/linux/mm_types.h | 19 +++++++++++++++++++
 mm/huge_memory.c         |  3 +++
 mm/memory.c              | 15 +++++++++++++++
 mm/migrate.c             |  4 ++++
 4 files changed, 41 insertions(+)

diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 3738e8877b73..94c1d1fedc9c 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -1521,6 +1521,21 @@ static inline void lru_gen_use_mm(struct mm_struct *mm)
 	WRITE_ONCE(mm->lru_gen.bitmap, -1);
 }
 
+/*
+ * A page of this mm appeared on node @nid (fault or migration): set the
+ * node's bitmap bit and clear the empty-walk skip. Probe first - the
+ * transitions are rare and unconditional locked RMWs would bounce the
+ * mm cache line.
+ */
+static inline void lru_gen_mm_accessed(struct mm_struct *mm, int nid)
+{
+	unsigned long key = nid % BITS_PER_TYPE(mm->lru_gen.bitmap);
+
+	if (!test_bit(key, &mm->lru_gen.bitmap))
+		set_bit(key, &mm->lru_gen.bitmap);
+	if (test_bit(key, &mm->lru_gen.empty_map))
+		clear_bit(key, &mm->lru_gen.empty_map);
+}
 #else /* !CONFIG_LRU_GEN_WALKS_MMU */
 
 static inline void lru_gen_add_mm(struct mm_struct *mm)
@@ -1543,6 +1558,10 @@ static inline void lru_gen_use_mm(struct mm_struct *mm)
 {
 }
 
+static inline void lru_gen_mm_accessed(struct mm_struct *mm, int nid)
+{
+}
+
 #endif /* CONFIG_LRU_GEN_WALKS_MMU */
 
 struct vma_iterator {
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index afbb5974bd22..23ae4626a980 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -1418,6 +1418,9 @@ static vm_fault_t __do_huge_pmd_anonymous_page(struct vm_fault *vmf)
 		map_anon_folio_pmd_pf(folio, vmf->pmd, vma, haddr);
 		mm_inc_nr_ptes(vma->vm_mm);
 		spin_unlock(vmf->ptl);
+		/* a new THP of this mm lands on this node */
+		if (lru_gen_enabled())
+			lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
 	}
 
 	return 0;
diff --git a/mm/memory.c b/mm/memory.c
index c54943302553..1d7ec0110ef2 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -4076,6 +4076,9 @@ static vm_fault_t wp_page_copy(struct vm_fault *vmf)
 		BUG_ON(unshare && pte_write(entry));
 		set_pte_at(mm, vmf->address, vmf->pte, entry);
 		update_mmu_cache_range(vmf, vma, vmf->address, vmf->pte, 1);
+		/* COW mapped a fresh folio of this mm on this node */
+		if (lru_gen_enabled())
+			lru_gen_mm_accessed(vma->vm_mm, folio_nid(new_folio));
 		if (old_folio) {
 			/*
 			 * Only after switching the pte to the new page may
@@ -5224,6 +5227,9 @@ vm_fault_t do_swap_page(struct vm_fault *vmf)
 	set_ptes(vma->vm_mm, address, ptep, pte, nr_pages);
 	arch_do_swap_page_nr(vma->vm_mm, vma, address,
 			pte, pte, nr_pages);
+	/* a swapped-in page of this mm lands on this node */
+	if (lru_gen_enabled())
+		lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
 
 	/*
 	 * Remove the swap entry and conditionally try to free up the swapcache.
@@ -5516,6 +5522,10 @@ static vm_fault_t do_anonymous_page(struct vm_fault *vmf)
 		folio_put(folio);
 		return handle_userfault(vmf, VM_UFFD_MISSING);
 	}
+	/* a new page of this mm lands on this node: invalidate any empty skip */
+	if (lru_gen_enabled())
+		lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
+
 	map_anon_folio_pte_pf(folio, vmf->pte, vma, addr,
 			      vmf_orig_pte_uffd_wp(vmf));
 unlock:
@@ -5776,6 +5786,11 @@ vm_fault_t finish_fault(struct vm_fault *vmf)
 		page = vmf->page;
 
 	folio = page_folio(page);
+
+	/* mapping a page of this mm on this node: invalidate any empty skip */
+	if (lru_gen_enabled())
+		lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
+
 	/*
 	 * check even for read faults because we might have lost our CoWed
 	 * page
diff --git a/mm/migrate.c b/mm/migrate.c
index 15b45832bcfa..26b8b34988d2 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -354,6 +354,10 @@ static bool remove_migration_pte(struct folio *folio,
 	struct rmap_walk_arg *rmap_walk_arg = arg;
 	DEFINE_FOLIO_VMA_WALK(pvmw, rmap_walk_arg->folio, vma, addr, PVMW_SYNC | PVMW_MIGRATION);
 
+	/* the folio ends up on folio_nid(): notify MGLRU for this mm */
+	if (lru_gen_enabled())
+		lru_gen_mm_accessed(vma->vm_mm, folio_nid(folio));
+
 	while (page_vma_mapped_walk(&pvmw)) {
 		rmap_t rmap_flags = RMAP_NONE;
 		unsigned long idx = 0;
-- 
2.54.0