mm/khugepaged: generalize collapse_huge_page for mTHP collapse

Pass an order to collapse_huge_page to support collapsing anon memory to
arbitrary orders within a PMD.  order indicates what mTHP size we are
attempting to collapse to.

For non-PMD collapse we must leave the anon VMA write locked until after
we collapse the mTHP-- in the PMD case all the pages are isolated, but in
the mTHP case this is not true, and we must keep the lock to prevent
access/changes to the page tables.  This can happen if the rmap walkers
hit a pmd_none while the PMD entry is currently unavailable due to being
temporarily removed during the collapse phase.

To properly establish the page table hierarchy without violating any
expectations from certain architectures (e.g.  MIPS), we must make sure to
have the PMD reinstalled before the PTEs, and hold both PTE/PMD locks
before calling update_mmu_cache_range() (if they are distinct locks).

Link: https://lore.kernel.org/20260605161422.213817-7-npache@redhat.com
Signed-off-by: Nico Pache <npache@redhat.com>
Acked-by: David Hildenbrand (Arm) <david@kernel.org>
Reviewed-by: Lorenzo Stoakes <ljs@kernel.org>
Reviewed-by: Lance Yang <lance.yang@linux.dev>
Cc: Alistair Popple <apopple@nvidia.com>
Cc: Andrea Arcangeli <aarcange@redhat.com>
Cc: Anshuman Khandual <anshuman.khandual@arm.com>
Cc: Bagas Sanjaya <bagasdotme@gmail.com>
Cc: Baolin Wang <baolin.wang@linux.alibaba.com>
Cc: Barry Song <baohua@kernel.org>
Cc: Brendan Jackman <jackmanb@google.com>
Cc: Byungchul Park <byungchul@sk.com>
Cc: Catalin Marinas <catalin.marinas@arm.com>
Cc: David Rientjes <rientjes@google.com>
Cc: Dev Jain <dev.jain@arm.com>
Cc: Gregory Price <gourry@gourry.net>
Cc: "Huang, Ying" <ying.huang@linux.alibaba.com>
Cc: Hugh Dickins <hughd@google.com>
Cc: Jan Kara <jack@suse.cz>
Cc: Jann Horn <jannh@google.com>
Cc: Johannes Weiner <hannes@cmpxchg.org>
Cc: Jonathan Corbet <corbet@lwn.net>
Cc: Joshua Hahn <joshua.hahnjy@gmail.com>
Cc: Kefeng Wang <wangkefeng.wang@huawei.com>
Cc: Liam R. Howlett <liam@infradead.org>
Cc: "Masami Hiramatsu (Google)" <mhiramat@kernel.org>
Cc: Mathieu Desnoyers <mathieu.desnoyers@efficios.com>
Cc: Matthew Brost <matthew.brost@intel.com>
Cc: Matthew Wilcox (Oracle) <willy@infradead.org>
Cc: Michal Hocko <mhocko@suse.com>
Cc: Mike Rapoport <rppt@kernel.org>
Cc: Nanyong Sun <sunnanyong@huawei.com>
Cc: Pedro Falcato <pfalcato@suse.de>
Cc: Peter Xu <peterx@redhat.com>
Cc: Rafael Aquini <raquini@redhat.com>
Cc: Rakie Kim <rakie.kim@sk.com>
Cc: Randy Dunlap <rdunlap@infradead.org>
Cc: Ryan Roberts <ryan.roberts@arm.com>
Cc: Shivank Garg <shivankg@amd.com>
Cc: Steven Rostedt <rostedt@goodmis.org>
Cc: Suren Baghdasaryan <surenb@google.com>
Cc: Takashi Iwai (SUSE) <tiwai@suse.de>
Cc: Thomas Hellström <thomas.hellstrom@linux.intel.com>
Cc: Usama Arif <usamaarif642@gmail.com>
Cc: Usama Arif <usama.arif@linux.dev>
Cc: Vishal Moola (Oracle) <vishal.moola@gmail.com>
Cc: Vlastimil Babka <vbabka@suse.cz>
Cc: Wei Yang <richard.weiyang@gmail.com>
Cc: Will Deacon <will@kernel.org>
Cc: Yang Shi <yang@os.amperecomputing.com>
Cc: Zach O'Keefe <zokeefe@google.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 72b5d3671abc818e82e59ff931072778d9d898ee..acbb925b01235bf1d41a60e1732517f7dfcfd044 100644 (file)
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -1228,22 +1228,24 @@ static enum scan_result alloc_charge_folio(struct folio **foliop, struct mm_stru
  * while allocating a THP, as that could trigger direct reclaim/compaction.
  * 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 address,
-               int referenced, int unmapped, struct collapse_control *cc)
+static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long start_addr,
+               int referenced, int unmapped, struct collapse_control *cc,
+               unsigned int order)
 {
+       const unsigned long pmd_addr = start_addr & HPAGE_PMD_MASK;
+       const unsigned long end_addr = start_addr + (PAGE_SIZE << order);
        LIST_HEAD(compound_pagelist);
        pmd_t *pmd, _pmd;
-       pte_t *pte;
+       pte_t *pte = NULL;
        pgtable_t pgtable;
        struct folio *folio;
        spinlock_t *pmd_ptl, *pte_ptl;
        enum scan_result result = SCAN_FAIL;
        struct vm_area_struct *vma;
        struct mmu_notifier_range range;
+       bool anon_vma_locked = false;
 
-       VM_BUG_ON(address & ~HPAGE_PMD_MASK);
-
-       result = alloc_charge_folio(&folio, mm, cc, HPAGE_PMD_ORDER);
+       result = alloc_charge_folio(&folio, mm, cc, order);
        if (result != SCAN_SUCCEED)
                goto out_nolock;
 
@@ -1253,14 +1255,14 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long a
        }
 
        mmap_read_lock(mm);
