mm: Use a folio in the softleaf_is_device_private path

From: Hongfu Li <[email protected]>

Use the folio APIs in the device_private migration path of do_swap_page(), replacing four calls to compound_head() with two page_folio() calls. The second one re-fetches the folio from vmf->page after migrate_to_ram(), which might have split the folio.

---
v5:
- Re-fetch the folio from vmf->page after migrate_to_ram(), which might have split it
- Update commit message
- Add Suggested-by and Reviewed-by tags
v4:
- Drop the do_non_swap_page() extraction (patch 2/2 from v2)
- Added Reviewed-by: Lorenzo Stoakes
- Link to v2: https://lore.kernel.org/all/[email protected]/
v3:
- Please ignore v3.
v2:
- Drop fault_folio, reuse the existing folio variable instead
- Refresh commit messages
- Update comment to use "folio" instead of "page" in the
device_private path
- Add Acked-by: David Hildenbrand
- Add patch 2/2 to extract do_non_swap_page() from do_swap_page()
- Link to v1: https://lore.kernel.org/all/[email protected]/

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

index c54943302553..c6a9242d1b5b 100644
@@ -4926,18 +4926,21 @@ vm_fault_t do_swap_page(struct vm_fault *vmf)
 				goto unlock;
 
 			/*
-			 * Get a page reference while we know the page can't be
-			 * freed.
+			 * Get a folio reference while we know the folio can't
+			 * be freed.
 			 */
-			if (trylock_page(vmf->page)) {
+			folio = page_folio(vmf->page);
+			if (folio_trylock(folio)) {
 				struct dev_pagemap *pgmap;
 
-				get_page(vmf->page);
+				folio_get(folio);
 				pte_unmap_unlock(vmf->pte, vmf->ptl);
 				pgmap = page_pgmap(vmf->page);
 				ret = pgmap->ops->migrate_to_ram(vmf);
-				unlock_page(vmf->page);
-				put_page(vmf->page);
+				/* migrate_to_ram() might have split the folio. */
+				folio = page_folio(vmf->page);
+				folio_unlock(folio);
+				folio_put(folio);
 			} else {
 				pte_unmap(vmf->pte);
 				softleaf_entry_wait_on_locked(entry, vmf->ptl);

2.54.0