Re: [PATCH mm-new v3 3/3] mm/khugepaged: merge PTE scanning logic into a new helper

[Wei Yang <richard.weiyang@gmail.com>]

On Wed, Oct 08, 2025 at 12:37:48PM +0800, Lance Yang wrote:
>From: Lance Yang <lance.yang@linux.dev>
>
>As David suggested, the PTE scanning logic in hpage_collapse_scan_pmd()
>and __collapse_huge_page_isolate() was almost duplicated.
>
>This patch cleans things up by moving all the common PTE checking logic
>into a new shared helper, thp_collapse_check_pte(). While at it, we use
>vm_normal_folio() instead of vm_normal_page().
>
>Suggested-by: David Hildenbrand <david@redhat.com>
>Suggested-by: Dev Jain <dev.jain@arm.com>
>Signed-off-by: Lance Yang <lance.yang@linux.dev>
>---
> mm/khugepaged.c | 243 ++++++++++++++++++++++++++----------------------
> 1 file changed, 130 insertions(+), 113 deletions(-)
>
>diff --git a/mm/khugepaged.c b/mm/khugepaged.c
>index b5c0295c3414..7116caae1fa4 100644
>--- a/mm/khugepaged.c
>+++ b/mm/khugepaged.c
>@@ -61,6 +61,12 @@ enum scan_result {
> 	SCAN_PAGE_FILLED,
> };
> 
>+enum pte_check_result {
>+	PTE_CHECK_SUCCEED,
>+	PTE_CHECK_CONTINUE,
>+	PTE_CHECK_FAIL,
>+};
>+
> #define CREATE_TRACE_POINTS
> #include <trace/events/huge_memory.h>
> 
>@@ -533,62 +539,139 @@ static void release_pte_pages(pte_t *pte, pte_t *_pte,
> 	}
> }
> 
>+/*
>+ * thp_collapse_check_pte - Check if a PTE is suitable for THP collapse
>+ * @pte:           The PTE to check
>+ * @vma:           The VMA the PTE belongs to
>+ * @addr:          The virtual address corresponding to this PTE
>+ * @foliop:        On success, used to return a pointer to the folio
>+ *                 Must be non-NULL
>+ * @none_or_zero:  Counter for none/zero PTEs. Must be non-NULL
>+ * @unmapped:      Counter for swap PTEs. Can be NULL if not scanning swaps
>+ * @shared:        Counter for shared pages. Must be non-NULL
>+ * @scan_result:   Used to return the failure reason (SCAN_*) on a
>+ *                 PTE_CHECK_FAIL return. Must be non-NULL
>+ * @cc:            Collapse control settings
>+ *
>+ * Returns:
>+ *   PTE_CHECK_SUCCEED  - PTE is suitable, proceed with further checks
>+ *   PTE_CHECK_CONTINUE - Skip this PTE and continue scanning
>+ *   PTE_CHECK_FAIL     - Abort collapse scan
>+ */
>+static inline int thp_collapse_check_pte(pte_t pte, struct vm_area_struct *vma,
>+		unsigned long addr, struct folio **foliop, int *none_or_zero,
>+		int *unmapped, int *shared, int *scan_result,
>+		struct collapse_control *cc)
>+{
>+	struct folio *folio = NULL;
>+
>+	if (pte_none(pte) || is_zero_pfn(pte_pfn(pte))) {
>+		(*none_or_zero)++;
>+		if (!userfaultfd_armed(vma) &&
>+		    (!cc->is_khugepaged ||
>+		     *none_or_zero <= khugepaged_max_ptes_none)) {
>+			return PTE_CHECK_CONTINUE;
>+		} else {
>+			*scan_result = SCAN_EXCEED_NONE_PTE;
>+			count_vm_event(THP_SCAN_EXCEED_NONE_PTE);
>+			return PTE_CHECK_FAIL;
>+		}
>+	} else if (!pte_present(pte)) {
>+		if (!unmapped) {
>+			*scan_result = SCAN_PTE_NON_PRESENT;
>+			return PTE_CHECK_FAIL;
>+		}
>+
>+		if (non_swap_entry(pte_to_swp_entry(pte))) {
>+			*scan_result = SCAN_PTE_NON_PRESENT;
>+			return PTE_CHECK_FAIL;
>+		}
>+
>+		(*unmapped)++;
>+		if (!cc->is_khugepaged ||
>+		    *unmapped <= khugepaged_max_ptes_swap) {
>+			/*
>+			 * Always be strict with uffd-wp enabled swap
>+			 * entries. Please see comment below for
>+			 * pte_uffd_wp().
>+			 */
>+			if (pte_swp_uffd_wp(pte)) {
>+				*scan_result = SCAN_PTE_UFFD_WP;
>+				return PTE_CHECK_FAIL;
>+			}
>+			return PTE_CHECK_CONTINUE;
>+		} else {
>+			*scan_result = SCAN_EXCEED_SWAP_PTE;
>+			count_vm_event(THP_SCAN_EXCEED_SWAP_PTE);
>+			return PTE_CHECK_FAIL;
>+		}
>+	} else if (pte_uffd_wp(pte)) {
>+		/*
>+		 * Don't collapse the page if any of the small PTEs are
>+		 * armed with uffd write protection. Here we can also mark
>+		 * the new huge pmd as write protected if any of the small
>+		 * ones is marked but that could bring unknown userfault
>+		 * messages that falls outside of the registered range.
>+		 * So, just be simple.
>+		 */
>+		*scan_result = SCAN_PTE_UFFD_WP;
>+		return PTE_CHECK_FAIL;
>+	}
>+
>+	folio = vm_normal_folio(vma, addr, pte);
>+	if (unlikely(!folio) || unlikely(folio_is_zone_device(folio))) {
>+		*scan_result = SCAN_PAGE_NULL;
>+		return PTE_CHECK_FAIL;
>+	}
>+
>+	if (!folio_test_anon(folio)) {
>+		VM_WARN_ON_FOLIO(true, folio);
>+		*scan_result = SCAN_PAGE_ANON;
>+		return PTE_CHECK_FAIL;
>+	}
>+
>+	/*
>+	 * We treat a single page as shared if any part of the THP
>+	 * is shared.
>+	 */
>+	if (folio_maybe_mapped_shared(folio)) {
>+		(*shared)++;
>+		if (cc->is_khugepaged && *shared > khugepaged_max_ptes_shared) {
>+			*scan_result = SCAN_EXCEED_SHARED_PTE;
>+			count_vm_event(THP_SCAN_EXCEED_SHARED_PTE);
>+			return PTE_CHECK_FAIL;
>+		}
>+	}
>+
>+	*foliop = folio;
>+
>+	return PTE_CHECK_SUCCEED;
>+}
>+

This one looks much better.

While my personal feeling is this is not a complete work to merge the scanning
logic. We still have folio_expected_ref_count() and pte_young() check present
both in __collapse_huge_page_isolate() and huge_collapse_scan_pmd().

-- 
Wei Yang
Help you, Help me