mm/migrate: calculate large folio page index using PFN

Rather than having to figure out the page index to use using
linear_page_index(), calculate it using PFN.

This is a more natural fit as the linear page index is immaterial to
determining the folio page index.

Derive the page index from the offset between migration entry PFN and
folio PFN - pvmw.pfn (set via DEFINE_FOLIO_VMA_WALK() which uses
folio_pfn() to obtain it).

Additionally remove a not so useful comment and clean the code layout up.

No functional change intended.

Link: https://lore.kernel.org/20260813-b4-scalable-cow-virt-pgoff-v5-13-c21581c0c3c8@kernel.org
Signed-off-by: Lorenzo Stoakes (ARM) <[email protected]>
Suggested-by: David Hildenbrand (Arm) <[email protected]>
Acked-by: David Hildenbrand (Arm) <[email protected]>
Cc: Adrian Hunter <[email protected]>
Cc: Alexander Deucher <[email protected]>
Cc: Alexander Gordeev <[email protected]>
Cc: Alexander Shishkin <[email protected]>
Cc: Alistair Popple <[email protected]>
Cc: Arnaldo Carvalho de Melo <[email protected]>
Cc: Arnd Bergmann <[email protected]>
Cc: Baolin Wang <[email protected]>
Cc: Baoquan He <[email protected]>
Cc: Barry Song <[email protected]>
Cc: Boris Brezillon <[email protected]>
Cc: Byungchul Park <[email protected]>
Cc: Chengming Zhou <[email protected]>
Cc: Chris Li <[email protected]>
Cc: Christan König <[email protected]>
Cc: Christian Borntraeger <[email protected]>
Cc: Claudio Imbrenda <[email protected]>
Cc: Dave Airlie <[email protected]>
Cc: Dev Jain <[email protected]>
Cc: Gerald Schaefer <[email protected]>
Cc: Greg Kroah-Hartman <[email protected]>
Cc: Gregory Price (Meta) <[email protected]>
Cc: Harry Yoo <[email protected]>
Cc: Heiko Carstens <[email protected]>
Cc: Huang Ray <[email protected]>
Cc: "Huang, Ying" <[email protected]>
Cc: Ian Rogers <[email protected]>
Cc: Ingo Molnar <[email protected]>
Cc: James Clark <[email protected]>
Cc: Jan Kara <[email protected]>
Cc: Jann Horn <[email protected]>
Cc: Janosch Frank <[email protected]>
Cc: Jason Gunthorpe <[email protected]>
Cc: Jiri Olsa <[email protected]>
Cc: John Hubbard <[email protected]>
Cc: Joshua Hahn <[email protected]>
Cc: Kairui Song <[email protected]>
Cc: Kees Cook <[email protected]>
Cc: Kemeng Shi <[email protected]>
Cc: Lance Yang <[email protected]>
Cc: Liam R. Howlett <[email protected]>
Cc: Liviu Dudau <[email protected]>
Cc: Maarten Lankhorst <[email protected]>
Cc: Marc Rutland <[email protected]>
Cc: "Masami Hiramatsu (Google)" <[email protected]>
Cc: Matthew Auld <[email protected]>
Cc: Matthew Brost <[email protected]>
Cc: Matthew Wilcox (Oracle) <[email protected]>
Cc: Maxime Ripard <[email protected]>
Cc: Miaohe Lin <[email protected]>
Cc: Michal Hocko <[email protected]>
Cc: Mike Rapoport <[email protected]>
Cc: Muchun Song <[email protected]>
Cc: Namhyung kim <[email protected]>
Cc: Naoya Horiguchi <[email protected]>
Cc: Nhat Pham <[email protected]>
Cc: Nico Pache <[email protected]>
Cc: Oleg Nesterov <[email protected]>
Cc: Oscar Salvador <[email protected]>
Cc: Pedro Falcato <[email protected]>
Cc: Peter Xu <[email protected]>
Cc: Peter Zijlstra <[email protected]>
Cc: Rakie Kim <[email protected]>
Cc: Rik van Riel <[email protected]>
Cc: Rodrigo Vivi <[email protected]>
Cc: Ryan Roberts <[email protected]>
Cc: Steven Price <[email protected]>
Cc: Suren Baghdasaryan <[email protected]>
Cc: Sven Schnelle <[email protected]>
Cc: Thomas Hellström <[email protected]>
Cc: Thomas Zimemrmann <[email protected]>
Cc: Vasily Gorbik <[email protected]>
Cc: Vlastimil Babka <[email protected]>
Cc: xu xin <[email protected]>
Cc: Zi Yan <[email protected]>
Signed-off-by: Andrew Morton <[email protected]>

mm/migrate.c

@@ -356,16 +356,11 @@ static bool remove_migration_pte(struct folio *folio,
 
 	while (page_vma_mapped_walk(&pvmw)) {
 		rmap_t rmap_flags = RMAP_NONE;
-		pte_t old_pte;
-		pte_t pte;
+		unsigned long idx = 0;
 		softleaf_t entry;
 		struct page *new;
-		unsigned long idx = 0;
-
-		/* pgoff is invalid for ksm pages, but they are never large */
-		if (folio_test_large(folio) && !folio_test_hugetlb(folio))
-			idx = linear_page_index(vma, pvmw.address) - pvmw.pgoff;
-		new = folio_page(folio, idx);
+		pte_t old_pte;
+		pte_t pte;
 
 #ifdef CONFIG_ARCH_HAS_PMD_SOFTLEAVES
 		/* PMD-mapped THP migration entry */
@@ -381,14 +376,18 @@ static bool remove_migration_pte(struct folio *folio,
 						pvmw.pte);
 		else
 			old_pte = ptep_get(pvmw.pte);
+
+		entry = softleaf_from_pte(old_pte);
+		if (folio_test_large(folio) && !folio_test_hugetlb(folio))
+			idx = softleaf_to_pfn(entry) - pvmw.pfn;
+
 		if (rmap_walk_arg->map_unused_to_zeropage &&
 		    try_to_map_unused_to_zeropage(&pvmw, folio, old_pte, idx))
 			continue;
 
 		folio_get(folio);
+		new = folio_page(folio, idx);
 		pte = mk_pte(new, READ_ONCE(vma->vm_page_prot));
-
-		entry = softleaf_from_pte(old_pte);
 		if (!softleaf_is_migration_young(entry))
 			pte = pte_mkold(pte);
 		if (folio_test_dirty(folio) && softleaf_is_migration_dirty(entry))