#include <linux/sched.h>
 #include <linux/page_table_check.h>
 
+static inline void emit_pte_barriers(void)
+{
+       /*
+        * These barriers are emitted under certain conditions after a pte entry
+        * was modified (see e.g. __set_pte_complete()). The dsb makes the store
+        * visible to the table walker. The isb ensures that any previous
+        * speculative "invalid translation" marker that is in the CPU's
+        * pipeline gets cleared, so that any access to that address after
+        * setting the pte to valid won't cause a spurious fault. If the thread
+        * gets preempted after storing to the pgtable but before emitting these
+        * barriers, __switch_to() emits a dsb which ensure the walker gets to
+        * see the store. There is no guarantee of an isb being issued though.
+        * This is safe because it will still get issued (albeit on a
+        * potentially different CPU) when the thread starts running again,
+        * before any access to the address.
+        */
+       dsb(ishst);
+       isb();
+}
+
+static inline void queue_pte_barriers(void)
+{
+       unsigned long flags;
+
+       VM_WARN_ON(in_interrupt());
+       flags = read_thread_flags();
+
+       if (flags & BIT(TIF_LAZY_MMU)) {
+               /* Avoid the atomic op if already set. */
+               if (!(flags & BIT(TIF_LAZY_MMU_PENDING)))
+                       set_thread_flag(TIF_LAZY_MMU_PENDING);
+       } else {
+               emit_pte_barriers();
+       }
+}
+
+#define  __HAVE_ARCH_ENTER_LAZY_MMU_MODE
+static inline void arch_enter_lazy_mmu_mode(void)
+{
+       VM_WARN_ON(in_interrupt());
+       VM_WARN_ON(test_thread_flag(TIF_LAZY_MMU));
+
+       set_thread_flag(TIF_LAZY_MMU);
+}
+
+static inline void arch_flush_lazy_mmu_mode(void)
+{
+       if (test_and_clear_thread_flag(TIF_LAZY_MMU_PENDING))
+               emit_pte_barriers();
+}
+
+static inline void arch_leave_lazy_mmu_mode(void)
+{
+       arch_flush_lazy_mmu_mode();
+       clear_thread_flag(TIF_LAZY_MMU);
+}
+
 #ifdef CONFIG_TRANSPARENT_HUGEPAGE
 #define __HAVE_ARCH_FLUSH_PMD_TLB_RANGE
 
         * Only if the new pte is valid and kernel, otherwise TLB maintenance
         * has the necessary barriers.
         */
-       if (pte_valid_not_user(pte)) {
-               dsb(ishst);
-               isb();
-       }
+       if (pte_valid_not_user(pte))
+               queue_pte_barriers();
 }
 
 static inline void __set_pte(pte_t *ptep, pte_t pte)
 
        WRITE_ONCE(*pmdp, pmd);
 
-       if (pmd_valid(pmd)) {
-               dsb(ishst);
-               isb();
-       }
+       if (pmd_valid(pmd))
+               queue_pte_barriers();
 }
 
 static inline void pmd_clear(pmd_t *pmdp)
 
        WRITE_ONCE(*pudp, pud);
 
-       if (pud_valid(pud)) {
-               dsb(ishst);
-               isb();
-       }
+       if (pud_valid(pud))
+               queue_pte_barriers();
 }
 
 static inline void pud_clear(pud_t *pudp)
        }
 
        WRITE_ONCE(*p4dp, p4d);
-       dsb(ishst);
-       isb();
+       queue_pte_barriers();
 }
 
 static inline void p4d_clear(p4d_t *p4dp)
        }
 
        WRITE_ONCE(*pgdp, pgd);
-       dsb(ishst);
-       isb();
+       queue_pte_barriers();
 }
 
 static inline void pgd_clear(pgd_t *pgdp)