Re: [PATCH] mm/migrate_device: avoid overflowing migrate_vma collection arrays
From: Alistair Popple
Date: Tue Jul 14 2026 - 01:31:13 EST
On 2026-07-14 at 10:10 +1000, Balbir Singh <balbirs@xxxxxxxxxx> wrote...
> On 7/14/26 8:24 AM, Alistair Popple wrote:
> > On 2026-07-14 at 01:39 +1000, Zi Yan <ziy@xxxxxxxxxx> wrote...
> >> On 13 Jul 2026, at 11:25, David Hildenbrand (Arm) wrote:
> >>
> >>> On 7/13/26 17:20, Zi Yan wrote:
> >>>> On 13 Jul 2026, at 8:57, David Hildenbrand (Arm) wrote:
> >>>>
> >>>>> On 7/10/26 17:35, Zi Yan wrote:
> >>>>>>
> >>>>>> Not all the time. After a folio split and PTE lock is dropped, there are
> >>>>>> different cases:
> >>>>>>
> >>>>>> 1. no page table change, the collection grabs the PTE lock and continues;
> >>>>>
> >>>>> Yes.
> >>>>>
> >>>>>>
> >>>>>> 2. some PTE(s) is changed, the same as 1 and the collection continues;
> >>>>>
> >>>>> Yes.
> >>>>>
> >>>>>>
> >>>>>> 3. the PMD pointing to the PTE page table is changed to either a leaf
> >>>>>> PMD or an invalid PMD, the collection cannot grab the PTE lock to
> >>>>>> work on the remaming PTEs, since they are gone. For the collected
> >>>>>> PFNs (they are mapped more than once) and folios with elevated
> >>>>>> refcount (they are mapped once and unmapped here as an optimization),
> >>>>>> they will be processed later in migrate_vma_unmap(). Since
> >>>>>> migrate_vma_collect() is best effort, there is no need to revert and
> >>>>>> try to recollect from the beginning (to get that possible large
> >>>>>> folio or skip).
> >>>>>
> >>>>> Yes.
> >>>>>
> >>>>>>
> >>>>>> BTW, recollection will be more feasible if migrate_vma_collect() does
> >>>>>> not do unmap singly-mapped optimization, since in the case, no PTE is
> >>>>>> changed, we just need to reset migrate->->cpages and migrate->npages and
> >>>>>> restart from the beginning.
> >>>>>>
> >>>>>>
> >>>>>> Right. The PTE page table should be gone and the original PMD, pointing
> >>>>>> to the PTE page table, becomes a leaf PMD or an invalid PMD, since
> >>>>>> someone changed it when the PTE lock is dropped.
> >>>>>
> >>>>> Right, so I guess we'd have to remember how far to rollback, and rollback would
> >>>>> mean resetting migrate->->cpages and migrate->npages.
> >>>>
> >>>> For case 1 and 2, we do not roll back. For case 3, since the PMD is changed,
> >>>> we just roll back to the start address and set migrate->cpages and migrate->npages
> >>>> both to 0.
> >>>
> >>> Can't we walk multiple PMDs and have data already in there?
> >>
> >> migrate_vma_collect() is the pmd_entry of struct mm_walk_ops, so it only
> >> walks a single PMD at a time.
> >>
> >>>
> >>>>>
> >>>>> And for the entries we rollback, we have to decide whether to folio_put() and
> >>>>> whether to folio_unlock() [fault_folio != folio]. That's the confusing bit given
> >>>>> that the code is rather "interesting".
> >>>>
> >>>> IMHO, it might be much cleaner to remove the “unmap the mapping if the folio is
> >>>> only mapped once” optimization, so there is no need to roll back at all.
> >>>
> >>> I guess most folios are only mapped once. Do we have any numbers on that?
> >>
> >> I will defer the question to Balbir, Alistair, and Matthew.
> >
> > Unfortunately previous tests that I've done showed the optimisation is critical
> > to performance in the common case (ie. singly mapped folios).
> >
> > I say unfortunately because I have similar feelings to others here about this
> > code :) I tested this a while ago because I was trying to clean up the collect
> > step by mostly removing it and replacing it with a hmm_range_fault() type thing,
> > but the optimisation proved too important due to unmap having to do another
> > rmap/page table walk.
> >
> > I don't have the numbers on hand atm, but will recreate them just in case my
> > recollection is faulty.
> >
>
> I had a patch to add tracepoints to the code and some of that instrumentation
> included this data. A combination of the trace points + hmm_tests will help us
> get some data to start with.
But hmm-test already has a nice bandwidth benchmark :) We just need to force
the slow path. That used to be easier before I simplified the trylock handling
but with the below patch applied I get quite a lot better perf from the
optimisation. Note the patch isn't entirely correct as it assumes any locked
page was locked by migrate_vma_collect_pmd but it's good enough for the purposes
here.
Optimised results:
HMM THP Migration Benchmark
---------------------------
System page size: 4096 bytes
=== Small Buffer (0.5 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.097 ms | 0.120 ms | 19.5%
Dev->Sys Migration | 0.061 ms | 0.062 ms | 0.4%
S->D Throughput | 5.05 GB/s | 4.07 GB/s | 24.2%
D->S Throughput | 7.96 GB/s | 7.92 GB/s | 0.4%
=== Half THP Size (1.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.217 ms | 0.212 ms | -2.1%
Dev->Sys Migration | 0.136 ms | 0.138 ms | 1.2%
S->D Throughput | 4.51 GB/s | 4.60 GB/s | -2.1%
D->S Throughput | 7.17 GB/s | 7.09 GB/s | 1.2%
=== Single THP Size (2.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.290 ms | 0.506 ms | 42.6%
Dev->Sys Migration | 0.095 ms | 0.283 ms | 66.3%
S->D Throughput | 6.72 GB/s | 3.86 GB/s | 74.3%
D->S Throughput | 20.52 GB/s | 6.91 GB/s | 197.0%
=== Two THP Size (4.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.634 ms | 1.026 ms | 38.2%
Dev->Sys Migration | 0.195 ms | 0.567 ms | 65.6%
S->D Throughput | 6.16 GB/s | 3.81 GB/s | 61.8%
D->S Throughput | 20.05 GB/s | 6.89 GB/s | 191.1%
=== Four THP Size (8.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 1.301 ms | 2.593 ms | 49.8%
Dev->Sys Migration | 0.541 ms | 1.444 ms | 62.5%
S->D Throughput | 6.01 GB/s | 3.01 GB/s | 99.3%
D->S Throughput | 14.45 GB/s | 5.41 GB/s | 167.0%
=== Eight THP Size (16.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 3.853 ms | 6.119 ms | 37.0%
Dev->Sys Migration | 1.731 ms | 3.354 ms | 48.4%
S->D Throughput | 4.06 GB/s | 2.55 GB/s | 58.8%
D->S Throughput | 9.03 GB/s | 4.66 GB/s | 93.8%
=== One twenty eight THP Size (256.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 70.022 ms | 105.921 ms | 33.9%
Dev->Sys Migration | 31.175 ms | 55.393 ms | 43.7%
S->D Throughput | 3.57 GB/s | 2.36 GB/s | 51.3%
D->S Throughput | 8.02 GB/s | 4.51 GB/s | 77.7%
Non-optmisied results (ignore the THP column, I wasn't able to test the slow
path so easily there so just skipped that bit of the test):
=== Small Buffer (0.5 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.000 ms | 0.321 ms | 100.0%
Dev->Sys Migration | 0.000 ms | 0.276 ms | 100.0%
S->D Throughput | 0.00 GB/s | 1.52 GB/s | -100.0%
D->S Throughput | 0.00 GB/s | 1.77 GB/s | -100.0%
=== Half THP Size (1.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.000 ms | 0.660 ms | 100.0%
Dev->Sys Migration | 0.000 ms | 0.571 ms | 100.0%
S->D Throughput | 0.00 GB/s | 1.48 GB/s | -100.0%
D->S Throughput | 0.00 GB/s | 1.71 GB/s | -100.0%
=== Single THP Size (2.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.000 ms | 1.385 ms | 100.0%
Dev->Sys Migration | 0.000 ms | 1.152 ms | 100.0%
S->D Throughput | 0.00 GB/s | 1.41 GB/s | -100.0%
D->S Throughput | 0.00 GB/s | 1.69 GB/s | -100.0%
=== Two THP Size (4.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.000 ms | 2.863 ms | 100.0%
Dev->Sys Migration | 0.000 ms | 2.323 ms | 100.0%
S->D Throughput | 0.00 GB/s | 1.36 GB/s | -100.0%
D->S Throughput | 0.00 GB/s | 1.68 GB/s | -100.0%
=== Four THP Size (8.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.000 ms | 6.199 ms | 100.0%
Dev->Sys Migration | 0.000 ms | 4.904 ms | 100.0%
S->D Throughput | 0.00 GB/s | 1.26 GB/s | -100.0%
D->S Throughput | 0.00 GB/s | 1.59 GB/s | -100.0%
=== Eight THP Size (16.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.000 ms | 14.260 ms | 100.0%
Dev->Sys Migration | 0.000 ms | 10.457 ms | 100.0%
S->D Throughput | 0.00 GB/s | 1.10 GB/s | -100.0%
D->S Throughput | 0.00 GB/s | 1.49 GB/s | -100.0%
=== One twenty eight THP Size (256.0 MB) ===
| With THP | Without THP | Improvement
---------------------------------------------------------------------
Sys->Dev Migration | 0.000 ms | 234.215 ms | 100.0%
Dev->Sys Migration | 0.000 ms | 167.061 ms | 100.0%
S->D Throughput | 0.00 GB/s | 1.07 GB/s | -100.0%
D->S Throughput | 0.00 GB/s | 1.50 GB/s | -100.0%
---
diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 554754eb26ff..9f5563af8cd9 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -397,7 +397,7 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
* optimisation to avoid walking the rmap later with
* try_to_migrate().
*/
- if (fault_folio == folio || folio_trylock(folio)) {
+ if (fault_folio == folio) {
bool anon_exclusive;
pte_t swp_pte;
@@ -466,8 +466,16 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
if (pte_present(pte))
unmapped++;
} else {
- folio_put(folio);
- mpfn = 0;
+ /*
+ * Slow path: couldn't lock the folio. Leave the PTE
+ * intact and collect it anyway; migrate_device_unmap()
+ * will lock it and use try_to_migrate() to install the
+ * migration entry via rmap walk.
+ *
+ * Keep the reference taken above for
+ * migrate_device_unmap() to drop.
+ */
+ migrate->cpages++;
}
next:
@@ -605,6 +613,16 @@ static unsigned long migrate_device_unmap(unsigned long *src_pfns,
folio_put(folio);
}
+ /*
+ * Fast-path pages (folio_trylock succeeded in collect) are
+ * already locked. Slow-path pages were not locked during
+ * collect; lock them now so try_to_migrate() and the finalize
+ * path can rely on the folio being locked. This applies to
+ * both regular and zone-device folios.
+ */
+ if (!folio_test_locked(folio))
+ folio_lock(folio);
+
if (folio_mapped(folio))
try_to_migrate(folio, 0);
diff --git a/tools/testing/selftests/mm/hmm-tests.c b/tools/testing/selftests/mm/hmm-tests.c
index e4c49699f3f7..3a673181b9ad 100644
--- a/tools/testing/selftests/mm/hmm-tests.c
+++ b/tools/testing/selftests/mm/hmm-tests.c
@@ -2936,8 +2936,8 @@ TEST_F_TIMEOUT(hmm, benchmark_thp_migration, 120)
break;
/* Test with THP */
- ASSERT_EQ(run_migration_benchmark(self->fd, 1, test_sizes[i],
- iterations, &thp_results), 0);
+ // ASSERT_EQ(run_migration_benchmark(self->fd, 1, test_sizes[i],
+ // iterations, &thp_results), 0);
/* Test without THP */
ASSERT_EQ(run_migration_benchmark(self->fd, 0, test_sizes[i],