On Wed, Oct 22, 2025 at 12:37:09PM -0600, Nico Pache wrote:

Pass an order and offset 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, and offset indicates were in the PMD to
start the collapse attempt.

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

NIT but is this -- a typo?

Presumably pmd_address is the PMD-aligned address?

+		int referenced, int unmapped, struct collapse_control *cc,
+		bool *mmap_locked, unsigned int order, unsigned long offset)

It'd be nice to pass through a helper struct at this point having so many params
but perhaps we can deal with that in a follow up series.

If PMD address is the PMD-aligned address, and mthp_address = pmd_address +
offset * PAGE_SIZE, couldn't we just pass in the mthp address and get the
PMD address by aligning down to PMD size and reduce the number of args by
1?

 {
 	LIST_HEAD(compound_pagelist);
 	pmd_t *pmd, _pmd;
-	pte_t *pte;
+	pte_t *pte = NULL, mthp_pte;

mthp_pte is only used in a single if () branch and can be declared there
AFAICT?

 	pgtable_t pgtable;
 	struct folio *folio;
 	spinlock_t *pmd_ptl, *pte_ptl;
 	int result = SCAN_FAIL;
 	struct vm_area_struct *vma;
 	struct mmu_notifier_range range;
+	bool anon_vma_locked = false;
+	const unsigned long nr_pages = 1UL << order;
+	unsigned long mthp_address = pmd_address + offset * PAGE_SIZE;

Do we ever update this? If not we can const-ify.

-	VM_BUG_ON(address & ~HPAGE_PMD_MASK);
+	VM_BUG_ON(pmd_address & ~HPAGE_PMD_MASK);

NIT: Be nice to convert this to a VM_WARN_ON_ONCE(), as VM_BUG_ON() is not
right here.

 	/*
 	 * Before allocating the hugepage, release the mmap_lock read lock.
 	 * The allocation can take potentially a long time if it involves
 	 * sync compaction, and we do not need to hold the mmap_lock during
 	 * that. We will recheck the vma after taking it again in write mode.
+	 * If collapsing mTHPs we may have already released the read_lock.
 	 */
-	mmap_read_unlock(mm);
+	if (*mmap_locked) {
+		mmap_read_unlock(mm);
+		*mmap_locked = false;
+	}

-	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;

 	mmap_read_lock(mm);
-	result = hugepage_vma_revalidate(mm, address, true, &vma, cc,
-					 HPAGE_PMD_ORDER);
+	*mmap_locked = true;
+	result = hugepage_vma_revalidate(mm, pmd_address, true, &vma, cc, order);
 	if (result != SCAN_SUCCEED) {
 		mmap_read_unlock(mm);

I don't really love the semantics of 'sometimes we set *mmap_locked false
when we unlock, sometimes we rely on out_nolock doing it'.

Let's just set it false when we unlock and VM_WARN_ON_ONCE(*mmap_locked) in
out_nolock.

Not your fault but so hate this _p** convention. One for a follow up I
suppose.

 	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, mthp_address, &pte_ptl);
 	if (pte) {
-		result = __collapse_huge_page_isolate(vma, address, pte, cc,
-						      HPAGE_PMD_ORDER,
-						      &compound_pagelist);
+		result = __collapse_huge_page_isolate(vma, mthp_address, pte, cc,
+						      order, &compound_pagelist);
 		spin_unlock(pte_ptl);
 	} else {
 		result = SCAN_PMD_NULL;
 	}

 	if (unlikely(result != SCAN_SUCCEED)) {
-		if (pte)
-			pte_unmap(pte);

OK I guess we drop this because it's handled in out_up_write. I assume no
issue keeping PTE mapped here?

Looking through __collapse_huge_page_copy() there doesn't seem to be any
issue with holding anon lock here.

As per above, let's just declare mthp_pte here.

+		mthp_pte = mk_pte(folio_page(folio, 0), vma->vm_page_prot);

Hm, so we make a PTE that references the first page of the folio? I guess
the folio will be an mTHP folio so we're just creating essentially a

+		mthp_pte = maybe_mkwrite(pte_mkdirty(mthp_pte), vma);

In set_pte_range() we have a whole host of other checks like dirty,
uffd_wp, etc. I wonder if we need to consider those?

+
+		spin_lock(pmd_ptl);

We're duplicating this in both branches, why not do outside if/else?

+		WARN_ON_ONCE(!pmd_none(*pmd));

Hmm so the PMD entry will still always be empty on mTHP collapse? Surely we
could be collapsing more than one mTHP into an existing PTE table no? I may
be missing something here/confused :)

+		folio_ref_add(folio, nr_pages - 1);

If we're setting the refcount here, where is the ref count being set in the
PMD path?

+		folio_add_new_anon_rmap(folio, vma, mthp_address, RMAP_EXCLUSIVE);
+		folio_add_lru_vma(folio, vma);
+		set_ptes(vma->vm_mm, mthp_address, pte, mthp_pte, nr_pages);
+		update_mmu_cache_range(NULL, vma, mthp_address, pte, nr_pages);

Prior to this change the only user of this are functions in memory.c, I
do wonder if this is the wrong abstraction here.

But maybe that's _yet another_ thing for a follow up (the THP code is a
mess).

+
+		smp_wmb(); /* make PTEs visible before PMD. See pmd_install() */

Feels like we could avoid open-coding this by just using pmd_install()?

Also are we therefore missing a mm_inc_nr_ptes() invocation here, or do we
update mm->pgtables_bytes elsewhere?

+		pmd_populate(mm, pmd, pmd_pgtable(_pmd));

Why are we referencing pmd in PMD branch and _pmd here?

+		spin_unlock(pmd_ptl);

The PMD case does this stuff in map_anon_pmd_nopf(), could we add one for
mTHP?

This function is already horribly overwrought (not your fault) so I'd like
to avoid adding open-coded blocks as much as possible.

+	}

 	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:
+	*mmap_locked = false;

See above comment about setting this prior to jumping to out_nolock.