#include <linux/hugetlb_cgroup.h>
 #include <linux/gfp.h>
 #include <linux/balloon_compaction.h>
+#include <linux/mmu_notifier.h>
 
 #include <asm/tlbflush.h>
 
                                struct page *page, int node)
 {
        spinlock_t *ptl;
-       unsigned long haddr = address & HPAGE_PMD_MASK;
        pg_data_t *pgdat = NODE_DATA(node);
        int isolated = 0;
        struct page *new_page = NULL;
        struct mem_cgroup *memcg = NULL;
        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;
 
        /*
        WARN_ON(PageLRU(new_page));
 
        /* 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:
                spin_unlock(ptl);
+               mmu_notifier_invalidate_range_end(mm, mmun_start, mmun_end);
 
                /* Reverse changes made by migrate_page_copy() */
                if (TestClearPageActive(new_page))
         * The SetPageUptodate on the new page and page_add_new_anon_rmap
         * guarantee the copy is visible before the pagetable update.
         */
-       flush_cache_range(vma, haddr, haddr + HPAGE_PMD_SIZE);
-       page_add_new_anon_rmap(new_page, vma, haddr);
-       pmdp_clear_flush(vma, haddr, pmd);
-       set_pmd_at(mm, haddr, pmd, entry);
+       flush_cache_range(vma, mmun_start, mmun_end);
+       page_add_new_anon_rmap(new_page, vma, mmun_start);
+       pmdp_clear_flush(vma, mmun_start, pmd);
+       set_pmd_at(mm, mmun_start, pmd, entry);
+       flush_tlb_range(vma, mmun_start, mmun_end);
        update_mmu_cache_pmd(vma, address, &entry);
 
        if (page_count(page) != 2) {
-               set_pmd_at(mm, haddr, pmd, orig_entry);
-               flush_tlb_range(vma, haddr, haddr + HPAGE_PMD_SIZE);
+               set_pmd_at(mm, mmun_start, pmd, orig_entry);
+               flush_tlb_range(vma, mmun_start, mmun_end);
                update_mmu_cache_pmd(vma, address, &entry);
                page_remove_rmap(new_page);
                goto fail_putback;
         */
        mem_cgroup_end_migration(memcg, page, new_page, true);
        spin_unlock(ptl);
+       mmu_notifier_invalidate_range_end(mm, mmun_start, mmun_end);
 
        unlock_page(new_page);
        unlock_page(page);
        ptl = pmd_lock(mm, pmd);
        if (pmd_same(*pmd, entry)) {
                entry = pmd_mknonnuma(entry);
-               set_pmd_at(mm, haddr, pmd, entry);
+               set_pmd_at(mm, mmun_start, pmd, entry);
                update_mmu_cache_pmd(vma, address, &entry);
        }
        spin_unlock(ptl);