[PATCH mm-stable] mm/khugepaged: avoid unnecessary checking for swap entries when collapsing a mTHP

[Pedro Demarchi Gomes <pedrodemargomes@gmail.com>]

mthp_collapse() tries to swap in PTEs when collapsing a mTHP if there are any
swap PTEs in the PMD range, even if none of those swap PTEs are
actually part of the mTHP's range.

Track swap PTEs in cc->mthp_unmapped_ptes so that mthp_collapse() can
tell whether the mTHP's own range contains a swap PTE, and skip the
swapin check when it doesn't.

Signed-off-by: Pedro Demarchi Gomes <pedrodemargomes@gmail.com>
---
 mm/khugepaged.c | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 11ff98d55c76..38596172f7dc 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -117,6 +117,9 @@ struct collapse_control {
 
 	/* Each bit represents a single occupied (!none/zero) page. */
 	DECLARE_BITMAP(mthp_present_ptes, MAX_PTRS_PER_PTE);
+
+	/* Each bit represents a single not present and not none/zero pte. */
+	DECLARE_BITMAP(mthp_unmapped_ptes, MAX_PTRS_PER_PTE);
 };
 
 /**
@@ -634,6 +637,7 @@ static void collapse_control_init_scan(struct collapse_control *cc)
 	memset(cc->node_load, 0, sizeof(cc->node_load));
 	nodes_clear(cc->alloc_nmask);
 	bitmap_zero(cc->mthp_present_ptes, MAX_PTRS_PER_PTE);
+	bitmap_zero(cc->mthp_unmapped_ptes, MAX_PTRS_PER_PTE);
 }
 
 static void release_pte_folio(struct folio *folio)
@@ -1291,7 +1295,7 @@ static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_stru
  * Note that the VMA must be rechecked after grabbing the mmap_lock again.
  */
 static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long start_addr,
-		int referenced, int unmapped, struct collapse_control *cc,
+		int referenced, bool swapin, struct collapse_control *cc,
 		unsigned int order)
 {
 	const unsigned long pmd_addr = start_addr & HPAGE_PMD_MASK;
@@ -1330,7 +1334,7 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long s
 		goto out_nolock;
 	}
 
-	if (unmapped) {
+	if (swapin) {
 		/*
 		 * __collapse_huge_page_swapin() will return with mmap_lock
 		 * released when it fails. So we jump out_nolock directly in
@@ -1502,10 +1506,10 @@ static unsigned int max_order_from_offset(unsigned int offset)
  * mTHP.
  */
 static enum scan_result mthp_collapse(struct mm_struct *mm,
-		unsigned long address, int referenced, int unmapped,
+		unsigned long address, int referenced,
 		struct collapse_control *cc, unsigned long enabled_orders)
 {
-	unsigned int nr_occupied_ptes, nr_ptes, max_ptes_none;
+	unsigned int nr_occupied_ptes, nr_unmapped_ptes, nr_ptes, max_ptes_none;
 	enum scan_result last_result = SCAN_FAIL;
 	int collapsed = 0;
 	bool alloc_failed = false;
@@ -1522,21 +1526,25 @@ static enum scan_result mthp_collapse(struct mm_struct *mm,
 		max_ptes_none = collapse_max_ptes_none(cc, NULL, order);
 		nr_occupied_ptes = bitmap_weight_from(cc->mthp_present_ptes, offset,
 						      offset + nr_ptes);
+		nr_unmapped_ptes = bitmap_weight_from(cc->mthp_unmapped_ptes, offset,
+						      offset + nr_ptes);
+
 
 		/*
-		 * Swap PTEs accepted during the scan are counted in @unmapped,
-		 * not in the present-PTE bitmap. Account them for the PMD-order
-		 * candidate.
+		 * Swap PTEs accepted during the scan are counted in
+		 * nr_unmapped_ptes, not in the present-PTE bitmap. Account
+		 * them for the PMD-order candidate.
 		 */
 		if (is_pmd_order(order))
-			nr_occupied_ptes += unmapped;
+			nr_occupied_ptes += nr_unmapped_ptes;
 
 		if (nr_occupied_ptes >= nr_ptes - max_ptes_none) {
 			enum scan_result ret;
+			bool swapin = nr_unmapped_ptes > 0;
 
 			collapse_address = address + offset * PAGE_SIZE;
 			ret = collapse_huge_page(mm, collapse_address, referenced,
-						 unmapped, cc, order);
+						 swapin, cc, order);
 
 			switch (ret) {
 			/* Cases where we continue to next collapse candidate */
@@ -1667,6 +1675,7 @@ static enum scan_result collapse_scan_pmd(struct mm_struct *mm,
 			continue;
 		}
 		if (!pte_present(pteval)) {
+			__set_bit(i, cc->mthp_unmapped_ptes);
 			if (++unmapped > max_ptes_swap) {
 				result = SCAN_EXCEED_SWAP_PTE;
 				count_collapse_event(HPAGE_PMD_ORDER, THP_SCAN_EXCEED_SWAP_PTE,
@@ -1785,7 +1794,7 @@ static enum scan_result collapse_scan_pmd(struct mm_struct *mm,
 		/* collapse_huge_page() expects the lock to be dropped before calling */
 		mmap_read_unlock(mm);
 		result = mthp_collapse(mm, start_addr, referenced,
-				       unmapped, cc, enabled_orders);
+				       cc, enabled_orders);
 		/* mmap_lock was released above, set lock_dropped */
 		*lock_dropped = true;
 	}
-- 
2.55.0