[PATCH v15 00/11] migrate on fault for device pages

[mpenttil@redhat.com]

From: Mika Penttilä <mpenttil@redhat.com>

A quick respin to address Sashiko's concerns..

Currently, the way device page faulting and migration works
is not optimal, if you want to do both fault handling and
migration at once.

Being able to migrate not present pages (or pages mapped with incorrect
permissions, eg. COW) to the GPU requires doing either of the
following sequences:

1. hmm_range_fault() - fault in non-present pages with correct permissions, etc.
2. migrate_vma_*() - migrate the pages

Or:

1. migrate_vma_*() - migrate present pages
2. If non-present pages detected by migrate_vma_*():
   a) call hmm_range_fault() to fault pages in
   b) call migrate_vma_*() again to migrate now present pages

The problem with the first sequence is that you always have to do two
page walks even when most of the time the pages are present or zero page
mappings so the common case takes a performance hit.

The second sequence is better for the common case, but far worse if
pages aren't present because now you have to walk the page tables three
times (once to find the page is not present, once so hmm_range_fault()
can find a non-present page to fault in and once again to setup the
migration). It is also tricky to code correctly. One page table walk
could costs over 1000 cpu cycles on X86-64, which is a significant hit.

We should be able to walk the page table once, faulting
pages in as required and replacing them with migration entries if
requested.

Add a new flag to HMM APIs, HMM_PFN_REQ_MIGRATE,
which tells to prepare for migration also during fault handling.
For the migrate_vma_setup() call paths, new flags, MIGRATE_VMA_FAULT,
and MIGRATE_VMA_WRITE are added to tell to add fault handling to migrate.

An extra benefit of migrating with hmm_range_fault() path
is the migrate_vma.vma gets populated, so no need to
retrieve that separataly.

This series avoids the problems in current collecting implementation,
where racing MADV_DONTNEED could make the collecting arrays overflow.
This is achieved mainly by having a rollback mechanism and
populating the collecting array as the last step.

Tested in X86-64 VM with HMM test device, passing the selftests.
For performance, the migrate throughput tests from the selftests
show similar numbers (within error margin) as unmodified kernel.
Tested also rebased on the
"Remove device private pages from physical address space" series:
https://lore.kernel.org/linux-mm/20260130111050.53670-1-jniethe@nvidia.com/
plus a small patch to adjust with no problems.

Changes in v15:
  - addressed Sashiko's concerns
  - collected Balbir's ack for patch 01
  
Sashiko's concerns :
https://sashiko.dev/#/patchset/20260922053421.4092027-1-mpenttil%40redhat.com
Concerns and responses, patch numbering refers to v14 series:

patch 01 mm/Kconfig: changes for migrate on fault for device pages
  Q "Does this regression violate Kconfig dependency rules by selecting HMM_MIRROR
  without inheriting its dependencies?"
  A DEVICE_MIGRATION is def_bool MIGRATION && ZONE_DEVICE, and MIGRATION depends on MMU,
  which HMM_MIRROR also depends on. DEVICE_MIGRATION is not user selectable, and nothing
  selects it explicitly now. If some config would select it in the future, it should have MMU.

patch 03 mm/hmm: preparations for HMM to participate in migration
 Q "Is the tracked start address unconditionally overwritten on every walk restart?"
 A yes, fixed
 Q "Could using the owner field as an initialization flag be problematic here?"
 A yes fixed,  range->dev_private_owner not used as sentinel anymore
 Q "How does this API change affect hmm_range_fault_unlocked_timeout()?"
 A migrate on fault and migrate_vma call hmm_range_fault()

patch 04 mm/hmm: do the plumbing for HMM to participate in migration
  Q "Does this mmu_notifier_invalidate_range_start() become unbalanced if the
  fault path is interrupted?"
  A HMM_FAULT_UNLOCKED can not happen with hmm_range_fault() + migration, however
  made the exit path flow thru hmm_vma_post_range_fault() in these cases also
  even though it is nop.
  Q "Can this early return bypass the TLB flush for pages unmapped in previous
  iterations?"
  A yes, fixed

patch 05 mm/hmm: implement folio split for migrate needs in HMM pagewalk
  Q "Does this create a livelock under contention?"
  Q "Will this spuriously fail migration if the folio is split concurrently?"
  Q "Does this correctly handle cases where fault_page was a tail page?"
  A this helper is removed in v15

