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

["Zi Yan" <ziy@nvidia.com>]

On Tue Jul 14, 2026 at 1:30 AM EDT, Alistair Popple wrote:
> 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.

Thanks for the numbers. If I read them correctly, the optimization
has about 3x throughput of unoptimzied version. The optimization has to
stay.

I will try to come up with a fix that resets collected PFNs, reverts
established migration entry along with putting folio ref, and restarts
from the beginning.

>
> 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],




-- 
Best Regards,
Yan, Zi