#else
        int             switch_pending;
 #endif
-       unsigned int    vmalloc_seq;
+       atomic_t        vmalloc_seq;
        unsigned long   sigpage;
 #ifdef CONFIG_VDSO
        unsigned long   vdso;
 
 
 void __check_vmalloc_seq(struct mm_struct *mm);
 
+#ifdef CONFIG_MMU
+static inline void check_vmalloc_seq(struct mm_struct *mm)
+{
+       if (!IS_ENABLED(CONFIG_ARM_LPAE) &&
+           unlikely(atomic_read(&mm->context.vmalloc_seq) !=
+                    atomic_read(&init_mm.context.vmalloc_seq)))
+               __check_vmalloc_seq(mm);
+}
+#endif
+
 #ifdef CONFIG_CPU_HAS_ASID
 
 void check_and_switch_context(struct mm_struct *mm, struct task_struct *tsk);
 static inline void check_and_switch_context(struct mm_struct *mm,
                                            struct task_struct *tsk)
 {
-       if (unlikely(mm->context.vmalloc_seq != init_mm.context.vmalloc_seq))
-               __check_vmalloc_seq(mm);
+       check_vmalloc_seq(mm);
 
        if (irqs_disabled())
                /*
 #endif
 }
 
+#ifdef CONFIG_VMAP_STACK
+static inline void enter_lazy_tlb(struct mm_struct *mm, struct task_struct *tsk)
+{
+       if (mm != &init_mm)
+               check_vmalloc_seq(mm);
+}
+#define enter_lazy_tlb enter_lazy_tlb
+#endif
+
 #include <asm-generic/mmu_context.h>
 
 #endif
 
 #include <asm/pgtable-3level-types.h>
 #else
 #include <asm/pgtable-2level-types.h>
-#endif
-
 #ifdef CONFIG_VMAP_STACK
 #define ARCH_PAGE_TABLE_SYNC_MASK      PGTBL_PMD_MODIFIED
 #endif
+#endif
 
 #endif /* CONFIG_MMU */
 
 
        die("kernel stack overflow", regs, 0);
 }
 
+#ifndef CONFIG_ARM_LPAE
 /*
  * Normally, we rely on the logic in do_translation_fault() to update stale PMD
  * entries covering the vmalloc space in a task's page tables when it first
  * So we need to ensure that these PMD entries are up to date *before* the MM
  * switch. As we already have some logic in the MM switch path that takes care
  * of this, let's trigger it by bumping the counter every time the core vmalloc
- * code modifies a PMD entry in the vmalloc region.
+ * code modifies a PMD entry in the vmalloc region. Use release semantics on
+ * the store so that other CPUs observing the counter's new value are
+ * guaranteed to see the updated page table entries as well.
  */
 void arch_sync_kernel_mappings(unsigned long start, unsigned long end)
 {
-       if (start > VMALLOC_END || end < VMALLOC_START)
-               return;
-
-       /*
-        * This hooks into the core vmalloc code to receive notifications of
-        * any PMD level changes that have been made to the kernel page tables.
-        * This means it should only be triggered once for every MiB worth of
-        * vmalloc space, given that we don't support huge vmalloc/vmap on ARM,
-        * and that kernel PMD level table entries are rarely (if ever)
-        * updated.
-        *
-        * This means that the counter is going to max out at ~250 for the
-        * typical case. If it overflows, something entirely unexpected has
-        * occurred so let's throw a warning if that happens.
-        */
-       WARN_ON(++init_mm.context.vmalloc_seq == UINT_MAX);
+       if (start < VMALLOC_END && end > VMALLOC_START)
+               atomic_inc_return_release(&init_mm.context.vmalloc_seq);
 }
-
+#endif
 #endif
 
        unsigned int cpu = smp_processor_id();
        u64 asid;
 
-       if (unlikely(mm->context.vmalloc_seq != init_mm.context.vmalloc_seq))
-               __check_vmalloc_seq(mm);
+       check_vmalloc_seq(mm);
 
        /*
         * We cannot update the pgd and the ASID atomicly with classic
 
 
 void __check_vmalloc_seq(struct mm_struct *mm)
 {
-       unsigned int seq;
+       int seq;
 
        do {
-               seq = init_mm.context.vmalloc_seq;
+               seq = atomic_read(&init_mm.context.vmalloc_seq);
                memcpy(pgd_offset(mm, VMALLOC_START),
                       pgd_offset_k(VMALLOC_START),
                       sizeof(pgd_t) * (pgd_index(VMALLOC_END) -
                                        pgd_index(VMALLOC_START)));
-               mm->context.vmalloc_seq = seq;
-       } while (seq != init_mm.context.vmalloc_seq);
+               /*
+                * Use a store-release so that other CPUs that observe the
+                * counter's new value are guaranteed to see the results of the
+                * memcpy as well.
+                */
+               atomic_set_release(&mm->context.vmalloc_seq, seq);
+       } while (seq != atomic_read(&init_mm.context.vmalloc_seq));
 }
 
 #if !defined(CONFIG_SMP) && !defined(CONFIG_ARM_LPAE)
                         * Note: this is still racy on SMP machines.
                         */
                        pmd_clear(pmdp);
-                       init_mm.context.vmalloc_seq++;
+                       atomic_inc_return_release(&init_mm.context.vmalloc_seq);
 
                        /*
                         * Free the page table, if there was one.
         * Ensure that the active_mm is up to date - we want to
         * catch any use-after-iounmap cases.
         */
-       if (current->active_mm->context.vmalloc_seq != init_mm.context.vmalloc_seq)
-               __check_vmalloc_seq(current->active_mm);
+       check_vmalloc_seq(current->active_mm);
 
        flush_tlb_kernel_range(virt, end);
 }