#ifdef CONFIG_ARCH_WANT_BATCHED_UNMAP_TLB_FLUSH
 void try_to_unmap_flush(void);
+void try_to_unmap_flush_dirty(void);
 #else
 static inline void try_to_unmap_flush(void)
 {
 }
+static inline void try_to_unmap_flush_dirty(void)
+{
+}
 
 #endif /* CONFIG_ARCH_WANT_BATCHED_UNMAP_TLB_FLUSH */
 #endif /* __MM_INTERNAL_H */
 
        }
        cpumask_clear(&tlb_ubc->cpumask);
        tlb_ubc->flush_required = false;
+       tlb_ubc->writable = false;
        put_cpu();
 }
 
+/* Flush iff there are potentially writable TLB entries that can race with IO */
+void try_to_unmap_flush_dirty(void)
+{
+       struct tlbflush_unmap_batch *tlb_ubc = ¤t->tlb_ubc;
+
+       if (tlb_ubc->writable)
+               try_to_unmap_flush();
+}
+
 static void set_tlb_ubc_flush_pending(struct mm_struct *mm,
-               struct page *page)
+               struct page *page, bool writable)
 {
        struct tlbflush_unmap_batch *tlb_ubc = ¤t->tlb_ubc;
 
        cpumask_or(&tlb_ubc->cpumask, &tlb_ubc->cpumask, mm_cpumask(mm));
        tlb_ubc->flush_required = true;
+
+       /*
+        * If the PTE was dirty then it's best to assume it's writable. The
+        * caller must use try_to_unmap_flush_dirty() or try_to_unmap_flush()
+        * before the page is queued for IO.
+        */
+       if (writable)
+               tlb_ubc->writable = true;
 }
 
 /*
 }
 #else
 static void set_tlb_ubc_flush_pending(struct mm_struct *mm,
-               struct page *page)
+               struct page *page, bool writable)
 {
 }
 
                 */
                pteval = ptep_get_and_clear(mm, address, pte);
 
-               /* Potentially writable TLBs must be flushed before IO */
-               if (pte_dirty(pteval))
-                       flush_tlb_page(vma, address);
-               else
-                       set_tlb_ubc_flush_pending(mm, page);
+               set_tlb_ubc_flush_pending(mm, page, pte_dirty(pteval));
        } else {
                pteval = ptep_clear_flush(vma, address, pte);
        }
 
                        if (!sc->may_writepage)
                                goto keep_locked;
 
-                       /* Page is dirty, try to write it out here */
+                       /*
+                        * Page is dirty. Flush the TLB if a writable entry
+                        * potentially exists to avoid CPU writes after IO
+                        * starts and then write it out here.
+                        */
+                       try_to_unmap_flush_dirty();
                        switch (pageout(page, mapping, sc)) {
                        case PAGE_KEEP:
                                goto keep_locked;