patch 06 mm/hmm: migrate collection in HMM pagewalk - pte level
  Q "Does taking the reference via folio_get() before blocking on folio_lock()
  create a livelock when multiple threads fault on the same large folio?"
  A The ref has to be taken before dropping ptl lock. That mirrors current
  implementation. That can be improved on with later patch if wanted
  Q "Can this split_folio() crash if a driver mapped a raw compound page into the
  VMA using vm_insert_page()?"
  A The  folio is either from vm_normal_page(), or a device private/coherent page
  Q "Does this error path unconditionally unlock a folio that the caller might hold
  the lock for?"
  A yes, fixed

patch 07 mm/hmm: migrate collection in HMM pagewalk - pmd level
  Q "Does this code handle non-present PMD leaf entries correctly?"
  Q "Is this branch actually reachable for device-private PMDs?"
  A  PMD size device pages are not present swap like entries

  Q "Will this use the correct subpage for folios larger than 2MB?"
  A only PMD sized THP pages are supported

patch 08 mm/hmm: add lazy MMU mode support for migration in HMM pagewalk
  Q Does this skip a necessary TLB flush on the error path during migration?
  A yes, fixed, same concern as for patch 04

patch 10 mm: enable device page migration from HMM pagewalk
  Q "Does this code cause a permanent MMU notifier imbalance when args->pgmap_owner
  is NULL?"
  A yes, this is same concern as for patch 03. fixed

patch 11 lib/test_hmm: add a new testcase for the migrate on fault
  Q "Does skipping the page table update here leave us vulnerable to concurrently
  unmaps during migration?"
  Q "Can this sequence lead to a use-after-free of device pages if a concurrent
  unmap occurs after we dropped the mutex in dmirror_range_fault()?"
  A yes, the sequence is fixed in v15

Link to v14: https://lore.kernel.org/linux-mm/20260922053421.4092027-1-mpenttil@redhat.com/

Cc: David Hildenbrand <david@kernel.org>
Cc: Jason Gunthorpe <jgg@nvidia.com>
Cc: Leon Romanovsky <leonro@nvidia.com>
Cc: Alistair Popple <apopple@nvidia.com>
Cc: Balbir Singh <balbirs@nvidia.com>
Cc: Zi Yan <ziy@nvidia.com>
Cc: Matthew Brost <matthew.brost@intel.com>
Cc: Andrew Morton <akpm@linux-foundation.org>
Cc: Lorenzo Stoakes <ljs@kernel.org>
Cc: "Liam R. Howlett" <Liam.Howlett@oracle.com>
Cc: Vlastimil Babka <vbabka@suse.cz>
Cc: Mike Rapoport <rppt@kernel.org>
Cc: Suren Baghdasaryan <surenb@google.com>
Cc: Michal Hocko <mhocko@suse.com>

Mika Penttilä (11):
  mm/Kconfig: changes for migrate on fault for device pages
  mm: add helper to convert HMM pfn to migrate pfn
  mm/hmm: preparations for HMM to participate in migration
  mm/hmm: do the plumbing for HMM to participate in migration
  mm/hmm: migrate collection in HMM pagewalk - pte level
  mm/hmm: migrate collection in HMM pagewalk - pmd level
  mm/hmm: add lazy MMU mode support for migration in HMM pagewalk
  mm/hmm: implement rollback for device page migration in HMM pagewalk
  mm: enable device page migration from HMM pagewalk
  lib/test_hmm: add a new testcase for the migrate on fault
  Documentation/mm/hmm: document migration through hmm_range_fault()

 Documentation/mm/hmm.rst               |  39 +
 include/linux/hmm.h                    |  51 +-
 include/linux/migrate.h                |  58 +-
 lib/test_hmm.c                         | 174 ++++-
 lib/test_hmm_uapi.h                    |  21 +-
 mm/Kconfig                             |   1 +
 mm/hmm.c                               | 956 +++++++++++++++++++++++--
 mm/migrate_device.c                    | 617 +++-------------
 tools/testing/selftests/mm/hmm-tests.c |  54 ++
 9 files changed, 1352 insertions(+), 619 deletions(-)

drm-tip
base-commit: dfe5a8188de9aaddd4e46b0f2410d0bccd5c1c04
-- 
2.55.0