[PATCH v2 0/4] mm/ksm: use per-VMA locking to reduce mmap_lock contention

[PATCH v2 0/4] mm/ksm: use per-VMA locking to reduce mmap_lock contention

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

KSM scans VM_MERGEABLE VMAs and currently protects each scan with
mmap_read_lock(). Per-VMA locking allows KSM to read-lock only the VMA
it is actually interested in, so that unrelated mmap()/munmap() activity
in the same mm no longer blocks ksmd.

This series is organized as follows:

  Patch 1 removes an unused 'vma' member from struct folio_walk.  It has
  never been used since its introduction and is pure cleanup.

  Patch 2 adds a 'walk_lock' member to struct folio_walk and extends
  folio_walk_start() to assert the required locking mode.  Existing
  callers are converted to pass PGWALK_RDLOCK, so there is no functional
  change.  This prepares folio_walk_start() for callers that hold a
  per-VMA read lock instead of mmap_read_lock(), which is needed by the
  Patch 4. No functional change.

  Patch 3 replaces the boolean 'lock_vma' argument of break_ksm() with a
  KSM-local enum ksm_walk_lock and builds the mm_walk_ops on the stack,
  so the function-pointer table stays const and there is no shared
  mutable state.  No functional change.

  Patch 4 introduces find_mergeable_vma_locked(), which uses the
  universal per-VMA locking helper vma_start_read_unlocked() to look up
  and read-lock a VM_MERGEABLE VMA without taking mmap_read_lock().  All
  KSM call sites that previously used find_mergeable_vma() under
  mmap_read_lock() are converted to the new helper, and the locking in
  get_mergeable_page() is switched to PGWALK_VMA_RDLOCK_VERIFY so that
  folio_walk_start() can verify the per-VMA lock is held.  Because ksmd
  no longer holds mmap_read_lock() while walking, __ksm_exit() must also
  write-lock every VMA (in addition to mmap_write_lock()) to keep
  serializing against an in-flight ksmd walk before the page tables are
  freed.

A microbenchmark was run to measure the time KSM takes to merge a
victim region under mmap_lock contention. Under interference from 4 churner
threads, the merge time of the per-VMA KSM-optimized kernel is
significantly reduced by 50%.

Changelog:
  v2:
  - Patch 1: rewrote the commit message to explain the introducing
    commit, how the VMA is actually used, and the audit performed;
    dropped the Fixes: tag.

  - Patch 3: 
    * replaced the boolean 'lock_vma' with a KSM-local enum
    ksm_walk_lock instead of exposing the internal enum page_walk_lock
    as a parameter;
    * build the mm_walk_ops on the stack rather than
    mutating a global struct, so the function-pointer table stays const
    and there is no shared mutable state; rewrote the commit message.

  - Patch 4:
    * Fixed a per-VMA read-lock leak in find_mergeable_vma_locked()
    when the VMA is not mergeable;
    made __ksm_exit() write-lock every VMA (vma_start_write()) in
    addition to mmap_write_lock(), so that it still serializes against a
    ksmd walk that holds only a per-VMA read lock.

Xu Xin (4):
  mm/pagewalk: delete the unused member
  mm: make folio_walk_start()'s locking asserts scalable
  mm/ksm: make break_ksm() more scalable
  ksm: add find_mergeable_vma_locked() to use per-VMA locking

 arch/s390/mm/fault.c     |   4 +-
 include/linux/pagewalk.h |   2 +-
 kernel/events/uprobes.c  |   4 +-
 mm/huge_memory.c         |   4 +-
 mm/ksm.c                 | 124 +++++++++++++++++++++++++++------------
 mm/migrate.c             |   8 ++-
 mm/pagewalk.c            |  10 +++-
 mm/rmap.c                |   4 +-
 8 files changed, 116 insertions(+), 44 deletions(-)

--
2.34.1

[PATCH v2 1/4] mm/pagewalk: delete the unused member

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

