khugepaged: optimize __collapse_huge_page_copy_succeeded() by PTE batching

Use PTE batching to batch process PTEs mapping the same large folio. An
improvement is expected due to batching refcount-mapcount manipulation on
the folios, and for arm64 which supports contig mappings, the number of
TLB flushes is also reduced.

Link: https://lkml.kernel.org/r/20250724052301.23844-3-dev.jain@arm.com
Signed-off-by: Dev Jain <dev.jain@arm.com>
Acked-by: David Hildenbrand <david@redhat.com>
Reviewed-by: Baolin Wang <baolin.wang@linux.alibaba.com>
Reviewed-by: Lorenzo Stoakes <lorenzo.stoakes@oracle.com>
Cc: Barry Song <baohua@kernel.org>
Cc: Liam Howlett <liam.howlett@oracle.com>
Cc: Mariano Pache <npache@redhat.com>
Cc: Ryan Roberts <ryan.roberts@arm.com>
Cc: Zi Yan <ziy@nvidia.com>
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index a55fb1dcd22451991b7770536af85524aedf39a1..f23e943506bc7fcfefd84141afb1934ca6efda69 100644 (file)
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -700,12 +700,15 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
                                                spinlock_t *ptl,
                                                struct list_head *compound_pagelist)
 {
+       unsigned long end = address + HPAGE_PMD_SIZE;
        struct folio *src, *tmp;
-       pte_t *_pte;
        pte_t pteval;
+       pte_t *_pte;
+       unsigned int nr_ptes;
 
-       for (_pte = pte; _pte < pte + HPAGE_PMD_NR;
-            _pte++, address += PAGE_SIZE) {
+       for (_pte = pte; _pte < pte + HPAGE_PMD_NR; _pte += nr_ptes,
+            address += nr_ptes * PAGE_SIZE) {
+               nr_ptes = 1;
                pteval = ptep_get(_pte);
                if (pte_none(pteval) || is_zero_pfn(pte_pfn(pteval))) {
                        add_mm_counter(vma->vm_mm, MM_ANONPAGES, 1);
@@ -722,18 +725,26 @@ static void __collapse_huge_page_copy_succeeded(pte_t *pte,
                        struct page *src_page = pte_page(pteval);
 
                        src = page_folio(src_page);
-                       if (!folio_test_large(src))
+
+                       if (folio_test_large(src)) {
+                               unsigned int max_nr_ptes = (end - address) >> PAGE_SHIFT;
+
+                               nr_ptes = folio_pte_batch(src, _pte, pteval, max_nr_ptes);
+                       } else {
                                release_pte_folio(src);
+                       }
+
                        /*
                         * ptl mostly unnecessary, but preempt has to
                         * be disabled to update the per-cpu stats
                         * inside folio_remove_rmap_pte().
                         */
                        spin_lock(ptl);
-                       ptep_clear(vma->vm_mm, address, _pte);
-                       folio_remove_rmap_pte(src, src_page, vma);
+                       clear_ptes(vma->vm_mm, address, _pte, nr_ptes);
+                       folio_remove_rmap_ptes(src, src_page, nr_ptes, vma);
                        spin_unlock(ptl);
-                       free_folio_and_swap_cache(src);
+                       free_swap_cache(src);
+                       folio_put_refs(src, nr_ptes);
                }
        }