Re: [PATCH] mm/migrate_device: avoid overflowing migrate_vma collection arrays

[Alistair Popple <apopple@nvidia.com>]

On 2026-07-14 at 10:10 +1000, Balbir Singh <balbirs@nvidia.com> wrote...
> On 7/14/26 8:24 AM, Alistair Popple wrote:
> > On 2026-07-14 at 01:39 +1000, Zi Yan <ziy@nvidia.com> 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],