-       result = hugepage_vma_revalidate(mm, address, true, &vma, cc,
-                                        HPAGE_PMD_ORDER);
+       result = hugepage_vma_revalidate(mm, pmd_addr, /*expect_anon=*/ true,
+                                        &vma, cc, order);
        if (result != SCAN_SUCCEED) {
                mmap_read_unlock(mm);
                goto out_nolock;
        }
 
-       result = find_pmd_or_thp_or_none(mm, address, &pmd);
+       result = find_pmd_or_thp_or_none(mm, pmd_addr, &pmd);
        if (result != SCAN_SUCCEED) {
                mmap_read_unlock(mm);
                goto out_nolock;
@@ -1272,8 +1274,8 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long a
                 * released when it fails. So we jump out_nolock directly in
                 * that case.  Continuing to collapse causes inconsistency.
                 */
-               result = __collapse_huge_page_swapin(mm, vma, address, pmd,
-                                                    referenced, HPAGE_PMD_ORDER);
+               result = __collapse_huge_page_swapin(mm, vma, start_addr, pmd,
+                                                    referenced, order);
                if (result != SCAN_SUCCEED)
                        goto out_nolock;
        }
@@ -1288,20 +1290,28 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long a
         * mmap_lock.
         */
        mmap_write_lock(mm);
-       result = hugepage_vma_revalidate(mm, address, true, &vma, cc,
-                                        HPAGE_PMD_ORDER);
+       result = hugepage_vma_revalidate(mm, pmd_addr, /*expect_anon=*/ true,
+                                        &vma, cc, order);
        if (result != SCAN_SUCCEED)
                goto out_up_write;
        /* check if the pmd is still valid */
        vma_start_write(vma);
-       result = check_pmd_still_valid(mm, address, pmd);
+       result = check_pmd_still_valid(mm, pmd_addr, pmd);
        if (result != SCAN_SUCCEED)
                goto out_up_write;
 
        anon_vma_lock_write(vma->anon_vma);
+       anon_vma_locked = true;
 
-       mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm, address,
-                               address + HPAGE_PMD_SIZE);
+       /*
+        * Only notify about the PTE range we will actually modify. While we
+        * temporary unmap the whole PTE table for mTHP collapse, we'll remap
+        * it later, leaving other PTEs effectively unmodified. The locks we
+        * hold prevent anybody from stumbling over such temporarily unmapped
+        * PTE tables.
+        */
+       mmu_notifier_range_init(&range, MMU_NOTIFY_CLEAR, 0, mm, start_addr,
+                               end_addr);
        mmu_notifier_invalidate_range_start(&range);
 
        pmd_ptl = pmd_lock(mm, pmd); /* probably unnecessary */
@@ -1313,26 +1323,23 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long a
         * Parallel GUP-fast is fine since GUP-fast will back off when
         * it detects PMD is changed.
         */
-       _pmd = pmdp_collapse_flush(vma, address, pmd);
+       _pmd = pmdp_collapse_flush(vma, pmd_addr, pmd);
        spin_unlock(pmd_ptl);
        mmu_notifier_invalidate_range_end(&range);
        tlb_remove_table_sync_one();
 
