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 8358dc144959..0f9e94537eee 100644
--- a/arch/sparc/include/asm/pgtable_64.h
+++ b/arch/sparc/include/asm/pgtable_64.h
@@ -619,7 +619,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;
 }
@@ -847,7 +847,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 3d1999458709..bbc8b12fa443 100644
--- a/arch/x86/include/asm/pgtable.h
+++ b/arch/x86/include/asm/pgtable.h
@@ -452,9 +452,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 f330d28e4d0e..b12079afbd5f 100644
--- a/include/asm-generic/pgtable.h
+++ b/include/asm-generic/pgtable.h
@@ -217,7 +217,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 bd299418a934..e5c49c30460f 100644
--- a/include/linux/mm_types.h
+++ b/include/linux/mm_types.h
@@ -443,6 +443,14 @@ struct mm_struct {
         /* 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;
 };
 
@@ -459,4 +467,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 728d5be9548c..5721f0e3f2da 100644
--- a/kernel/fork.c
+++ b/kernel/fork.c
@@ -537,6 +537,7 @@ static struct mm_struct *mm_init(struct mm_struct *mm, struct task_struct *p)
         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;
diff --git a/mm/huge_memory.c b/mm/huge_memory.c
index 7de1bf85f683..3d2783e10596 100644
--- a/mm/huge_memory.c
+++ b/mm/huge_memory.c
@@ -1377,6 +1377,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 f8421722acb9..bb53a6591aea 100644
--- a/mm/mprotect.c
+++ b/mm/mprotect.c
@@ -188,6 +188,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))
@@ -199,6 +200,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 e84cad27a801..a8b919925934 100644
--- a/mm/pgtable-generic.c
+++ b/mm/pgtable-generic.c
@@ -110,9 +110,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;
 }