int page_lru = page_is_file_cache(page);
        unsigned long mmun_start = address & HPAGE_PMD_MASK;
        unsigned long mmun_end = mmun_start + HPAGE_PMD_SIZE;
-       pmd_t orig_entry;
 
        /*
         * Rate-limit the amount of data that is being migrated to a node.
        /* Recheck the target PMD */
        mmu_notifier_invalidate_range_start(mm, mmun_start, mmun_end);
        ptl = pmd_lock(mm, pmd);
-       if (unlikely(!pmd_same(*pmd, entry) || page_count(page) != 2)) {
-fail_putback:
+       if (unlikely(!pmd_same(*pmd, entry) || !page_ref_freeze(page, 2))) {
                spin_unlock(ptl);
                mmu_notifier_invalidate_range_end(mm, mmun_start, mmun_end);
 
                goto out_unlock;
        }
 
-       orig_entry = *pmd;
        entry = mk_huge_pmd(new_page, vma->vm_page_prot);
        entry = maybe_pmd_mkwrite(pmd_mkdirty(entry), vma);
 
        set_pmd_at(mm, mmun_start, pmd, entry);
        update_mmu_cache_pmd(vma, address, &entry);
 
-       if (page_count(page) != 2) {
-               set_pmd_at(mm, mmun_start, pmd, orig_entry);
-               flush_pmd_tlb_range(vma, mmun_start, mmun_end);
-               mmu_notifier_invalidate_range(mm, mmun_start, mmun_end);
-               update_mmu_cache_pmd(vma, address, &entry);
-               page_remove_rmap(new_page, true);
-               goto fail_putback;
-       }
-
+       page_ref_unfreeze(page, 2);
        mlock_migrate_page(new_page, page);
        page_remove_rmap(page, true);
        set_page_owner_migrate_reason(new_page, MR_NUMA_MISPLACED);