[PATCH v2] mm/migrate_device: do not access pgmap for non zone device pages

[Balbir Singh <balbirs@nvidia.com>]

page_pgmap() is referenced before checking if the page is a zone device
page and this triggers the warning in page_pgmap().  Refactor the code to
use the helper function after relevant checks.

Link: https://lkml.kernel.org/r/20250219231337.364133-1-balbirs@nvidia.com
Signed-off-by: Balbir Singh <balbirs@nvidia.com>
Cc: Alistair Popple <apopple@nvidia.com>
Cc: Jason Gunthorpe <jgg@ziepe.ca>
Cc: David Hildenbrand <david@redhat.com>
Cc: Dan Williams <dan.j.williams@intel.com>
Acked-by: David Hildenbrand <david@redhat.com>
Reviewed-by: Alistair Popple <apopple@nvidia.com>
---

This patch replaces 11ba213916b9 and fixes 7f1cfd71153b
("mm: allow compound zone device pages") on mm-unstable

Changelog v2:
 - Added {} around the if statement as per coding style
 - Added Reviewed-by/Acked-by tags

 mm/migrate_device.c | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/mm/migrate_device.c b/mm/migrate_device.c
index 6771893d4601..7d0d64f67cdf 100644
--- a/mm/migrate_device.c
+++ b/mm/migrate_device.c
@@ -153,14 +153,17 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
 				goto next;
 			}
 			page = vm_normal_page(migrate->vma, addr, pte);
-			pgmap = page_pgmap(page);
 			if (page && !is_zone_device_page(page) &&
-			    !(migrate->flags & MIGRATE_VMA_SELECT_SYSTEM))
-				goto next;
-			else if (page && is_device_coherent_page(page) &&
-			    (!(migrate->flags & MIGRATE_VMA_SELECT_DEVICE_COHERENT) ||
-			     pgmap->owner != migrate->pgmap_owner))
+			    !(migrate->flags & MIGRATE_VMA_SELECT_SYSTEM)) {
 				goto next;
+			} else if (page && is_device_coherent_page(page)) {
+				pgmap = page_pgmap(page);
+
+				if (!(migrate->flags &
+					MIGRATE_VMA_SELECT_DEVICE_COHERENT) ||
+					pgmap->owner != migrate->pgmap_owner)
+					goto next;
+			}
 			mpfn = migrate_pfn(pfn) | MIGRATE_PFN_MIGRATE;
 			mpfn |= pte_write(pte) ? MIGRATE_PFN_WRITE : 0;
 		}
-- 
2.48.1