unsigned long haddr, bool freeze)
 {
        struct mm_struct *mm = vma->vm_mm;
+       struct folio *folio;
        struct page *page;
        pgtable_t pgtable;
        pmd_t old_pmd, _pmd;
                uffd_wp = pmd_swp_uffd_wp(old_pmd);
        } else {
                page = pmd_page(old_pmd);
+               folio = page_folio(page);
                if (pmd_dirty(old_pmd)) {
                        dirty = true;
-                       SetPageDirty(page);
+                       folio_set_dirty(folio);
                }
                write = pmd_write(old_pmd);
                young = pmd_young(old_pmd);
                soft_dirty = pmd_soft_dirty(old_pmd);
                uffd_wp = pmd_uffd_wp(old_pmd);
 
-               VM_BUG_ON_PAGE(!page_count(page), page);
+               VM_WARN_ON_FOLIO(!folio_ref_count(folio), folio);
+               VM_WARN_ON_FOLIO(!folio_test_anon(folio), folio);
 
                /*
                 * Without "freeze", we'll simply split the PMD, propagating the
                 *
                 * See page_try_share_anon_rmap(): invalidate PMD first.
                 */
-               anon_exclusive = PageAnon(page) && PageAnonExclusive(page);
+               anon_exclusive = PageAnonExclusive(page);
                if (freeze && anon_exclusive && page_try_share_anon_rmap(page))
                        freeze = false;
-               if (!freeze)
-                       page_ref_add(page, HPAGE_PMD_NR - 1);
+               if (!freeze) {
+                       rmap_t rmap_flags = RMAP_NONE;
+
+                       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);
+               }
        }
 
        /*
                        entry = mk_pte(page + i, READ_ONCE(vma->vm_page_prot));
                        if (write)
                                entry = pte_mkwrite(entry, vma);
-                       if (anon_exclusive)
-                               SetPageAnonExclusive(page + i);
                        if (!young)
                                entry = pte_mkold(entry);
                        /* NOTE: this may set soft-dirty too on some archs */
                                entry = pte_mksoft_dirty(entry);
                        if (uffd_wp)
                                entry = pte_mkuffd_wp(entry);
-                       page_add_anon_rmap(page + i, vma, addr, RMAP_NONE);
                }
                VM_BUG_ON(!pte_none(ptep_get(pte)));
                set_pte_at(mm, addr, pte, entry);