*/
        bool is_lazy;
 
+       /*
+        * If set we changed the page tables in such a way that we
+        * needed an invalidation of all contexts (aka. PCIDs / ASIDs).
+        * This tells us to go invalidate all the non-loaded ctxs[]
+        * on the next context switch.
+        *
+        * The current ctx was kept up-to-date as it ran and does not
+        * need to be invalidated.
+        */
+       bool invalidate_other;
+
        /*
         * Access to this CR4 shadow and to H/W CR4 is protected by
         * disabling interrupts when modifying either one.
        return this_cpu_read(cpu_tlbstate.cr4);
 }
 
+/*
+ * Mark all other ASIDs as invalid, preserves the current.
+ */
+static inline void invalidate_other_asid(void)
+{
+       this_cpu_write(cpu_tlbstate.invalidate_other, true);
+}
+
 /*
  * Save some of cr4 feature set we're using (e.g.  Pentium 4MB
  * enable and PPro Global page enable), so that any CPU's that boot
                 */
                __flush_tlb();
        }
-
-       /*
-        * Note: if we somehow had PCID but not PGE, then this wouldn't work --
-        * we'd end up flushing kernel translations for the current ASID but
-        * we might fail to flush kernel translations for other cached ASIDs.
-        *
-        * To avoid this issue, we force PCID off if PGE is off.
-        */
 }
 
 /*
 {
        count_vm_tlb_event(NR_TLB_LOCAL_FLUSH_ONE);
        __flush_tlb_single(addr);
+
+       if (!static_cpu_has(X86_FEATURE_PTI))
+               return;
+
+       /*
+        * __flush_tlb_single() will have cleared the TLB entry for this ASID,
+        * but since kernel space is replicated across all, we must also
+        * invalidate all others.
+        */
+       invalidate_other_asid();
 }
 
 #define TLB_FLUSH_ALL  -1UL
 
  *     Implement flush IPI by CALL_FUNCTION_VECTOR, Alex Shi
  */
 
+/*
+ * We get here when we do something requiring a TLB invalidation
+ * but could not go invalidate all of the contexts.  We do the
+ * necessary invalidation by clearing out the 'ctx_id' which
+ * forces a TLB flush when the context is loaded.
+ */
+void clear_asid_other(void)
+{
+       u16 asid;
+
+       /*
+        * This is only expected to be set if we have disabled
+        * kernel _PAGE_GLOBAL pages.
+        */
+       if (!static_cpu_has(X86_FEATURE_PTI)) {
+               WARN_ON_ONCE(1);
+               return;
+       }
+
+       for (asid = 0; asid < TLB_NR_DYN_ASIDS; asid++) {
+               /* Do not need to flush the current asid */
+               if (asid == this_cpu_read(cpu_tlbstate.loaded_mm_asid))
+                       continue;
+               /*
+                * Make sure the next time we go to switch to
+                * this asid, we do a flush:
+                */
+               this_cpu_write(cpu_tlbstate.ctxs[asid].ctx_id, 0);
+       }
+       this_cpu_write(cpu_tlbstate.invalidate_other, false);
+}
+
 atomic64_t last_mm_ctx_id = ATOMIC64_INIT(1);
 
 
                return;
        }
 
+       if (this_cpu_read(cpu_tlbstate.invalidate_other))
+               clear_asid_other();
+
        for (asid = 0; asid < TLB_NR_DYN_ASIDS; asid++) {
                if (this_cpu_read(cpu_tlbstate.ctxs[asid].ctx_id) !=
                    next->context.ctx_id)