3 files changed, 24 insertions, 18 deletions

diff --git a/mm/hugetlb.c b/mm/hugetlb.c
index 5cf99bf8cce2..7c5eb85ec645 100644
--- a/mm/hugetlb.c
+++ b/mm/hugetlb.c
@@ -2490,7 +2490,7 @@ void unmap_hugepage_range(struct vm_area_struct *vma, unsigned long start,
 
         mm = vma->vm_mm;
 
-        tlb_gather_mmu(&tlb, mm, 0);
+        tlb_gather_mmu(&tlb, mm, start, end);
         __unmap_hugepage_range(&tlb, vma, start, end, ref_page);
         tlb_finish_mmu(&tlb, start, end);
 }
diff --git a/mm/memory.c b/mm/memory.c
index 5e5080005bc4..5a35443c01ad 100644
--- a/mm/memory.c
+++ b/mm/memory.c
@@ -211,14 +211,15 @@ static int tlb_next_batch(struct mmu_gather *tlb)
  *      tear-down from @mm. The @fullmm argument is used when @mm is without
  *      users and we're going to destroy the full address space (exit/execve).
  */
-void tlb_gather_mmu(struct mmu_gather *tlb, struct mm_struct *mm, bool fullmm)
+void tlb_gather_mmu(struct mmu_gather *tlb, struct mm_struct *mm, unsigned long start, unsigned long end)
 {
         tlb->mm = mm;
 
-        tlb->fullmm     = fullmm;
+        /* Is it from 0 to ~0? */
+        tlb->fullmm     = !(start | (end+1));
         tlb->need_flush_all = 0;
-        tlb->start      = -1UL;
-        tlb->end        = 0;
+        tlb->start      = start;
+        tlb->end        = end;
         tlb->need_flush = 0;
         tlb->local.next = NULL;
         tlb->local.nr   = 0;
@@ -258,8 +259,6 @@ void tlb_finish_mmu(struct mmu_gather *tlb, unsigned long start, unsigned long e
 {
         struct mmu_gather_batch *batch, *next;
 
-        tlb->start = start;
-        tlb->end   = end;
         tlb_flush_mmu(tlb);
 
         /* keep the page table cache within bounds */
@@ -1101,7 +1100,6 @@ static unsigned long zap_pte_range(struct mmu_gather *tlb,
         spinlock_t *ptl;
         pte_t *start_pte;
         pte_t *pte;
-        unsigned long range_start = addr;
 
 again:
         init_rss_vec(rss);
@@ -1204,17 +1202,25 @@ again:
          * and page-free while holding it.
          */
         if (force_flush) {
+                unsigned long old_end;
+
                 force_flush = 0;
 
-#ifdef HAVE_GENERIC_MMU_GATHER
-                tlb->start = range_start;
+                /*
+                 * Flush the TLB just for the previous segment,
+                 * then update the range to be the remaining
+                 * TLB range.
+                 */
+                old_end = tlb->end;
                 tlb->end = addr;
-#endif
+
                 tlb_flush_mmu(tlb);
-                if (addr != end) {
-                        range_start = addr;
+
+                tlb->start = addr;
+                tlb->end = old_end;
+
+                if (addr != end)
                         goto again;
-                }
         }
 
         return addr;
@@ -1399,7 +1405,7 @@ void zap_page_range(struct vm_area_struct *vma, unsigned long start,
         unsigned long end = start + size;
 
         lru_add_drain();
-        tlb_gather_mmu(&tlb, mm, 0);
+        tlb_gather_mmu(&tlb, mm, start, end);
         update_hiwater_rss(mm);
         mmu_notifier_invalidate_range_start(mm, start, end);
         for ( ; vma && vma->vm_start < end; vma = vma->vm_next)
@@ -1425,7 +1431,7 @@ static void zap_page_range_single(struct vm_area_struct *vma, unsigned long addr
         unsigned long end = address + size;
 
         lru_add_drain();
-        tlb_gather_mmu(&tlb, mm, 0);
+        tlb_gather_mmu(&tlb, mm, address, end);
         update_hiwater_rss(mm);
         mmu_notifier_invalidate_range_start(mm, address, end);
         unmap_single_vma(&tlb, vma, address, end, details);
diff --git a/mm/mmap.c b/mm/mmap.c
index 7dbe39745be9..8d25fdc653be 100644
--- a/mm/mmap.c
+++ b/mm/mmap.c
@@ -2356,7 +2356,7 @@ static void unmap_region(struct mm_struct *mm,
         struct mmu_gather tlb;
 
         lru_add_drain();
-        tlb_gather_mmu(&tlb, mm, 0);
+        tlb_gather_mmu(&tlb, mm, start, end);
         update_hiwater_rss(mm);
         unmap_vmas(&tlb, vma, start, end);
         free_pgtables(&tlb, vma, prev ? prev->vm_end : FIRST_USER_ADDRESS,
@@ -2735,7 +2735,7 @@ void exit_mmap(struct mm_struct *mm)
 
         lru_add_drain();
         flush_cache_mm(mm);
-        tlb_gather_mmu(&tlb, mm, 1);
+        tlb_gather_mmu(&tlb, mm, 0, -1);
         /* update_hiwater_rss(mm) here? but nobody should be looking */
         /* Use -1 here to ensure all VMAs in the mm are unmapped */
         unmap_vmas(&tlb, vma, 0, -1);