Re: [PATCH v4 10/13] mm/collapse: open-code collapse_single_pmd() in its two callers

["David Hildenbrand (Arm)" <david@kernel.org>]

On 9/28/26 12:06, Kiryl Shutsemau wrote:
> From: "Kiryl Shutsemau (Meta)" <kas@kernel.org>
> 
> A scan and a collapse want different things from mmap_lock.  The scan
> reads one PTE table under the lock the caller holds, refuses most of the
> time, and the caller moves on to the next table without letting go.  The
> collapse allocates, may sleep in writeback and takes the lock for write
> itself, so the lock it is handed is of no use to it.
> 
> collapse_single_pmd() kept that boundary inside itself.  It dropped the
> lock on some paths and not others, and reported which by way of a bool
> its callers had to carry along and then act on. 

Let me think this through. collapse_scan_file: it doesn't actually need the MM at
all. The only reason is to do tracing. Rather stupid, it just should not consume
the MM at all anymore. Consequently it doesn't even need the mmap lock. But the
caller needs the mmap lock to figure out the file + range from the vma (the
per-vma lock would also be sufficient for that).

Also, I guess we can convert some of the scanning to use per-vma locks in the
future, whereby we would actually want to scan with the per-vma lock held.

I do wonder about one thing: should we really care so much about keeping the
mmap lock locked? Meaning, why not provide a single collapse_single_pmd() that

* Is always called without the mmap lock (as is)
* Always returns with the mmap lock unlocked (change)

Sure, we drop+re-acquire the mmap lock a couple of times and lookup the vma, but
isn't that actually being nice to the other parts of the system? In the future
 it would simply get called with the per-vma lock and would return with it unlocked.

In the good old days, looking up VMAs was expensive, but nowadays ... not sure
if it still matters?

khugepaged? Not sure if this matters. madvise? I suspect many real users operate
on a single PMD only (e.g., tcmalloc, jemalloc). For the other ones, not sure if
dropping the lock every PMD is really a problem?

IOW, how bad would the following simplification be (prototype that needs more
work and thought):

From a688929aa137bedfdb6b1420a1a7015e1087664a Mon Sep 17 00:00:00 2001
From: "David Hildenbrand (Arm)" <david@kernel.org>
Date: Thu, 1 Oct 2026 11:22:26 +0200
Subject: [PATCH] tmp

Signed-off-by: David Hildenbrand (Arm) <david@kernel.org>
---
 mm/khugepaged.c | 105 ++++++++++++++++++------------------------------
 1 file changed, 39 insertions(+), 66 deletions(-)

