8 files changed, 69 insertions, 7 deletions

diff --git a/arch/sparc/include/asm/pgtable_64.h b/arch/sparc/include/asm/pgtable_64.h
index 7619f2f792af..dfb0019bf05b 100644
--- a/arch/sparc/include/asm/pgtable_64.h
+++ b/arch/sparc/include/asm/pgtable_64.h
@@ -616,7 +616,7 @@ static inline unsigned long pte_present(pte_t pte)
 }
 
 #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;
 }
@@ -806,7 +806,7 @@ static inline void __set_pte_at(struct mm_struct *mm, unsigned long addr,
          * 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);
 }
 
diff --git a/arch/x86/include/asm/pgtable.h b/arch/x86/include/asm/pgtable.h
index 1e672234c4ff..5460bf923e16 100644
--- a/arch/x86/include/asm/pgtable.h
+++ b/arch/x86/include/asm/pgtable.h
@@ -415,9 +415,16 @@ static inline int pte_present(pte_t a)
 }
 
 #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)
diff --git a/include/asm-generic/pgtable.h b/include/asm-generic/pgtable.h
index a59ff51b0166..b58268a5ddd4 100644
--- a/include/asm-generic/pgtable.h
+++ b/include/asm-generic/pgtable.h
@@ -220,7 +220,7 @@ static inline int pmd_same(pmd_t pmd_a, pmd_t pmd_b)
 #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
diff --git a/include/linux/mm_types.h b/include/linux/mm_types.h
index 4a189ba6b128..49f0ada525a8 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -437,6 +437,14 @@ struct mm_struct {
          */
         int first_nid;
 #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;
 };
 
@@ -457,4 +465,40 @@ static inline cpumask_t *mm_cpumask(struct mm_struct *mm)
         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 */
diff --git a/kernel/fork.c b/kernel/fork.c
index 80d92e987f21..ff7be9dac4c1 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -544,6 +544,7 @@ static struct mm_struct *mm_init(struct mm_struct *mm, struct task_struct *p)
         mm->cached_hole_size = ~0UL;
         mm_init_aio(mm);
         mm_init_owner(mm, p);
+        clear_tlb_flush_pending(mm);
 
         if (likely(!mm_alloc_pgd(mm))) {
                 mm->def_flags = 0;
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index b2e803e14ea9..6bd22902d289 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -1352,6 +1352,13 @@ int do_huge_pmd_numa_page(struct mm_struct *mm, struct vm_area_struct *vma,
         }
 
         /*
+         * 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.
          */
diff --git a/mm/mprotect.c b/mm/mprotect.c
index d4d5399c7aba..e9f65aaa3182 100644
--- a/mm/mprotect.c
+++ b/mm/mprotect.c
@@ -206,6 +206,7 @@ static unsigned long change_protection_range(struct vm_area_struct *vma,
         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))
@@ -217,6 +218,7 @@ static unsigned long change_protection_range(struct vm_area_struct *vma,
         /* 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;
 }
diff --git a/mm/pgtable-generic.c b/mm/pgtable-generic.c
index eb900bbaa713..4b62a16fc3c1 100644
--- a/mm/pgtable-generic.c
+++ b/mm/pgtable-generic.c
@@ -86,9 +86,10 @@ int pmdp_clear_flush_young(struct vm_area_struct *vma,
 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;
 }