struct page_vma_mapped_walk;
 
 #ifdef CONFIG_ARCH_ENABLE_THP_MIGRATION
-extern void set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
+extern int set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
                struct page *page);
 
 extern void remove_migration_pmd(struct page_vma_mapped_walk *pvmw,
        return !pmd_present(pmd) && is_migration_entry(pmd_to_swp_entry(pmd));
 }
 #else
-static inline void set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
+static inline int set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
                struct page *page)
 {
        BUILD_BUG();
 
 #endif
 
 #ifdef CONFIG_ARCH_ENABLE_THP_MIGRATION
-void set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
+int set_pmd_migration_entry(struct page_vma_mapped_walk *pvmw,
                struct page *page)
 {
        struct vm_area_struct *vma = pvmw->vma;
        pmd_t pmdswp;
 
        if (!(pvmw->pmd && !pvmw->pte))
-               return;
+               return 0;
 
        flush_cache_range(vma, address, address + HPAGE_PMD_SIZE);
        pmdval = pmdp_invalidate(vma, address, pvmw->pmd);
        anon_exclusive = PageAnon(page) && PageAnonExclusive(page);
        if (anon_exclusive && page_try_share_anon_rmap(page)) {
                set_pmd_at(mm, address, pvmw->pmd, pmdval);
-               return;
+               return -EBUSY;
        }
 
        if (pmd_dirty(pmdval))
        page_remove_rmap(page, vma, true);
        put_page(page);
        trace_set_migration_pmd(address, pmd_val(pmdswp));
+
+       return 0;
 }
 
 void remove_migration_pmd(struct page_vma_mapped_walk *pvmw, struct page *new)
 
                        VM_BUG_ON_FOLIO(folio_test_hugetlb(folio) ||
                                        !folio_test_pmd_mappable(folio), folio);
 
-                       set_pmd_migration_entry(&pvmw, subpage);
+                       if (set_pmd_migration_entry(&pvmw, subpage)) {
+                               ret = false;
+                               page_vma_mapped_walk_done(&pvmw);
+                               break;
+                       }
                        continue;
                }
 #endif