The 'vma' member of struct folio_walk has never been used since it was
introduced by commit aa39ca6940f1a ("mm/pagewalk: introduce
folio_walk_start() + folio_walk_end()").

The VMA is always passed to folio_walk_start() as a separate argument,
and neither folio_walk_start() nor folio_walk_end() ever reads or writes
fw->vma. I audited all the current users of struct folio_walk (the
callers of folio_walk_start()) and none of them assign to or read the
member, so it is safe to remove.

Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 include/linux/pagewalk.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/include/linux/pagewalk.h b/include/linux/pagewalk.h
index b41d7265c01b..1c397be0d092 100644
--- a/include/linux/pagewalk.h
+++ b/include/linux/pagewalk.h
@@ -183,7 +183,6 @@ struct folio_walk {
 		pmd_t pmd;
 	};
 	/* private */
-	struct vm_area_struct *vma;
 	spinlock_t *ptl;
 };
 
-- 
2.34.1

[PATCH v2 2/4] mm: make folio_walk_start()'s locking asserts scalable

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

Add an additional member 'walk_lock' to folio_walk to indicate
locking requirements for the walk. Similar to commit 49b0638502da0
("mm: enable page walking API to lock vmas during the walk").
But no change is made on any existing locking behavior,
all existing folio_walk_start() callers are currently still
under mmap_read_lock() protection.

This change is prepared for the later patch to enable VMA locking
asserts. folio_walk_start now operate under write-locked mmap_lock.
With introduction of vma locks at the next patch, the vmas have to
be locked as well during such walks to prevent concurrent page
faults in these areas.

No functional change intended.

Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 arch/s390/mm/fault.c     |  4 +++-
 include/linux/pagewalk.h |  1 +
 kernel/events/uprobes.c  |  4 +++-
 mm/huge_memory.c         |  4 +++-
 mm/ksm.c                 |  4 +++-
 mm/migrate.c             |  8 ++++++--
 mm/pagewalk.c            | 10 +++++++++-
 mm/rmap.c                |  4 +++-
 8 files changed, 31 insertions(+), 8 deletions(-)

diff --git a/arch/s390/mm/fault.c b/arch/s390/mm/fault.c
index 46d828926009..968f128a5232 100644
--- a/arch/s390/mm/fault.c
+++ b/arch/s390/mm/fault.c
@@ -412,7 +412,9 @@ __context_unsafe(/* folio_walk_end() not instrumented */)
 	unsigned long addr = get_fault_address(regs);
 	struct mm_struct *mm = current->mm;
 	struct vm_area_struct *vma;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct folio *folio;
 	int rc;
 
diff --git a/include/linux/pagewalk.h b/include/linux/pagewalk.h
index 1c397be0d092..455428eaf1fc 100644
--- a/include/linux/pagewalk.h
+++ b/include/linux/pagewalk.h
@@ -184,6 +184,7 @@ struct folio_walk {
 	};
 	/* private */
 	spinlock_t *ptl;