-       pte = pte_offset_map_lock(mm, &_pmd, address, &pte_ptl);
+       pte = pte_offset_map_lock(mm, &_pmd, start_addr, &pte_ptl);
        if (pte) {
-               result = __collapse_huge_page_isolate(vma, address, pte, cc,
-                                                     HPAGE_PMD_ORDER,
-                                                     &compound_pagelist);
+               result = __collapse_huge_page_isolate(vma, start_addr, pte, cc,
+                                                     order, &compound_pagelist);
                spin_unlock(pte_ptl);
        } else {
                result = SCAN_NO_PTE_TABLE;
        }
 
        if (unlikely(result != SCAN_SUCCEED)) {
-               if (pte)
-                       pte_unmap(pte);
                spin_lock(pmd_ptl);
-               BUG_ON(!pmd_none(*pmd));
+               VM_WARN_ON_ONCE(!pmd_none(*pmd));
                /*
                 * We can only use set_pmd_at when establishing
                 * hugepmds and never for establishing regular pmds that
@@ -1340,21 +1347,24 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long a
                 */
                pmd_populate(mm, pmd, pmd_pgtable(_pmd));
                spin_unlock(pmd_ptl);
-               anon_vma_unlock_write(vma->anon_vma);
                goto out_up_write;
        }
 
        /*
-        * All pages are isolated and locked so anon_vma rmap
-        * can't run anymore.
+        * For PMD collapse all pages are isolated and locked so anon_vma
+        * rmap can't run anymore. For mTHP collapse the PMD entry has been
+        * removed and not all pages are isolated and locked, so we must hold
+        * the lock to prevent neighboring folios from attempting to access
+        * this PMD until its reinstalled.
         */
-       anon_vma_unlock_write(vma->anon_vma);
+       if (is_pmd_order(order)) {
+               anon_vma_unlock_write(vma->anon_vma);
+               anon_vma_locked = false;
+       }
 
        result = __collapse_huge_page_copy(pte, folio, pmd, _pmd,
-                                          vma, address, pte_ptl,
-                                          HPAGE_PMD_ORDER,
-                                          &compound_pagelist);
-       pte_unmap(pte);
+                                          vma, start_addr, pte_ptl,
+                                          order, &compound_pagelist);
        if (unlikely(result != SCAN_SUCCEED))
                goto out_up_write;
 
@@ -1364,18 +1374,37 @@ static enum scan_result collapse_huge_page(struct mm_struct *mm, unsigned long a
         * write.
         */
        __folio_mark_uptodate(folio);
-       pgtable = pmd_pgtable(_pmd);
-
        spin_lock(pmd_ptl);
-       BUG_ON(!pmd_none(*pmd));
-       pgtable_trans_huge_deposit(mm, pmd, pgtable);
-       map_anon_folio_pmd_nopf(folio, pmd, vma, address);
+       VM_WARN_ON_ONCE(!pmd_none(*pmd));
+       if (is_pmd_order(order)) {
+               pgtable = pmd_pgtable(_pmd);
+               pgtable_trans_huge_deposit(mm, pmd, pgtable);
+               map_anon_folio_pmd_nopf(folio, pmd, vma, pmd_addr);
+       } else {
+               /*
+                * Some architectures (e.g. MIPS) walk the live page table in
+                * their implementation. update_mmu_cache_range() must be called
+                * with a valid page table hierarchy and the PTE lock held.
+                * Acquire it nested inside pmd_ptl when they are distinct locks.
+                */
+               if (pte_ptl != pmd_ptl)
+                       spin_lock_nested(pte_ptl, SINGLE_DEPTH_NESTING);
+               pmd_populate(mm, pmd, pmd_pgtable(_pmd));
+               map_anon_folio_pte_nopf(folio, pte, vma, start_addr,
+                                         /*uffd_wp=*/ false);
+               if (pte_ptl != pmd_ptl)
+                       spin_unlock(pte_ptl);
+       }
        spin_unlock(pmd_ptl);
 
        folio = NULL;
 
        result = SCAN_SUCCEED;
 out_up_write:
+       if (anon_vma_locked)
+               anon_vma_unlock_write(vma->anon_vma);
+       if (pte)
+               pte_unmap(pte);
        mmap_write_unlock(mm);
 out_nolock:
        if (folio)
@@ -1555,7 +1584,7 @@ out_unmap:
                /* collapse_huge_page expects the lock to be dropped before calling */
                mmap_read_unlock(mm);
                result = collapse_huge_page(mm, start_addr, referenced,
-                                           unmapped, cc);
+                                           unmapped, cc, HPAGE_PMD_ORDER);
                /* collapse_huge_page will return with the mmap_lock released */
                *lock_dropped = true;
        }