@@ -343,7 +343,6 @@ unsigned long thp_get_unmapped_area_vmflags(struct file *filp, unsigned long add
 		vm_flags_t vm_flags);
 
 bool can_split_folio(struct folio *folio, int caller_pins, int *pextra_pins);
-int split_device_private_folio(struct folio *folio);
 int __split_huge_page_to_list_to_order(struct page *page, struct list_head *list,
 		unsigned int new_order, bool unmapped);
 int min_order_for_split(struct folio *folio);

@@ -1625,13 +1625,22 @@ static vm_fault_t dmirror_devmem_fault(struct vm_fault *vmf)
 	 * the mirror but here we use it to hold the page for the simulated
 	 * device memory and that page holds the pointer to the mirror.
 	 */
-	rpage = vmf->page->zone_device_data;
+	rpage = folio_page(page_folio(vmf->page), 0)->zone_device_data;
 	dmirror = rpage->zone_device_data;
 
 	/* FIXME demonstrate how we can adjust migrate range */
 	order = folio_order(page_folio(vmf->page));
 	nr = 1 << order;
 
+	/*
+	 * When folios are partially mapped, we can't rely on the folio
+	 * order of vmf->page as the folio might not be fully split yet
+	 */
+	if (vmf->pte) {
+		order = 0;
+		nr = 1;
+	}
+
 	/*
 	 * Consider a per-cpu cache of src and dst pfns, but with
 	 * large number of cpus that might not scale well.

@@ -72,10 +72,6 @@ static unsigned long deferred_split_count(struct shrinker *shrink,
 					  struct shrink_control *sc);
 static unsigned long deferred_split_scan(struct shrinker *shrink,
 					 struct shrink_control *sc);
-static int __split_unmapped_folio(struct folio *folio, int new_order,
-		struct page *split_at, struct xa_state *xas,
-		struct address_space *mapping, bool uniform_split);
-
 static bool split_underused_thp = true;
 
 static atomic_t huge_zero_refcount;

@@ -2924,51 +2920,6 @@ static void __split_huge_zero_page_pmd(struct vm_area_struct *vma,
 	pmd_populate(mm, pmd, pgtable);
 }
 
-/**
- * split_huge_device_private_folio - split a huge device private folio into
- * smaller pages (of order 0), currently used by migrate_device logic to
- * split folios for pages that are partially mapped
- *
- * @folio: the folio to split
- *
- * The caller has to hold the folio_lock and a reference via folio_get
- */
-int split_device_private_folio(struct folio *folio)
-{
-	struct folio *end_folio = folio_next(folio);
-	struct folio *new_folio;
-	int ret = 0;
-
-	/*
-	 * Split the folio now. In the case of device
-	 * private pages, this path is executed when
-	 * the pmd is split and since freeze is not true
-	 * it is likely the folio will be deferred_split.
-	 *
-	 * With device private pages, deferred splits of
-	 * folios should be handled here to prevent partial
-	 * unmaps from causing issues later on in migration
-	 * and fault handling flows.
-	 */
-	folio_ref_freeze(folio, 1 + folio_expected_ref_count(folio));
-	ret = __split_unmapped_folio(folio, 0, &folio->page, NULL, NULL, true);
-	VM_WARN_ON(ret);
-	for (new_folio = folio_next(folio); new_folio != end_folio;
-					new_folio = folio_next(new_folio)) {
-		zone_device_private_split_cb(folio, new_folio);
-		folio_ref_unfreeze(new_folio, 1 + folio_expected_ref_count(
-								new_folio));
-	}
-
-	/*
-	 * Mark the end of the folio split for device private THP
-	 * split
-	 */
-	zone_device_private_split_cb(folio, NULL);
-	folio_ref_unfreeze(folio, 1 + folio_expected_ref_count(folio));
-	return ret;
-}
-
 static void __split_huge_pmd_locked(struct vm_area_struct *vma, pmd_t *pmd,
 		unsigned long haddr, bool freeze)
 {

@@ -3064,30 +3015,15 @@ static void __split_huge_pmd_locked(struct vm_area_struct *vma, pmd_t *pmd,
 				freeze = false;
 			if (!freeze) {
 				rmap_t rmap_flags = RMAP_NONE;
-				unsigned long addr = haddr;
-				struct folio *new_folio;
-				struct folio *end_folio = folio_next(folio);
 
 				if (anon_exclusive)
 					rmap_flags |= RMAP_EXCLUSIVE;
 
-				folio_lock(folio);
-				folio_get(folio);
-
-				split_device_private_folio(folio);
-
-				for (new_folio = folio_next(folio);
-					new_folio != end_folio;
-					new_folio = folio_next(new_folio)) {
-					addr += PAGE_SIZE;
-					folio_unlock(new_folio);
-					folio_add_anon_rmap_ptes(new_folio,
-						&new_folio->page, 1,
-						vma, addr, rmap_flags);
-				}
-				folio_unlock(folio);
-				folio_add_anon_rmap_ptes(folio, &folio->page,
-						1, vma, haddr, rmap_flags);
+				folio_ref_add(folio, HPAGE_PMD_NR - 1);
+				if (anon_exclusive)
+					rmap_flags |= RMAP_EXCLUSIVE;
+				folio_add_anon_rmap_ptes(folio, page, HPAGE_PMD_NR,
+						 vma, haddr, rmap_flags);
 			}
 		}
 

@@ -4065,7 +4001,7 @@ static int __folio_split(struct folio *folio, unsigned int new_order,
 	if (nr_shmem_dropped)
 		shmem_uncharge(mapping->host, nr_shmem_dropped);
 
-	if (!ret && is_anon)
+	if (!ret && is_anon && !folio_is_device_private(folio))
 		remap_flags = RMP_USE_SHARED_ZEROPAGE;
 
 	remap_page(folio, 1 << order, remap_flags);

@@ -248,6 +248,8 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
 			 * page table entry. Other special swap entries are not
 			 * migratable, and we ignore regular swapped page.
 			 */
+			struct folio *folio;
+
 			entry = pte_to_swp_entry(pte);
 			if (!is_device_private_entry(entry))
 				goto next;

@@ -259,6 +261,55 @@ static int migrate_vma_collect_pmd(pmd_t *pmdp,
 			    pgmap->owner != migrate->pgmap_owner)
 				goto next;
 
+			folio = page_folio(page);
+			if (folio_test_large(folio)) {
+				struct folio *new_folio;
+				struct folio *new_fault_folio;
+
+				/*
+				 * The reason for finding pmd present with a
+				 * device private pte and a large folio for the
+				 * pte is partial unmaps. Split the folio now
+				 * for the migration to be handled correctly
+				 */
+				pte_unmap_unlock(ptep, ptl);
+
+				folio_get(folio);
+				if (folio != fault_folio)
+					folio_lock(folio);
+				if (split_folio(folio)) {
+					if (folio != fault_folio)
+						folio_unlock(folio);
+					ptep = pte_offset_map_lock(mm, pmdp, addr, &ptl);
+					goto next;
+				}
+
+				/*
+				 * After the split, get back the extra reference
+				 * on the fault_page, this reference is checked during
+				 * folio_migrate_mapping()
+				 */
+				if (migrate->fault_page) {
+					new_fault_folio = page_folio(migrate->fault_page);
+					folio_get(new_fault_folio);
+				}
+
+				new_folio = page_folio(page);
+				pfn = page_to_pfn(page);
+
+				/*
+				 * Ensure the lock is held on the correct
+				 * folio after the split
+				 */
+				if (folio != new_folio) {
+					folio_unlock(folio);
+					folio_lock(new_folio);
+				}
+				folio_put(folio);
+				addr = start;
+				goto again;
+			}
+
 			mpfn = migrate_pfn(page_to_pfn(page)) |
 					MIGRATE_PFN_MIGRATE;
 			if (is_writable_device_private_entry(entry))