}
 
 #define pte_accessible pte_accessible
-static inline unsigned long pte_accessible(pte_t a)
+static inline unsigned long pte_accessible(struct mm_struct *mm, pte_t a)
 {
        return pte_val(a) & _PAGE_VALID;
 }
         * SUN4V NOTE: _PAGE_VALID is the same value in both the SUN4U
         *             and SUN4V pte layout, so this inline test is fine.
         */
-       if (likely(mm != &init_mm) && pte_accessible(orig))
+       if (likely(mm != &init_mm) && pte_accessible(mm, orig))
                tlb_batch_add(mm, addr, ptep, orig, fullmm);
 }
 
 
 }
 
 #define pte_accessible pte_accessible
-static inline int pte_accessible(pte_t a)
+static inline bool pte_accessible(struct mm_struct *mm, pte_t a)
 {
-       return pte_flags(a) & _PAGE_PRESENT;
+       if (pte_flags(a) & _PAGE_PRESENT)
+               return true;
+
+       if ((pte_flags(a) & (_PAGE_PROTNONE | _PAGE_NUMA)) &&
+                       mm_tlb_flush_pending(mm))
+               return true;
+
+       return false;
 }
 
 static inline int pte_hidden(pte_t pte)
 
 #endif
 
 #ifndef pte_accessible
-# define pte_accessible(pte)           ((void)(pte),1)
+# define pte_accessible(mm, pte)       ((void)(pte), 1)
 #endif
 
 #ifndef flush_tlb_fix_spurious_fault
 
 
        /* numa_scan_seq prevents two threads setting pte_numa */
        int numa_scan_seq;
+#endif
+#if defined(CONFIG_NUMA_BALANCING) || defined(CONFIG_COMPACTION)
+       /*
+        * An operation with batched TLB flushing is going on. Anything that
+        * can move process memory needs to flush the TLB when moving a
+        * PROT_NONE or PROT_NUMA mapped page.
+        */
+       bool tlb_flush_pending;
 #endif
        struct uprobes_state uprobes_state;
 };
        return mm->cpu_vm_mask_var;
 }
 
+#if defined(CONFIG_NUMA_BALANCING) || defined(CONFIG_COMPACTION)
+/*
+ * Memory barriers to keep this state in sync are graciously provided by
+ * the page table locks, outside of which no page table modifications happen.
+ * The barriers below prevent the compiler from re-ordering the instructions
+ * around the memory barriers that are already present in the code.
+ */
+static inline bool mm_tlb_flush_pending(struct mm_struct *mm)
+{
+       barrier();
+       return mm->tlb_flush_pending;
+}
+static inline void set_tlb_flush_pending(struct mm_struct *mm)
+{
+       mm->tlb_flush_pending = true;
+       barrier();
+}
+/* Clearing is done after a TLB flush, which also provides a barrier. */
+static inline void clear_tlb_flush_pending(struct mm_struct *mm)
+{
+       barrier();
+       mm->tlb_flush_pending = false;
+}
+#else
+static inline bool mm_tlb_flush_pending(struct mm_struct *mm)
+{
+       return false;
+}
+static inline void set_tlb_flush_pending(struct mm_struct *mm)
+{
+}
+static inline void clear_tlb_flush_pending(struct mm_struct *mm)
+{
+}
+#endif
+
 #endif /* _LINUX_MM_TYPES_H */
 
        spin_lock_init(&mm->page_table_lock);
        mm_init_aio(mm);
        mm_init_owner(mm, p);
+       clear_tlb_flush_pending(mm);
 
        if (likely(!mm_alloc_pgd(mm))) {
                mm->def_flags = 0;
 
                goto clear_pmdnuma;
        }
 
+       /*
+        * The page_table_lock above provides a memory barrier
+        * with change_protection_range.
+        */
+       if (mm_tlb_flush_pending(mm))
+               flush_tlb_range(vma, haddr, haddr + HPAGE_PMD_SIZE);
+
        /*
         * Migrate the THP to the requested node, returns with page unlocked
         * and pmd_numa cleared.
 
        BUG_ON(addr >= end);
        pgd = pgd_offset(mm, addr);
        flush_cache_range(vma, addr, end);
+       set_tlb_flush_pending(mm);
        do {
                next = pgd_addr_end(addr, end);
                if (pgd_none_or_clear_bad(pgd))
        /* Only flush the TLB if we actually modified any entries: */
        if (pages)
                flush_tlb_range(vma, start, end);
+       clear_tlb_flush_pending(mm);
 
        return pages;
 }
 
 pte_t ptep_clear_flush(struct vm_area_struct *vma, unsigned long address,
                       pte_t *ptep)
 {
+       struct mm_struct *mm = (vma)->vm_mm;
        pte_t pte;
-       pte = ptep_get_and_clear((vma)->vm_mm, address, ptep);
-       if (pte_accessible(pte))
+       pte = ptep_get_and_clear(mm, address, ptep);
+       if (pte_accessible(mm, pte))
                flush_tlb_page(vma, address);
        return pte;
 }