+	enum page_walk_lock walk_lock;
 };
 
 struct folio *folio_walk_start(struct folio_walk *fw,
diff --git a/kernel/events/uprobes.c b/kernel/events/uprobes.c
index caf444762a27..3e3927109ee9 100644
--- a/kernel/events/uprobes.c
+++ b/kernel/events/uprobes.c
@@ -507,7 +507,9 @@ int uprobe_write(struct arch_uprobe *auprobe, struct vm_area_struct *vma,
 	int ret, ref_ctr_updated = 0;
 	unsigned int gup_flags = FOLL_FORCE;
 	struct mmu_notifier_range range;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct folio *folio;
 	struct page *page;
 
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index bd44d58e8684..62ec81396c9d 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -4816,7 +4816,9 @@ static int split_huge_pages_pid(int pid, unsigned long vaddr_start,
 	 */
 	for (addr = vaddr_start; addr < vaddr_end; addr += PAGE_SIZE) {
 		struct vm_area_struct *vma = vma_lookup(mm, addr);
-		struct folio_walk fw;
+		struct folio_walk fw = {
+			.walk_lock  = PGWALK_RDLOCK,
+		};
 		struct folio *folio;
 		struct address_space *mapping;
 		unsigned int target_order = new_order;
diff --git a/mm/ksm.c b/mm/ksm.c
index d6f736a1aa2f..d4fc0b0965aa 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -819,8 +819,10 @@ static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
 	unsigned long addr = rmap_item->address;
 	struct vm_area_struct *vma;
 	struct page *page = NULL;
-	struct folio_walk fw;
 	struct folio *folio;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 
 	mmap_read_lock(mm);
 	vma = find_mergeable_vma(mm, addr);
diff --git a/mm/migrate.c b/mm/migrate.c
index 875079464930..89662d5bedf4 100644
--- a/mm/migrate.c
+++ b/mm/migrate.c
@@ -2301,7 +2301,9 @@ static int add_folio_for_migration(struct mm_struct *mm, const void __user *p,
 		int node, struct list_head *pagelist, bool migrate_all)
 {
 	struct vm_area_struct *vma;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct folio *folio;
 	unsigned long addr;
 	int err = -EFAULT;
@@ -2463,7 +2465,9 @@ static void do_pages_stat_array(struct mm_struct *mm, unsigned long nr_pages,
 	for (i = 0; i < nr_pages; i++) {
 		unsigned long addr = (unsigned long)untagged_addr_remote(mm, *pages);
 		struct vm_area_struct *vma;
-		struct folio_walk fw;
+		struct folio_walk fw = {
+			.walk_lock  = PGWALK_RDLOCK,
+		};
 		struct folio *folio;
 		int err = -EFAULT;
 
diff --git a/mm/pagewalk.c b/mm/pagewalk.c
index e6493bbe6919..7d85289fcb87 100644
--- a/mm/pagewalk.c
+++ b/mm/pagewalk.c
@@ -910,7 +910,15 @@ struct folio *folio_walk_start(struct folio_walk *fw,
 	pgd_t *pgdp;
 	p4d_t *p4dp;
 
-	mmap_assert_locked(vma->vm_mm);
+	/*
+	 * Other locking modes except for mmap or vma read locking are not
+	 * expected.
+	 */
+	if (fw->walk_lock != PGWALK_RDLOCK && fw->walk_lock != PGWALK_VMA_RDLOCK_VERIFY)
+		WARN_ONCE(1, "walk_lock is not expected!\n");
+	process_mm_walk_lock(vma->vm_mm, fw->walk_lock);
+	process_vma_walk_lock(vma, fw->walk_lock);
+
 	vma_pgtable_walk_begin(vma);
 
 	if (WARN_ON_ONCE(addr < vma->vm_start || addr >= vma->vm_end))
diff --git a/mm/rmap.c b/mm/rmap.c
index 1131b76bbbc2..6e3bc2acd1d7 100644
--- a/mm/rmap.c
+++ b/mm/rmap.c
@@ -2862,7 +2862,9 @@ struct page *make_device_exclusive(struct mm_struct *mm, unsigned long addr,
 	struct mmu_notifier_range range;
 	struct folio *folio, *fw_folio;
 	struct vm_area_struct *vma;
-	struct folio_walk fw;
+	struct folio_walk fw = {
+		.walk_lock  = PGWALK_RDLOCK,
+	};
 	struct page *page;
 	swp_entry_t entry;
 	pte_t swp_pte;
-- 
2.34.1

[PATCH v2 3/4] mm/ksm: make break_ksm() more scalable

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

break_ksm() currently takes a boolean 'lock_vma' to select between two
static mm_walk_ops tables.  This is fragile: the ops struct is full of
function pointers and making it mutable (to switch walk_lock at runtime)
turns it into a juicy target for attackers.

Replace the boolean with a KSM-local enum ksm_walk_lock and build the
mm_walk_ops on the stack, so the table stays const and there is no
shared mutable state.  This also prepares for the later patch that
switches break_ksm() to per-VMA read-locking, where a third locking
mode will be needed.

No functional change intended.

Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 mm/ksm.c | 42 +++++++++++++++++++++++++++---------------
 1 file changed, 27 insertions(+), 15 deletions(-)

diff --git a/mm/ksm.c b/mm/ksm.c
index d4fc0b0965aa..84d363dabdc8 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -661,15 +661,25 @@ static int break_ksm_pmd_entry(pmd_t *pmdp, unsigned long addr, unsigned long en
 	return found;
 }
 
-static const struct mm_walk_ops break_ksm_ops = {
-	.pmd_entry = break_ksm_pmd_entry,
-	.walk_lock = PGWALK_RDLOCK,
+/*
+ * Locking contexts in which break_ksm() may walk the page tables.
+ */
+enum ksm_walk_lock {
+	/* walk under mmap_read_lock() */
+	KSM_WALK_MMAP_RDLOCK,
+	/* walk under mmap_write_lock() */
+	KSM_WALK_MMAP_WRLOCK,
 };
 
-static const struct mm_walk_ops break_ksm_lock_vma_ops = {
-	.pmd_entry = break_ksm_pmd_entry,
-	.walk_lock = PGWALK_WRLOCK,
-};
+static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
+{
+	switch (lock) {
+	case KSM_WALK_MMAP_WRLOCK:
+		return PGWALK_WRLOCK;
+	default:
+		return PGWALK_RDLOCK;
+	}
+}
 
 /*
  * Though it's very tempting to unmerge rmap_items from stable tree rather
@@ -697,17 +707,19 @@ static const struct mm_walk_ops break_ksm_lock_vma_ops = {
  * protection keys here anyway.
  */
 static int break_ksm(struct vm_area_struct *vma, unsigned long addr,
-		unsigned long end, bool lock_vma)
+		unsigned long end, enum ksm_walk_lock walk_lock)
 {
 	vm_fault_t ret = 0;
-	const struct mm_walk_ops *ops = lock_vma ?
-				&break_ksm_lock_vma_ops : &break_ksm_ops;
+	struct mm_walk_ops ops = {
+		.pmd_entry = break_ksm_pmd_entry,
+		.walk_lock = ksm_walk_lock_to_pgwalk(walk_lock),
+	};
 
 	do {
 		int ksm_page;
 
 		cond_resched();
-		ksm_page = walk_page_range_vma(vma, addr, end, ops, &addr);
+		ksm_page = walk_page_range_vma(vma, addr, end, &ops, &addr);
 		if (ksm_page <= 0)
 			return ksm_page;
 		ret = handle_mm_fault(vma, addr,
@@ -809,7 +821,7 @@ static void break_cow(struct ksm_rmap_item *rmap_item)
 	mmap_read_lock(mm);
 	vma = find_mergeable_vma(mm, addr);
 	if (vma)
-		break_ksm(vma, addr, addr + PAGE_SIZE, false);
+		break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_MMAP_RDLOCK);
 	mmap_read_unlock(mm);
 }
 
@@ -1247,7 +1259,7 @@ static int unmerge_and_remove_all_rmap_items(void)
 		for_each_vma(vmi, vma) {
 			if (!(vma->vm_flags & VM_MERGEABLE) || !vma_has_anon_rmap(vma))
 				continue;
-			err = break_ksm(vma, vma->vm_start, vma->vm_end, false);
+			err = break_ksm(vma, vma->vm_start, vma->vm_end, KSM_WALK_MMAP_RDLOCK);
 			if (err)
 				goto error;
 		}
@@ -2896,7 +2908,7 @@ static int __ksm_del_vma(struct vm_area_struct *vma)
 		return 0;
 
 	if (vma_has_anon_rmap(vma)) {
-		err = break_ksm(vma, vma->vm_start, vma->vm_end, true);
+		err = break_ksm(vma, vma->vm_start, vma->vm_end, KSM_WALK_MMAP_WRLOCK);
 		if (err)
 			return err;
 	}
@@ -3048,7 +3060,7 @@ int ksm_madvise(struct vm_area_struct *vma, unsigned long start,
 			return 0;		/* just ignore the advice */
 
 		if (vma_has_anon_rmap(vma)) {
-			err = break_ksm(vma, start, end, true);
+			err = break_ksm(vma, start, end, KSM_WALK_MMAP_WRLOCK);
 			if (err)
 				return err;
 		}
-- 
2.34.1

[PATCH v2 4/4] ksm: add find_mergeable_vma_locked() to use per-VMA locking

[xu.xin]

From: Xu Xin <xu.xin@linux.dev>

Longlong Xia reported that get_mergeable_page() takes mmap_read_lock()
on every merge, which becomes a contention point under load [1]. David
Hildenbrand suggested switching to per-VMA locking to avoid the
process-wide mmap lock [2].

Add find_mergeable_vma_locked(), which is similar to find_tcp_vma(),
using the universal per-VMA locking helper vma_start_read_unlocked() to
look up and read-lock a VM_MERGEABLE VMA without taking
mmap_read_lock().  All KSM call sites that previously used
find_mergeable_vma() under mmap_read_lock() are converted to the new
helper, and the locking in get_mergeable_page() is switched to
PGWALK_VMA_RDLOCK_VERIFY so that folio_walk_start() can verify the
per-VMA lock is held.

Switching to per-VMA locking also changes how ksmd is synchronized with
a process that is exiting.  __ksm_exit() previously relied on a plain
mmap_write_lock() to serialize against ksmd, but that no longer blocks
a ksmd walk that holds only a per-VMA read lock.  So __ksm_exit() now
also write-locks every VMA (vma_start_write()) while holding
mmap_write_lock(): this blocks until any in-flight ksmd walk drops its
per-VMA read lock, and prevents new ones from starting, since
vma_start_read_unlocked() backs out when a VMA is write-locked.  The
per-VMA write locks are released by mmap_write_unlock() via
vma_end_write_all().  This is safe because __mmput() runs ksm_exit()
after mm_users has gone to zero and before exit_mmap() frees the page
tables, so any new ksmd walk backs out via ksm_test_exit().

Performance
===========
As a test, I construct a testcase which follows the approach: Create
one victim and several churner threads sharing one mm_struct [3]; The
victim registers a 32 MiB anonymous VM_MERGEABLE region containing
8192 pages totally:churners hammer mmap_lock via mmap/munmap
repeatedly; ksmd merges the victim's pages; Compare latency baseline
VS this per-VMA patch.

                 Before patched         After Patched
0 churner:     1.627 seconds          1.426 seconds
4 churners:    72.45 seconds          36.61 seconds

In conclusion, when no mmap_lock contention (0 churner), there is little
difference between the baseline kernel and the per-VMA optimized kernel;
But under interference from 4 churner threads, the merge time of the
per-VMA KSM-optimized kernel is significantly reduced by 50%.

[1] https://lore.kernel.org/all/20260909023507.1408925-1-xialonglong2025@163.com/
[2] https://lore.kernel.org/all/91cf8129-d705-4ebd-b420-129bab7ed261@kernel.org/
[3] https://lore.kernel.org/all/20260911162242418ODBbr9wspxxIXDS8uvpeZ@zte.com.cn/

Reported-by: Longlong Xia <xialonglong@kylinos.cn>
Signed-off-by: Xu Xin <xu.xin@linux.dev>
---
 mm/ksm.c | 82 +++++++++++++++++++++++++++++++++++++++++---------------
 1 file changed, 60 insertions(+), 22 deletions(-)

diff --git a/mm/ksm.c b/mm/ksm.c
index 84d363dabdc8..64f9e803b99e 100644
--- a/mm/ksm.c
+++ b/mm/ksm.c
@@ -669,6 +669,8 @@ enum ksm_walk_lock {
 	KSM_WALK_MMAP_RDLOCK,
 	/* walk under mmap_write_lock() */
 	KSM_WALK_MMAP_WRLOCK,
+	/* walk with the vma already read-locked (per-VMA locking) */
+	KSM_WALK_VMA_RDLOCK,
 };
 
 static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
@@ -676,6 +678,8 @@ static enum page_walk_lock ksm_walk_lock_to_pgwalk(enum ksm_walk_lock lock)
 	switch (lock) {
 	case KSM_WALK_MMAP_WRLOCK:
 		return PGWALK_WRLOCK;
+	case KSM_WALK_VMA_RDLOCK:
+		return PGWALK_VMA_RDLOCK_VERIFY;
 	default:
 		return PGWALK_RDLOCK;
 	}
@@ -784,15 +788,36 @@ static bool vma_ksm_compatible(struct vm_area_struct *vma)
 	return ksm_compatible(vma->vm_file, vma->flags);
 }
 
-static struct vm_area_struct *find_mergeable_vma(struct mm_struct *mm,
-		unsigned long addr)
+/**
+ * find_mergeable_vma_locked() - Find the VMA covering 'address' which is
+ * VM_MERGEABLE and read-lock it by per-VMA locks. Please use vma_end_read()
+ * to unlock vma after finishing reading the VMA (non-NULL).
+ *
+ * Return: If a VMA exists which spans @address, return that VMA, read-locked.
+ * If no VMA is mapped there or, very unlikely, a reference count overflow
+ * occurred, return NULL, and no read-locked.
+ *
+ * IMPORTANT: If a VMA exists but is not VM_MERGEABLE or has no anon_vma,
+ * this function releases the per-VMA read lock before returning NULL.
+ * Callers must NOT call vma_end_read() on a NULL return value.
+ */
+static struct vm_area_struct *find_mergeable_vma_locked(struct mm_struct *mm,
+		unsigned long address)
 {
 	struct vm_area_struct *vma;
+
 	if (ksm_test_exit(mm))
 		return NULL;
-	vma = vma_lookup(mm, addr);
-	if (!vma || !(vma->vm_flags & VM_MERGEABLE) || !vma_has_anon_rmap(vma))
+
+	vma = vma_start_read_unlocked(mm, address);
+	if (!vma)
 		return NULL;
+
+	if (!(vma->vm_flags & VM_MERGEABLE) || !vma->anon_vma) {
+		vma_end_read(vma);
+		return NULL;
+	}
+
 	return vma;
 }
 
@@ -818,11 +843,12 @@ static void break_cow(struct ksm_rmap_item *rmap_item)
 	 */
 	rmap_item->linear_page_index = 0;
 
-	mmap_read_lock(mm);
-	vma = find_mergeable_vma(mm, addr);
-	if (vma)
-		break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_MMAP_RDLOCK);
-	mmap_read_unlock(mm);
+	vma = find_mergeable_vma_locked(mm, addr);
+	if (!vma)
+		return;
+
+	break_ksm(vma, addr, addr + PAGE_SIZE, KSM_WALK_VMA_RDLOCK);
+	vma_end_read(vma);
 }
 
 static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
@@ -833,13 +859,12 @@ static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
 	struct page *page = NULL;
 	struct folio *folio;
 	struct folio_walk fw = {
-		.walk_lock  = PGWALK_RDLOCK,
+		.walk_lock  = PGWALK_VMA_RDLOCK_VERIFY,
 	};
 
-	mmap_read_lock(mm);
-	vma = find_mergeable_vma(mm, addr);
+	vma = find_mergeable_vma_locked(mm, addr);
 	if (!vma)
-		goto out;
+		return NULL;
 
 	folio = folio_walk_start(&fw, vma, addr, 0);
 	if (folio) {
@@ -850,12 +875,12 @@ static struct page *get_mergeable_page(struct ksm_rmap_item *rmap_item)
 		}
 		folio_walk_end(&fw, vma);
 	}
-out:
+
 	if (page) {
 		flush_anon_page(vma, page, addr);
 		flush_dcache_page(page);
 	}
-	mmap_read_unlock(mm);
+	vma_end_read(vma);
 	return page;
 }
 
@@ -1596,14 +1621,14 @@ static int try_to_merge_with_zero_page(struct ksm_rmap_item *rmap_item,
 	if (ksm_use_zero_pages && (rmap_item->oldchecksum == zero_checksum)) {
 		struct vm_area_struct *vma;
 
-		mmap_read_lock(mm);
-		vma = find_mergeable_vma(mm, rmap_item->address);
+		vma = find_mergeable_vma_locked(mm, rmap_item->address);
 		if (vma) {
 			err = try_to_merge_one_page(vma, page,
 					ZERO_PAGE(rmap_item->address));
 			trace_ksm_merge_one_page(
 				page_to_pfn(ZERO_PAGE(rmap_item->address)),
 				rmap_item, mm, err);
+			vma_end_read(vma);
 		} else {
 			/*
 			 * If the vma is out of date, we do not need to
@@ -1611,7 +1636,6 @@ static int try_to_merge_with_zero_page(struct ksm_rmap_item *rmap_item,
 			 */
 			err = 0;
 		}
-		mmap_read_unlock(mm);
 	}
 
 	return err;
@@ -1630,10 +1654,9 @@ static int try_to_merge_with_ksm_page(struct ksm_rmap_item *rmap_item,
 	struct vm_area_struct *vma;
 	int err = -EFAULT;
 
-	mmap_read_lock(mm);
-	vma = find_mergeable_vma(mm, rmap_item->address);
+	vma = find_mergeable_vma_locked(mm, rmap_item->address);
 	if (!vma)
-		goto out;
+		goto out_trace;
 
 	err = try_to_merge_one_page(vma, page, kpage);
 	if (err)
@@ -1653,7 +1676,8 @@ static int try_to_merge_with_ksm_page(struct ksm_rmap_item *rmap_item,
 	rmap_item->linear_page_index = linear_anon_page_index(vma, rmap_item->address);
 	get_anon_vma(vma->anon_vma);
 out:
-	mmap_read_unlock(mm);
+	vma_end_read(vma);
+out_trace:
 	trace_ksm_merge_with_ksm_page(kpage, page_to_pfn(kpage ? kpage : page),
 				rmap_item, mm, err);
 	return err;
@@ -3119,6 +3143,7 @@ void __ksm_exit(struct mm_struct *mm)
 {
 	struct ksm_mm_slot *mm_slot = NULL;
 	struct mm_slot *slot;
+	struct vm_area_struct *vma;
 	int easy_to_free = 0;
 
 	/*
@@ -3153,7 +3178,20 @@ void __ksm_exit(struct mm_struct *mm)
 		mm_flags_clear(MMF_VM_MERGEABLE, mm);
 		mmdrop(mm);
 	} else if (mm_slot) {
+		/*
+		 * ksmd may be walking this mm's page tables under a per-VMA
+		 * read lock (see find_mergeable_vma_locked()).  A plain
+		 * mmap_write_lock() does not exclude such readers, so also
+		 * write-lock every VMA: this blocks until any in-flight ksmd
+		 * walk drops its per-VMA read lock, and prevents new ones
+		 * from starting (vma_start_read_unlocked() backs out when a
+		 * VMA is write-locked).  The per-VMA write locks are released
+		 * by mmap_write_unlock() via vma_end_write_all().
+		 */
 		mmap_write_lock(mm);
+		VMA_ITERATOR(vmi, mm, 0);
+		for_each_vma(vmi, vma)
+			vma_start_write(vma);
 		mmap_write_unlock(mm);
 	}
 
-- 
2.34.1