diff --git a/mm/khugepaged.c b/mm/khugepaged.c
index 87bbba6ce59a..3085e8b34ff7 100644
--- a/mm/khugepaged.c
+++ b/mm/khugepaged.c
@@ -2858,43 +2858,42 @@ static enum scan_result collapse_run_pmd(struct mm_struct *mm,
 
 /*
  * Try to collapse a single PMD starting at a PMD aligned addr, and return
- * the results.
+ * the results. mmap_lock must be held for reading on entry and is always
+ * dropped on return.
  */
 static enum scan_result collapse_single_pmd(unsigned long addr,
-		struct vm_area_struct *vma, bool *lock_dropped,
+		struct vm_area_struct *vma,
 		struct collapse_control *cc)
 {
 	struct mm_struct *mm = vma->vm_mm;
 	enum scan_result result;
 
 	result = collapse_scan_pmd(vma, addr, cc);
-	if (result != SCAN_SUCCEED && result != SCAN_PTE_MAPPED_HUGEPAGE)
-		return result;
 
 	/* The collapse takes its own locks, so give this up */
 	mmap_read_unlock(mm);
-	*lock_dropped = true;
 
-	return collapse_run_pmd(mm, addr, result, cc);
+	if (result == SCAN_SUCCEED || result == SCAN_PTE_MAPPED_HUGEPAGE)
+		result = collapse_run_pmd(mm, addr, result, cc);
+
+	return result;
 }
 
-static void collapse_scan_mm_slot(unsigned int progress_max,
-		enum scan_result *result, struct collapse_control *cc)
+static enum scan_result collapse_scan_mm_slot(struct collapse_control *cc)
 	__releases(&khugepaged_mm_lock)
 	__acquires(&khugepaged_mm_lock)
 {
 	struct vma_iterator vmi;
 	struct mm_slot *slot;
 	struct mm_struct *mm;
-	struct vm_area_struct *vma;
+	struct vm_area_struct *vma = NULL;
 	unsigned int progress_prev = cc->progress;
+	enum scan_result result = SCAN_FAIL;
 
 	lockdep_assert_held(&khugepaged_mm_lock);
-	*result = SCAN_FAIL;
 
-	if (khugepaged_scan.mm_slot) {
-		slot = khugepaged_scan.mm_slot;
-	} else {
+	slot = khugepaged_scan.mm_slot;
+	if (!slot) {
 		slot = list_first_entry(&khugepaged_scan.mm_head,
 				     struct mm_slot, mm_node);
 		khugepaged_scan.address = 0;
@@ -2903,17 +2902,12 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 	spin_unlock(&khugepaged_mm_lock);
 
 	mm = slot->mm;
-	/*
-	 * Don't wait for semaphore (to avoid long wait times).  Just move to
-	 * the next mm on the list.
-	 */
-	vma = NULL;
+	/* Don't wait for mmap_lock. Move to the next mm if it is contended. */
 	if (unlikely(!mmap_read_trylock(mm)))
-		goto breakouterloop_mmap_lock;
+		goto out;
 
-	cc->progress++;
 	if (unlikely(collapse_test_exit_or_disable(mm)))
-		goto breakouterloop;
+		goto out_unlock;
 
 	vma_iter_init(&vmi, mm, khugepaged_scan.address);
 	for_each_vma(vmi, vma) {
@@ -2939,39 +2933,22 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 			khugepaged_scan.address = hstart;
 		VM_BUG_ON(khugepaged_scan.address & ~HPAGE_PMD_MASK);
 
-		while (khugepaged_scan.address < hend) {
-			bool lock_dropped = false;
+		if (khugepaged_scan.address >= hend)
+			continue;
 
-			cond_resched();
-			if (unlikely(collapse_test_exit_or_disable(mm)))
-				goto breakouterloop;
-
-			VM_WARN_ON_ONCE(khugepaged_scan.address < hstart ||
-				  khugepaged_scan.address + HPAGE_PMD_SIZE >
-				  hend);
-
-			*result = collapse_single_pmd(khugepaged_scan.address,
-						      vma, &lock_dropped, cc);
-			if (*result == SCAN_SUCCEED)
-				khugepaged_pages_collapsed++;
-			/* move to next address */
-			khugepaged_scan.address += HPAGE_PMD_SIZE;
-			if (lock_dropped)
-				/*
-				 * We released mmap_lock so break loop.  Note
-				 * that we drop mmap_lock before all hugepage
-				 * allocations, so if allocation fails, we are
-				 * guaranteed to break here and report the
-				 * correct result back to caller.
-				 */
-				goto breakouterloop_mmap_lock;
-			if (cc->progress >= progress_max)
-				goto breakouterloop;
-		}
+		result = collapse_single_pmd(khugepaged_scan.address, vma, cc);
+		if (result == SCAN_SUCCEED)
+			khugepaged_pages_collapsed++;
+		/* Move to the next address and look up its VMA again. */
+		khugepaged_scan.address += HPAGE_PMD_SIZE;
+		goto out;
 	}
-breakouterloop:
+out_unlock:
 	mmap_read_unlock(mm); /* exit_mmap will destroy ptes after this */
-breakouterloop_mmap_lock:
+out:
+	/* PMD scans and skipped VMAs account for their own progress. */
+	if (cc->progress == progress_prev)
+		cc->progress++;
 
 	spin_lock(&khugepaged_mm_lock);
 	VM_BUG_ON(khugepaged_scan.mm_slot != slot);
@@ -2998,6 +2975,7 @@ static void collapse_scan_mm_slot(unsigned int progress_max,
 
 	trace_mm_khugepaged_scan(mm, cc->progress - progress_prev,
 				 khugepaged_scan.mm_slot == NULL);
+	return result;
 }
 
 static int khugepaged_has_work(void)
@@ -3034,7 +3012,7 @@ static void khugepaged_do_scan(struct collapse_control *cc)
 			pass_through_head++;
 		if (khugepaged_has_work() &&
 		    pass_through_head < 2)
-			collapse_scan_mm_slot(progress_max, &result, cc);
+			result = collapse_scan_mm_slot(cc);
 		else
 			cc->progress = progress_max;
 		spin_unlock(&khugepaged_mm_lock);
@@ -3232,7 +3210,6 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	unsigned long hstart, hend, addr;
 	enum scan_result last_fail = SCAN_FAIL;
 	int thps = 0;
-	bool mmap_unlocked = false;
 
 	BUG_ON(vma->vm_start > start);
 	BUG_ON(vma->vm_end < end);
@@ -3255,24 +3232,23 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 	lru_add_drain_all();
 
 	for (addr = hstart; addr < hend; addr += HPAGE_PMD_SIZE) {
-		enum scan_result result = SCAN_FAIL;
+		enum scan_result result;
 
-		if (mmap_unlocked) {
+		if (addr != hstart) {
 			cond_resched();
 			mmap_read_lock(mm);
-			mmap_unlocked = false;
-			*lock_dropped = true;
 			result = hugepage_vma_revalidate(mm, addr, false, &vma,
 							 cc, HPAGE_PMD_ORDER);
 			if (result != SCAN_SUCCEED) {
 				last_fail = result;
-				goto out_nolock;
+				goto out;
 			}
 
 			hend = min(hend, vma->vm_end & HPAGE_PMD_MASK);
 		}
 
-		result = collapse_single_pmd(addr, vma, &mmap_unlocked, cc);
+		result = collapse_single_pmd(addr, vma, cc);
+		*lock_dropped = true;
 
 		switch (result) {
 		case SCAN_SUCCEED:
@@ -3295,17 +3271,14 @@ int madvise_collapse(struct vm_area_struct *vma, unsigned long start,
 		default:
 			last_fail = result;
 			/* Other error, exit */
-			goto out_maybelock;
+			goto out_lock;
 		}
 	}
 
-out_maybelock:
+out_lock:
 	/* Caller expects us to hold mmap_lock on return */
-	if (mmap_unlocked) {
-		*lock_dropped = true;
-		mmap_read_lock(mm);
-	}
-out_nolock:
+	mmap_read_lock(mm);
+out:
 	mmap_assert_locked(mm);
 	kfree(cc);
 
-- 
2.43.0


Based on that, I'd rather want to see collapse_single_pmd() to just inline the file
and anon paths, and see how we can further optimize the locking internally (e.g., perform
the pagecache scanning without the mmap lock).

-- 
Cheers,

David