PCID and improved laziness

Nadav noticed the reentrancy issue in a different context, but neither of us realized that it caused a problem yet.

Cc Nadav Amit <[email protected]>

Cc Dave Hansen <[email protected]>

Reported-by "Levin, Alexander (Sasha Levin)" <[email protected]>

Fixes 3d28ebceaffa ("x86/mm: Rework lazy TLB to track the actual loaded mm")

Signed-off-by Andy Lutomirski <[email protected]>

It was historically possible to have two concurrent TLB flushes targetting the same CPU: one initiated locally and one initiated remotely. This can now cause an OOPS in leave_mm() at arch/x86/mm/tlb.c:47:

        if (this_cpu_read(cpu_tlbstate.state) == TLBSTATE_OK)
                BUG();

with this call trace:

 flush_tlb_func_local arch/x86/mm/tlb.c:239 [inline]
 flush_tlb_mm_range+0x26d/0x370 arch/x86/mm/tlb.c:317

Without reentrancy, this OOPS is impossible: leave_mm() is only called if we're not in TLBSTATE_OK, but then we're unexpectedly in TLBSTATE_OK in leave_mm().

This can be caused by flush_tlb_func_remote() happening between the two checks and calling leave_mm(), resulting in two consecutive leave_mm() calls on the same CPU with no intervening switch_mm() calls.

We never saw this OOPS before because the old leave_mm() implementation didn't put us back in TLBSTATE_OK, so the assertion didn't fire.

Nadav noticed the reentrancy issue in a different context, but neither of us realized that it caused a problem yet.

Cc Nadav Amit <[email protected]>

Cc Dave Hansen <[email protected]>

Reported-by "Levin, Alexander (Sasha Levin)" <[email protected]>

Fixes 3d28ebceaffa ("x86/mm: Rework lazy TLB to track the actual loaded mm")

Signed-off-by Andy Lutomirski <[email protected]>

---
 arch/x86/mm/tlb.c | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

arch/x86/mm/tlb.c +3 -0

index 2a5e851f2035..f06239c6919f 100644
@@ -208,6 +208,9 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
 static void flush_tlb_func_common(const struct flush_tlb_info *f,
 				  bool local, enum tlb_flush_reason reason)
 {
+	/* This code cannot presently handle being reentered. */
+	VM_WARN_ON(!irqs_disabled());
+
 	if (this_cpu_read(cpu_tlbstate.state) != TLBSTATE_OK) {
 		leave_mm(smp_processor_id());
 		return;

Quoted 8 lines · Andy Lutomirski

It was historically possible to have two concurrent TLB flushes targetting the same CPU: one initiated locally and one initiated remotely. This can now cause an OOPS in leave_mm() at arch/x86/mm/tlb.c:47:

        if (this_cpu_read(cpu_tlbstate.state) == TLBSTATE_OK)
                BUG();

Quoted 3 lines

with this call trace:

 flush_tlb_func_local arch/x86/mm/tlb.c:239 [inline]
 flush_tlb_mm_range+0x26d/0x370 arch/x86/mm/tlb.c:317

These line numbers would most likely mean nothing soon. I think you should rather explain why the bug can happen so that future lookers at that code can find the spot...

I'm assuming this is going away in a future patch, as disabling IRQs around a TLB flush is kinda expensive. I guess I'll see if I continue reading...

index 2a5e851f2035..f06239c6919f 100644

arch/x86/mm/tlb.c +3 -0

@@ -208,6 +208,9 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
 static void flush_tlb_func_common(const struct flush_tlb_info *f,
 				  bool local, enum tlb_flush_reason reason)
 {
+	/* This code cannot presently handle being reentered. */
+	VM_WARN_ON(!irqs_disabled());
+
 	if (this_cpu_read(cpu_tlbstate.state) != TLBSTATE_OK) {
 		leave_mm(smp_processor_id());
 		return;

arch/x86/mm/tlb.c +4 -0

+		local_irq_disable();
 		flush_tlb_func_local(&info, TLB_LOCAL_MM_SHOOTDOWN);
+		local_irq_enable();
+	}
+
 	if (cpumask_any_but(mm_cpumask(mm), cpu) < nr_cpu_ids)
 		flush_tlb_others(mm_cpumask(mm), &info);
 	put_cpu();

arch/x86/mm/tlb.c +4 -0

+		local_irq_disable();
 		flush_tlb_func_local(&info, TLB_LOCAL_SHOOTDOWN);
+		local_irq_enable();
+	}
+
 	if (cpumask_any_but(&batch->cpumask, cpu) < nr_cpu_ids)
 		flush_tlb_others(&batch->cpumask, &info);
 	cpumask_clear(&batch->cpumask);

Signature · 1 line

2.9.4

Quoted 14 lines · Andy Lutomirski

index 2a5e851f2035..f06239c6919f 100644

arch/x86/mm/tlb.c +3 -0

@@ -208,6 +208,9 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
 static void flush_tlb_func_common(const struct flush_tlb_info *f,
                                bool local, enum tlb_flush_reason reason)
 {
+     /* This code cannot presently handle being reentered. */
+     VM_WARN_ON(!irqs_disabled());
+
      if (this_cpu_read(cpu_tlbstate.state) != TLBSTATE_OK) {
              leave_mm(smp_processor_id());
              return;

arch/x86/mm/tlb.c void flush_tlb_mm_range(struct mm_struct *mm, unsigned long start, · line 319

Quoted 5 lines

arch/x86/mm/tlb.c +1 -1

@@ -313,8 +316,12 @@ void flush_tlb_mm_range(struct mm_struct *mm, unsigned long start,
              info.end = TLB_FLUSH_ALL;
      }
 
-     if (mm == this_cpu_read(cpu_tlbstate.loaded_mm))
+     if (mm == this_cpu_read(cpu_tlbstate.loaded_mm)) {

Perhaps you want to add:

        VM_WARN_ON(irqs_disabled());

here

Quoted 10 lines

arch/x86/mm/tlb.c +4 -0

+             local_irq_disable();
              flush_tlb_func_local(&info, TLB_LOCAL_MM_SHOOTDOWN);
+             local_irq_enable();
+     }
+
      if (cpumask_any_but(mm_cpumask(mm), cpu) < nr_cpu_ids)
              flush_tlb_others(mm_cpumask(mm), &info);
      put_cpu();

@@ -370,8 +377,12 @@ void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch)
 

arch/x86/mm/tlb.c void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch) · line 380

Quoted 4 lines

arch/x86/mm/tlb.c +1 -1

      int cpu = get_cpu();
 
-     if (cpumask_test_cpu(cpu, &batch->cpumask))
+     if (cpumask_test_cpu(cpu, &batch->cpumask)) {

and here?

Quoted 38 lines · Nadav Amit

Quoted 14 lines · Andy Lutomirski

index 2a5e851f2035..f06239c6919f 100644

arch/x86/mm/tlb.c +3 -0

@@ -208,6 +208,9 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
 static void flush_tlb_func_common(const struct flush_tlb_info *f,
                                bool local, enum tlb_flush_reason reason)
 {
+     /* This code cannot presently handle being reentered. */
+     VM_WARN_ON(!irqs_disabled());
+
      if (this_cpu_read(cpu_tlbstate.state) != TLBSTATE_OK) {
              leave_mm(smp_processor_id());
              return;

arch/x86/mm/tlb.c void flush_tlb_mm_range(struct mm_struct *mm, unsigned long start, · line 319

Quoted 5 lines

arch/x86/mm/tlb.c +1 -1

@@ -313,8 +316,12 @@ void flush_tlb_mm_range(struct mm_struct *mm, unsigned long start,
              info.end = TLB_FLUSH_ALL;
      }
 
-     if (mm == this_cpu_read(cpu_tlbstate.loaded_mm))
+     if (mm == this_cpu_read(cpu_tlbstate.loaded_mm)) {

Perhaps you want to add:

        VM_WARN_ON(irqs_disabled());

here

Quoted 10 lines

arch/x86/mm/tlb.c +4 -0

+             local_irq_disable();
              flush_tlb_func_local(&info, TLB_LOCAL_MM_SHOOTDOWN);
+             local_irq_enable();
+     }
+
      if (cpumask_any_but(mm_cpumask(mm), cpu) < nr_cpu_ids)
              flush_tlb_others(mm_cpumask(mm), &info);
      put_cpu();

@@ -370,8 +377,12 @@ void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch)
 

Quoted 7 lines

arch/x86/mm/tlb.c void arch_tlbbatch_flush(struct arch_tlbflush_unmap_batch *batch) · line 380

Quoted 4 lines

arch/x86/mm/tlb.c +1 -1

      int cpu = get_cpu();
 
-     if (cpumask_test_cpu(cpu, &batch->cpumask))
+     if (cpumask_test_cpu(cpu, &batch->cpumask)) {

and here?

Will do.

What I really want is lockdep_assert_irqs_disabled() or, even better, for this to be implicit when calling local_irq_disable(). Ingo?

When we redo lazy mode to stop flush IPIs without switching to init_mm, though, the current logic would become incorrect: it will be possible to have real_prev == next but nonetheless have a stale LDT descriptor.

Simplify the code to update LDTR if either the previous or the next mm has an LDT, i.e. effectively restore the historical logic.. While we're at it, clean up the code by moving all the ifdeffery to a header where it belongs.

Acked-by Rik van Riel <[email protected]>

Signed-off-by Andy Lutomirski <[email protected]>

Originally, Linux reloaded the LDT whenever the prev mm or the next mm had an LDT. It was changed in 0bbed3beb4f2 ("[PATCH] Thread-Local Storage (TLS) support") (from the historical tree) like this:

(quoted diff) +2 -2

-		/* load_LDT, if either the previous or next thread
-		 * has a non-default LDT.
+		/*
+		 * load the LDT, if the LDT is different:

		 */

(quoted diff) +1 -1

-		if (next->context.size+prev->context.size)
+		if (unlikely(prev->context.ldt != next->context.ldt))

			load_LDT(&next->context);

The current code is unlikely to avoid any LDT reloads, since different mms won't share an LDT.

When we redo lazy mode to stop flush IPIs without switching to init_mm, though, the current logic would become incorrect: it will be possible to have real_prev == next but nonetheless have a stale LDT descriptor.

Simplify the code to update LDTR if either the previous or the next mm has an LDT, i.e. effectively restore the historical logic.. While we're at it, clean up the code by moving all the ifdeffery to a header where it belongs.

Acked-by Rik van Riel <[email protected]>

Signed-off-by Andy Lutomirski <[email protected]>

---
 arch/x86/include/asm/mmu_context.h | 26 ++++++++++++++++++++++++++
 arch/x86/mm/tlb.c                  | 20 ++------------------
 2 files changed, 28 insertions(+), 18 deletions(-)

The only call site also calls idle_task_exit(), and idle_task_exit() puts us into a clean state by explicitly switching to init_mm.

Reviewed-by Rik van Riel <[email protected]>

Signed-off-by Andy Lutomirski <[email protected]>

The only call site also calls idle_task_exit(), and idle_task_exit() puts us into a clean state by explicitly switching to init_mm.

Reviewed-by Rik van Riel <[email protected]>

Signed-off-by Andy Lutomirski <[email protected]>

---
 arch/x86/include/asm/tlbflush.h | 8 --------
 arch/x86/kernel/smpboot.c       | 1 -
 2 files changed, 9 deletions(-)

This adds two new variables to mmu_context_t: ctx_id and tlb_gen. ctx_id uniquely identifies the mm_struct and will never be reused. For a given mm_struct (and hence ctx_id), tlb_gen is a monotonic count of the number of times that a TLB flush has been requested. The pair (ctx_id, tlb_gen) can be used as an identifier for TLB flush actions and will be used in subsequent patches to reliably determine whether all needed TLB flushes have occurred on a given CPU.

This patch is split out for ease of review. By itself, it has no real effect other than creating and updating the new variables.

Signed-off-by Andy Lutomirski <[email protected]>

- * The x86 doesn't have a mmu context, but
- * we put the segment information here.
+ * x86 has arch-specific MMU state beyond what lives in mm_struct. */
typedef struct {

arch/x86/include/asm/mmu.h +16 -0

+	/*
+	 * ctx_id uniquely identifies this mm_struct.  A ctx_id will never
+	 * be reused, and zero is not a valid ctx_id.
+	 */
+	u64 ctx_id;
+
+	/*
+	 * Any code that needs to do any sort of TLB flushing for this
+	 * mm will first make its changes to the page tables, then
+	 * increment tlb_gen, then flush.  This lets the low-level
+	 * flushing code keep track of what needs flushing.
+	 *
+	 * This is not used on Xen PV.
+	 */
+	atomic64_t tlb_gen;
+
 #ifdef CONFIG_MODIFY_LDT_SYSCALL
 	struct ldt_struct *ldt;
 #endif

arch/x86/include/asm/mmu.h +2 -0

+	}
+
 void leave_mm(int cpu);
 
 #endif /* _ASM_X86_MMU_H */

arch/x86/include/asm/mmu_context.h +0 -0

index ecfcb6643c9b..e5295d485899 100644

arch/x86/include/asm/mmu_context.h +4 -0

+
 static inline int init_new_context(struct task_struct *tsk,
 				   struct mm_struct *mm)
 {
+	mm->context.ctx_id = atomic64_inc_return(&last_mm_ctx_id);
+	atomic64_set(&mm->context.tlb_gen, 0);
+
 	#ifdef CONFIG_X86_INTEL_MEMORY_PROTECTION_KEYS
 	if (cpu_feature_enabled(X86_FEATURE_OSPKE)) {
 		/* pkey 0 is the default and always allocated */

arch/x86/include/asm/tlbflush.h +0 -0

index 50ea3482e1d1..1eb946c0507e 100644

arch/x86/include/asm/tlbflush.h +10 -0

+	 * their read of mm_cpumask after their writes to the paging
+	 * structures.
+	 */
+	smp_mb__before_atomic();
+	new_tlb_gen = atomic64_inc_return(&mm->context.tlb_gen);
+	smp_mb__after_atomic();
+
+	return new_tlb_gen;
+}
+
 #ifdef CONFIG_PARAVIRT
 #include <asm/paravirt.h>
 #else

arch/x86/include/asm/mmu.h typedef struct { · line 58

Quoted 3 lines · Andy Lutomirski

arch/x86/include/asm/mmu.h +3 -0

+#define INIT_MM_CONTEXT(mm)                                          \
+     .context = {                                                    \
+             .ctx_id = 1,                                            \

So ctx_id of 0 is invalid?

Let's state that explicitly. We could even use it to sanity-check mms or whatever.

Quoted 14 lines

arch/x86/include/asm/mmu.h +2 -0

+     }
+
 void leave_mm(int cpu);
 
 #endif /* _ASM_X86_MMU_H */

arch/x86/include/asm/mmu_context.h +0 -0

index ecfcb6643c9b..e5295d485899 100644
@@ -129,9 +129,14 @@ static inline void enter_lazy_tlb(struct mm_struct *mm, struct task_struct *tsk)
              this_cpu_write(cpu_tlbstate.state, TLBSTATE_LAZY);
 }
 

arch/x86/include/asm/mmu_context.h static inline void enter_lazy_tlb(struct mm_struct *mm, struct task_struct *tsk) · line 132

Quoted 1 line

arch/x86/include/asm/mmu_context.h +1 -0

+extern atomic64_t last_mm_ctx_id;

I think we prefer externs/variable defines at the beginning of the file, not intermixed with functions.

arch/x86/include/asm/mmu_context.h static inline void enter_lazy_tlb(struct mm_struct *mm, struct task_struct *tsk) · line 133

Quoted 2 lines

arch/x86/include/asm/mmu_context.h +1 -0

+static inline u64 bump_mm_tlb_gen(struct mm_struct *mm)

inc_mm_tlb_gen() I guess. git grep says like "inc" more :-)

arch/x86/include/asm/mmu_context.h static inline void enter_lazy_tlb(struct mm_struct *mm, struct task_struct *tsk) · line 134

Quoted 1 line

arch/x86/include/asm/mmu_context.h +1 -0

+      * that synchronizes with switch_mm: callers are required to order

Please end function names with parentheses.

It's stated explicitly in the comment where it's declared in the same file.

- * The x86 doesn't have a mmu context, but
- * we put the segment information here.
+ * x86 has arch-specific MMU state beyond what lives in mm_struct. */
typedef struct {

arch/x86/include/asm/mmu.h +5 -0

+	/*
+	 * ctx_id uniquely identifies this mm_struct.  A ctx_id will never
+	 * be reused, and zero is not a valid ctx_id.
+	 */
+	u64 ctx_id;

Quoted 10 lines · Andy Lutomirski

- * The x86 doesn't have a mmu context, but
- * we put the segment information here.
+ * x86 has arch-specific MMU state beyond what lives in mm_struct. */
typedef struct {

arch/x86/include/asm/mmu.h +5 -0

+     /*
+      * ctx_id uniquely identifies this mm_struct.  A ctx_id will never
+      * be reused, and zero is not a valid ctx_id.
+      */
+     u64 ctx_id;

arch/x86/include/asm/mmu.h line 26

Quoted 10 lines

arch/x86/include/asm/mmu.h +10 -0

+
+     /*
+      * Any code that needs to do any sort of TLB flushing for this
+      * mm will first make its changes to the page tables, then
+      * increment tlb_gen, then flush.  This lets the low-level
+      * flushing code keep track of what needs flushing.
+      *
+      * This is not used on Xen PV.
+      */
+     atomic64_t tlb_gen;

Btw, can this just be a 4-byte int instead? I.e., simply atomic_t. I mean, it should be enough for all the TLB generations in flight, no?

arch/x86/include/asm/tlbflush.h +6 -0

+	 *   TLB_FLUSH_ALL, and new_tlb_gen will be the tlb_gen to which the
+	 *   IPI sender is trying to catch us up.
+	 *
+	 * - Partially flush a single mm.  flush_mm will be set, flush_start
+	 *   and flush_end will indicate the range, and new_tlb_gen will be
+	 *   set such that the changes between generation new_tlb_gen-1 and

arch/x86/mm/tlb.c +1 -2

+		 */
 		leave_mm(smp_processor_id());
 		return;
 	}
 
-	if (f->end == TLB_FLUSH_ALL) {
-		local_flush_tlb();

arch/x86/include/asm/tlbflush.h static inline void __flush_tlb_one(unsigned long addr) · line 274

Quoted 5 lines · Andy Lutomirski

struct flush_tlb_info {

arch/x86/include/asm/tlbflush.h +4 -0

+     /*
+      * We support several kinds of flushes.
+      *
+      * - Fully flush a single mm.  flush_mm will be set, flush_end will be

flush_mm is the *mm member in the struct, right? You might rename that as a preparatory step so comments and implementation match.

Quoted 7 lines

arch/x86/include/asm/tlbflush.h +6 -0

+      *   TLB_FLUSH_ALL, and new_tlb_gen will be the tlb_gen to which the
+      *   IPI sender is trying to catch us up.
+      *
+      * - Partially flush a single mm.  flush_mm will be set, flush_start
+      *   and flush_end will indicate the range, and new_tlb_gen will be
+      *   set such that the changes between generation new_tlb_gen-1 and

arch/x86/include/asm/tlbflush.h static inline void __flush_tlb_one(unsigned long addr) · line 290

Quoted 10 lines

arch/x86/include/asm/tlbflush.h +7 -0

+      *   new_tlb_gen are entirely contained in the indicated range.
+      *
+      * - Fully flush all mms whose tlb_gens have been updated.  flush_mm
+      *   will be NULL, flush_end will be TLB_FLUSH_ALL, and new_tlb_gen
+      *   will be zero.
+      */
      struct mm_struct *mm;
      unsigned long start;
      unsigned long end;
+     u64 new_tlb_gen;

Nit. While at it could you please make that struct tabular aligned as we usually do in x86?

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 208

Quoted 13 lines

static void flush_tlb_func_common(const struct flush_tlb_info *f,
bool local, enum tlb_flush_reason reason)
{

arch/x86/mm/tlb.c +9 -0

+     struct mm_struct *loaded_mm = this_cpu_read(cpu_tlbstate.loaded_mm);
+
+     /*
+      * Our memory ordering requirement is that any TLB fills that
+      * happen after we flush the TLB are ordered after we read
+      * active_mm's tlb_gen.  We don't need any explicit barrier
+      * because all x86 flush operations are serializing and the
+      * atomic64_read operation won't be reordered by the compiler.
+      */

Can you please move the comment above the loaded_mm assignment?

Quoted 7 lines

arch/x86/mm/tlb.c +3 -0

+     u64 mm_tlb_gen = atomic64_read(&loaded_mm->context.tlb_gen);
+     u64 local_tlb_gen = this_cpu_read(cpu_tlbstate.ctxs[0].tlb_gen);
+
      /* This code cannot presently handle being reentered. */
      VM_WARN_ON(!irqs_disabled());

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 221

Quoted 7 lines

arch/x86/mm/tlb.c +6 -0

+     VM_WARN_ON(this_cpu_read(cpu_tlbstate.ctxs[0].ctx_id) !=
+                loaded_mm->context.ctx_id);
+
      if (this_cpu_read(cpu_tlbstate.state) != TLBSTATE_OK) {
+             /*
+              * leave_mm() is adequate to handle any type of flush, and
+              * we would prefer not to receive further IPIs.

While I know what you mean, it might be useful to have a more elaborate explanation why this prevents new IPIs.

Quoted 8 lines

arch/x86/mm/tlb.c +1 -2

+              */
              leave_mm(smp_processor_id());
              return;
      }
 
-     if (f->end == TLB_FLUSH_ALL) {
-             local_flush_tlb();

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 232

Quoted 10 lines

arch/x86/mm/tlb.c +6 -4

-             if (local)
-                     count_vm_tlb_event(NR_TLB_LOCAL_FLUSH_ALL);
-             trace_tlb_flush(reason, TLB_FLUSH_ALL);
-     } else {
+     if (local_tlb_gen == mm_tlb_gen) {
+             /*
+              * There's nothing to do: we're already up to date.  This can
+              * happen if two concurrent flushes happen -- the first IPI to
+              * be handled can catch us all the way up, leaving no work for
+              * the second IPI to be handled.

That not restricted to IPIs, right? A local flush / IPI combo can do that as well.

Other than those nits;

Reviewed-by Thomas Gleixner <[email protected]>

There are two kernel features that would benefit from tracking how up-to-date each CPU's TLB is in the case where IPIs aren't keeping it up to date in real time:

- Lazy mm switching currently works by switching to init_mm when it would otherwise flush. This is wasteful: there isn't fundamentally any need to update CR3 at all when going lazy or when returning from lazy mode, nor is there any need to receive flush IPIs at all. Instead, we should just stop trying to keep the TLB coherent when we go lazy and, when unlazying, check whether we missed any flushes.

- PCID will let us keep recent user contexts alive in the TLB. If we start doing this, we need a way to decide whether those contexts are up to date.

On some paravirt systems, remote TLBs can be flushed without IPIs. This won't update the target CPUs' tlb_gens, which may cause unnecessary local flushes later on. We can address this if it becomes a problem by carefully updating the target CPU's tlb_gen directly.

By itself, this patch is a very minor optimization that avoids unnecessary flushes when multiple TLB flushes targetting the same CPU race.

Signed-off-by Andy Lutomirski <[email protected]>

---
 arch/x86/include/asm/tlbflush.h | 37 +++++++++++++++++++
 arch/x86/mm/tlb.c               | 79 +++++++++++++++++++++++++++++++++++++----
 2 files changed, 109 insertions(+), 7 deletions(-)

arch/x86/mm/tlb.c +19 -0

 	WARN_ON_ONCE(cpumask_test_cpu(cpu, mm_cpumask(next)));
 	cpumask_set_cpu(cpu, mm_cpumask(next));

@@ -194,20 +197,73 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
 static void flush_tlb_func_common(const struct flush_tlb_info *f,
 				  bool local, enum tlb_flush_reason reason)
 {
+	struct mm_struct *loaded_mm = this_cpu_read(cpu_tlbstate.loaded_mm);
+
+	/*
+	 * Our memory ordering requirement is that any TLB fills that
+	 * happen after we flush the TLB are ordered after we read
+	 * active_mm's tlb_gen.  We don't need any explicit barrier
+	 * because all x86 flush operations are serializing and the
+	 * atomic64_read operation won't be reordered by the compiler.
+	 */
+	u64 mm_tlb_gen = atomic64_read(&loaded_mm->context.tlb_gen);
+	u64 local_tlb_gen = this_cpu_read(cpu_tlbstate.ctxs[0].tlb_gen);
+
 	/* This code cannot presently handle being reentered. */
 	VM_WARN_ON(!irqs_disabled());
 
+	VM_WARN_ON(this_cpu_read(cpu_tlbstate.ctxs[0].ctx_id) !=
+		   loaded_mm->context.ctx_id);
+
 	if (this_cpu_read(cpu_tlbstate.state) != TLBSTATE_OK) {
+		/*
+		 * leave_mm() is adequate to handle any type of flush, and
+		 * we would prefer not to receive further IPIs.
+		 */
 		leave_mm(smp_processor_id());

arch/x86/mm/tlb.c +28 -0

+		/*
+		 * There's nothing to do: we're already up to date.  This can
+		 * happen if two concurrent flushes happen -- the first IPI to
+		 * be handled can catch us all the way up, leaving no work for
+		 * the second IPI to be handled.
+		 */
+		return;
+	}
 
 
+
+	WARN_ON_ONCE(local_tlb_gen > mm_tlb_gen);
+	WARN_ON_ONCE(f->new_tlb_gen > mm_tlb_gen);
+
+	/*
+	 * If we get to this point, we know that our TLB is out of date.
+	 * This does not strictly imply that we need to flush (it's
+	 * possible that f->new_tlb_gen <= local_tlb_gen), but we're
+	 * going to need to flush in the very near future, so we might
+	 * as well get it over with.
+	 *
+	 * The only question is whether to do a full or partial flush.
+	 *
+	 * A partial TLB flush is safe and worthwhile if two conditions are
+	 * met:
+	 *
+	 * 1. We wouldn't be skipping a tlb_gen.  If the requester bumped
+	 *    the mm's tlb_gen from p to p+1, a partial flush is only correct
+	 *    if we would be bumping the local CPU's tlb_gen from p to p+1 as
+	 *    well.

arch/x86/mm/tlb.c +1 -0

+		/* Partial flush */
 		unsigned long addr;

unsigned long nr_pages = (f->end - f->start) >> PAGE_SHIFT;

arch/x86/mm/tlb.c +9 -0

 		addr = f->start;

@@ -218,7 +274,16 @@ static void flush_tlb_func_common(const struct flush_tlb_info *f,
 		if (local)
 			count_vm_tlb_events(NR_TLB_LOCAL_FLUSH_ONE, nr_pages);
 		trace_tlb_flush(reason, nr_pages);
+	} else {
+		/* Full flush. */
+		local_flush_tlb();
+		if (local)
+			count_vm_tlb_event(NR_TLB_LOCAL_FLUSH_ALL);
+		trace_tlb_flush(reason, TLB_FLUSH_ALL);
 	}
+
+	/* Both paths above update our state to mm_tlb_gen. */
+	this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen, mm_tlb_gen);
 }
 
 static void flush_tlb_func_local(void *info, enum tlb_flush_reason reason)

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 110

Quoted 3 lines · Andy Lutomirski

arch/x86/mm/tlb.c +3 -0

+     this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id, next->context.ctx_id);
+     this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
+                    atomic64_read(&next->context.tlb_gen));

Just let it stick out:

        this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id,  next->context.ctx_id);
        this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen, atomic64_read(&next->context.tlb_gen));

Should be a bit better readable this way.

Quoted 1 line

+ if (local_tlb_gen == mm_tlb_gen) {

        if (unlikely(...

maybe?

Sounds to me like the concurrent flushes case would be the uncommon one...

Quoted 20 lines

arch/x86/mm/tlb.c +20 -0

+
+     WARN_ON_ONCE(local_tlb_gen > mm_tlb_gen);
+     WARN_ON_ONCE(f->new_tlb_gen > mm_tlb_gen);
+
+     /*
+      * If we get to this point, we know that our TLB is out of date.
+      * This does not strictly imply that we need to flush (it's
+      * possible that f->new_tlb_gen <= local_tlb_gen), but we're
+      * going to need to flush in the very near future, so we might
+      * as well get it over with.
+      *
+      * The only question is whether to do a full or partial flush.
+      *
+      * A partial TLB flush is safe and worthwhile if two conditions are
+      * met:
+      *
+      * 1. We wouldn't be skipping a tlb_gen.  If the requester bumped
+      *    the mm's tlb_gen from p to p+1, a partial flush is only correct
+      *    if we would be bumping the local CPU's tlb_gen from p to p+1 as
+      *    well.

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 265

Quoted 10 lines

arch/x86/mm/tlb.c +10 -0

+      *
+      * 2. If there are no more flushes on their way.  Partial TLB
+      *    flushes are not all that much cheaper than full TLB
+      *    flushes, so it seems unlikely that it would be a
+      *    performance win to do a partial flush if that won't bring
+      *    our TLB fully up to date.
+      */
+     if (f->end != TLB_FLUSH_ALL &&
+         f->new_tlb_gen == local_tlb_gen + 1 &&
+         f->new_tlb_gen == mm_tlb_gen) {

I'm certainly still missing something here:

We have f->new_tlb_gen and mm_tlb_gen to control the flushing, i.e., we do once

        bump_mm_tlb_gen(mm);

and once

        info.new_tlb_gen = bump_mm_tlb_gen(mm);

and in both cases, the bumping is done on mm->context.tlb_gen.

So why isn't that enough to do the flushing and we have to consult info.new_tlb_gen too?

Quoted 4 lines

+ /* Partial flush */

              unsigned long addr;
              unsigned long nr_pages = (f->end - f->start) >> PAGE_SHIFT;

<---- newline here.

Quoted 13 lines

I'm certainly still missing something here:

We have f->new_tlb_gen and mm_tlb_gen to control the flushing, i.e., we do once

        bump_mm_tlb_gen(mm);

and once

        info.new_tlb_gen = bump_mm_tlb_gen(mm);

and in both cases, the bumping is done on mm->context.tlb_gen.

Quoted 2 lines

So why isn't that enough to do the flushing and we have to consult info.new_tlb_gen too?

The issue is a possible race. Suppose we start at tlb_gen == 1 and then two concurrent flushes happen. The first flush is a full flush and sets tlb_gen to 2. The second is a partial flush and sets tlb_gen to 3. If the second flush gets propagated to a given CPU first and it

Quoted 13 lines · Andy Lutomirski

Quoted 13 lines

I'm certainly still missing something here:

We have f->new_tlb_gen and mm_tlb_gen to control the flushing, i.e., we do once

        bump_mm_tlb_gen(mm);

and once

        info.new_tlb_gen = bump_mm_tlb_gen(mm);

and in both cases, the bumping is done on mm->context.tlb_gen.

Quoted 7 lines

Quoted 2 lines

So why isn't that enough to do the flushing and we have to consult info.new_tlb_gen too?

The issue is a possible race. Suppose we start at tlb_gen == 1 and then two concurrent flushes happen. The first flush is a full flush and sets tlb_gen to 2. The second is a partial flush and sets tlb_gen to 3. If the second flush gets propagated to a given CPU first and it

Maybe I'm still missing something, which is likely...

but if the second flush gets propagated to the CPU first, the CPU will have local tlb_gen 1 and thus enforce a full flush anyway because we will go 1 -> 3 on that particular CPU. Or?

Quoted 24 lines · Borislav Petkov

Quoted 13 lines · Andy Lutomirski

Quoted 13 lines

I'm certainly still missing something here:

We have f->new_tlb_gen and mm_tlb_gen to control the flushing, i.e., we do once

        bump_mm_tlb_gen(mm);

and once

        info.new_tlb_gen = bump_mm_tlb_gen(mm);

and in both cases, the bumping is done on mm->context.tlb_gen.

Quoted 7 lines

Quoted 2 lines

So why isn't that enough to do the flushing and we have to consult info.new_tlb_gen too?

The issue is a possible race. Suppose we start at tlb_gen == 1 and then two concurrent flushes happen. The first flush is a full flush and sets tlb_gen to 2. The second is a partial flush and sets tlb_gen to 3. If the second flush gets propagated to a given CPU first and it

Maybe I'm still missing something, which is likely...

Quoted 4 lines

but if the second flush gets propagated to the CPU first, the CPU will have local tlb_gen 1 and thus enforce a full flush anyway because we will go 1 -> 3 on that particular CPU. Or?

Yes, exactly. Which means I'm probably just misunderstanding your original question. Can you re-ask it?

Quoted 25 lines · Andy Lutomirski

Quoted 24 lines · Borislav Petkov

Quoted 13 lines · Andy Lutomirski

Quoted 13 lines

I'm certainly still missing something here:

We have f->new_tlb_gen and mm_tlb_gen to control the flushing, i.e., we do once

        bump_mm_tlb_gen(mm);

and once

        info.new_tlb_gen = bump_mm_tlb_gen(mm);

and in both cases, the bumping is done on mm->context.tlb_gen.

Quoted 7 lines

Quoted 2 lines

So why isn't that enough to do the flushing and we have to consult info.new_tlb_gen too?

The issue is a possible race. Suppose we start at tlb_gen == 1 and then two concurrent flushes happen. The first flush is a full flush and sets tlb_gen to 2. The second is a partial flush and sets tlb_gen to 3. If the second flush gets propagated to a given CPU first and it

Maybe I'm still missing something, which is likely...

Quoted 7 lines

Quoted 4 lines

but if the second flush gets propagated to the CPU first, the CPU will have local tlb_gen 1 and thus enforce a full flush anyway because we will go 1 -> 3 on that particular CPU. Or?

Yes, exactly. Which means I'm probably just misunderstanding your original question. Can you re-ask it?

Ah, simple: we control the flushing with info.new_tlb_gen and mm->context.tlb_gen. I.e., this check:

        if (f->end != TLB_FLUSH_ALL &&
            f->new_tlb_gen == local_tlb_gen + 1 &&
            f->new_tlb_gen == mm_tlb_gen) {

why can't we write:

        if (f->end != TLB_FLUSH_ALL &&
            mm_tlb_gen == local_tlb_gen + 1)

?

If mm_tlb_gen is + 2, then we'll do a full flush, if it is + 1, then partial.

If the second flush, as you say is a partial one and still gets propagated first, the check will force a full flush anyway.

When the first flush propagates after the second, we'll ignore it because local_tlb_gen has advanced adready due to the second flush.

As a matter of fact, we could simplify the logic: if local_tlb_gen is only mm_tlb_gen - 1, then do the requested flush type.

Quoted 8 lines

Ah, simple: we control the flushing with info.new_tlb_gen and mm->context.tlb_gen. I.e., this check:

        if (f->end != TLB_FLUSH_ALL &&
            f->new_tlb_gen == local_tlb_gen + 1 &&
            f->new_tlb_gen == mm_tlb_gen) {

Quoted 6 lines

why can't we write:

        if (f->end != TLB_FLUSH_ALL &&
            mm_tlb_gen == local_tlb_gen + 1)

?

Ah, I thought you were asking about why I needed mm_tlb_gen == local_tlb_gen + 1. This is just an optimization, or at least I hope it is. The idea is that, if we know that another flush is coming, it seems likely that it would be faster to do a full flush and increase local_tlb_gen all the way to mm_tlb_gen rather than doing a partial flush, increasing local_tlb_gen to something less than mm_tlb_gen, and needing to flush again very soon.

Hmm. I'd be nervous that there are more subtle races if we do this. For example, suppose that a partial flush increments tlb_gen from 1 to 2 and a full flush increments tlb_gen from 2 to 3. Meanwhile, the CPU is busy switching back and forth between mms, so the partial flush sees the cpu set in mm_cpumask but the full flush doesn't see the cpu set in mm_cpumask.

The flush IPI hits after a switch_mm_irqs_off() call notices the change from 1 to 2. switch_mm_irqs_off() will do a full flush and increment the local tlb_gen to 2, and the IPI handler for the partial flush will see local_tlb_gen == mm_tlb_gen - 1 (because local_tlb_gen == 2 and mm_tlb_gen == 3) and do a partial flush.

The problem here is that it's not obvious to me that this actually ends up flushing everything that's needed. Maybe all the memory ordering gets this right, but I can imagine scenarios in which switch_mm_irqs_off() does its flush early enough that the TLB picks up an entry that was supposed to get zapped by the full flush.

IOW it *might* be valid, but I think it would need very careful review and documentation.

Quoted 15 lines · Andy Lutomirski

Quoted 8 lines

Ah, simple: we control the flushing with info.new_tlb_gen and mm->context.tlb_gen. I.e., this check:

        if (f->end != TLB_FLUSH_ALL &&
            f->new_tlb_gen == local_tlb_gen + 1 &&
            f->new_tlb_gen == mm_tlb_gen) {

Quoted 6 lines

why can't we write:

        if (f->end != TLB_FLUSH_ALL &&
            mm_tlb_gen == local_tlb_gen + 1)

?

Quoted 7 lines

Ah, I thought you were asking about why I needed mm_tlb_gen == local_tlb_gen + 1. This is just an optimization, or at least I hope it is. The idea is that, if we know that another flush is coming, it seems likely that it would be faster to do a full flush and increase local_tlb_gen all the way to mm_tlb_gen rather than doing a partial flush, increasing local_tlb_gen to something less than mm_tlb_gen, and needing to flush again very soon.

Thus the f->new_tlb_gen check whether it is local_tlb_gen + 1.

Btw, do you see how confusing this check is: you have new_tlb_gen from a variable passed from the function call IPI, local_tlb_gen which is the CPU's own and then there's also mm_tlb_gen which we've written into new_tlb_gen from:

        info.new_tlb_gen = bump_mm_tlb_gen(mm);

which incremented mm_tlb_gen too.

Quoted 7 lines

Hmm. I'd be nervous that there are more subtle races if we do this. For example, suppose that a partial flush increments tlb_gen from 1 to 2 and a full flush increments tlb_gen from 2 to 3. Meanwhile, the CPU is busy switching back and forth between mms, so the partial flush sees the cpu set in mm_cpumask but the full flush doesn't see the cpu set in mm_cpumask.

Lemme see if I understand this correctly: you mean, the full flush will exit early due to the

        if (!cpumask_test_cpu(smp_processor_id(), mm_cpumask(loaded_mm))) {

test?

Quoted 6 lines

The flush IPI hits after a switch_mm_irqs_off() call notices the change from 1 to 2. switch_mm_irqs_off() will do a full flush and increment the local tlb_gen to 2, and the IPI handler for the partial flush will see local_tlb_gen == mm_tlb_gen - 1 (because local_tlb_gen == 2 and mm_tlb_gen == 3) and do a partial flush.

Why, the 2->3 flush has f->end == TLB_FLUSH_ALL.

That's why you have this thing in addition to the tlb_gen.

What we end up doing in this case, is promote the partial flush to a full one and thus have a partial and a full flush which are close by converted to two full flushes.

Quoted 6 lines

The problem here is that it's not obvious to me that this actually ends up flushing everything that's needed. Maybe all the memory ordering gets this right, but I can imagine scenarios in which switch_mm_irqs_off() does its flush early enough that the TLB picks up an entry that was supposed to get zapped by the full flush.

See above.

And I don't think that having two full flushes back-to-back is going to cost a lot as the second one won't flush a whole lot.

From limited benchmarking on new Intel chips, a full flush is very

Quoted 3 lines

IOW it *might* be valid, but I think it would need very careful review and documentation.

Always.

Btw, I get the sense this TLB flush avoiding scheme becomes pretty complex for diminishing reasons.

[ Or maybe I'm not seeing them - I'm always open to corrections. ]

Especially if intermediary levels from the pagetable walker are cached and reestablishing the TLB entries seldom means a full walk.

You should do a full fledged benchmark to see whether this whole complexity is even worth it, methinks.

Yes, I agree it's confusing. There really are three numbers. Those numbers are: the latest generation, the generation that this CPU has caught up to, and the generation that the requester of the flush we're currently handling has asked us to catch up to. I don't see a way to reduce the complexity.

Quoted 6 lines

Quoted 5 lines

The flush IPI hits after a switch_mm_irqs_off() call notices the change from 1 to 2. switch_mm_irqs_off() will do a full flush and increment the local tlb_gen to 2, and the IPI handler for the partial flush will see local_tlb_gen == mm_tlb_gen - 1 (because local_tlb_gen == 2 and mm_tlb_gen == 3) and do a partial flush.

Quoted 3 lines

Why, the 2->3 flush has f->end == TLB_FLUSH_ALL.

That's why you have this thing in addition to the tlb_gen.

Yes. The idea is that we only do remote partial flushes when it's 100% obvious that it's safe.

It could be converted to two full flushes or to just one, I think, depending on what order everything happens in.

But this approach of using three separate tlb_gen values seems to cover all the bases, and I don't think it's *that* bad.

Quoted 5 lines · Andy Lutomirski

Yes, I agree it's confusing. There really are three numbers. Those numbers are: the latest generation, the generation that this CPU has caught up to, and the generation that the requester of the flush we're currently handling has asked us to catch up to. I don't see a way to reduce the complexity.

Yeah, can you pls put that clarification what what is, over it. It explains it nicely what the check is supposed to do.

Quoted 7 lines

Quoted 7 lines

Quoted 6 lines

The flush IPI hits after a switch_mm_irqs_off() call notices the change from 1 to 2. switch_mm_irqs_off() will do a full flush and increment the local tlb_gen to 2, and the IPI handler for the partial flush will see local_tlb_gen == mm_tlb_gen - 1 (because local_tlb_gen == 2 and mm_tlb_gen == 3) and do a partial flush.

Quoted 6 lines

Quoted 3 lines

Why, the 2->3 flush has f->end == TLB_FLUSH_ALL.

That's why you have this thing in addition to the tlb_gen.

Yes. The idea is that we only do remote partial flushes when it's 100% obvious that it's safe.

So why wouldn't my simplified suggestion work then?

        if (f->end != TLB_FLUSH_ALL &&
             mm_tlb_gen == local_tlb_gen + 1)

1->2 is a partial flush - gets promoted to a full one
2->3 is a full flush - it will get executed as one due to the f->end setting to TLB_FLUSH_ALL.

Quoted 3 lines

It could be converted to two full flushes or to just one, I think, depending on what order everything happens in.

Right. One flush at the right time would be optimal.

Quoted 2 lines

But this approach of using three separate tlb_gen values seems to cover all the bases, and I don't think it's *that* bad.

Sure.

As I said in IRC, let's document that complexity then so that when we stumble over it in the future, we at least know why it was done this way.

-/*
- * The flush IPI assumes that a thread switch happens in this order:
- * [cpu0: the cpu that switches]
- * 1) switch_mm() either 1a) or 1b)
- * 1a) thread switch to a different mm
- * 1a1) set cpu_tlbstate to TLBSTATE_OK
- * Now the tlb flush NMI handler flush_tlb_func won't call leave_mm
- * if cpu0 was in lazy tlb mode.
- * 1a2) update cpu active_mm
- * Now cpu0 accepts tlb flushes for the new mm.
- * 1a3) cpu_set(cpu, new_mm->cpu_vm_mask);
- * Now the other cpus will send tlb flush ipis.
- * 1a4) change cr3.
- * 1a5) cpu_clear(cpu, old_mm->cpu_vm_mask);
- * Stop ipi delivery for the old mm. This is not synchronized with
- * the other cpus, but flush_tlb_func ignore flush ipis for the wrong
- * mm, and in the worst case we perform a superfluous tlb flush.
- * 1b) thread switch without mm change
- * cpu active_mm is correct, cpu0 already handles flush ipis.
- * 1b1) set cpu_tlbstate to TLBSTATE_OK
- * 1b2) test_and_set the cpu bit in cpu_vm_mask.
- * Atomically set the bit [other cpus will start sending flush ipis],
- * and test the bit.
- * 1b3) if the bit was 0: leave_mm was called, flush the tlb.
- * 2) switch %%esp, ie current
- *

- * The interrupt must handle 2 special cases:
- * - cr3 is changed before %%esp, ie. it cannot use current->{active_,}mm.
- * - the cpu performs speculative tlb reads, i.e. even if the cpu only
- * runs in kernel space, the cpu could load tlb entries for user space
- * pages.
- *
- * The good news is that cpu_tlbstate is local to each cpu, no
- * write/read ordering problems.

Quoted 26 lines · Andy Lutomirski

-/*
- * The flush IPI assumes that a thread switch happens in this order:
- * [cpu0: the cpu that switches]
- * 1) switch_mm() either 1a) or 1b)
- * 1a) thread switch to a different mm
- * 1a1) set cpu_tlbstate to TLBSTATE_OK
- * Now the tlb flush NMI handler flush_tlb_func won't call leave_mm
- * if cpu0 was in lazy tlb mode.
- * 1a2) update cpu active_mm
- * Now cpu0 accepts tlb flushes for the new mm.
- * 1a3) cpu_set(cpu, new_mm->cpu_vm_mask);
- * Now the other cpus will send tlb flush ipis.
- * 1a4) change cr3.
- * 1a5) cpu_clear(cpu, old_mm->cpu_vm_mask);

- *   Stop ipi delivery for the old mm. This is not synchronized with
- *   the other cpus, but flush_tlb_func ignore flush ipis for the wrong
- *   mm, and in the worst case we perform a superfluous tlb flush.

- * 1b) thread switch without mm change
- * cpu active_mm is correct, cpu0 already handles flush ipis.
- * 1b1) set cpu_tlbstate to TLBSTATE_OK
- * 1b2) test_and_set the cpu bit in cpu_vm_mask.
- * Atomically set the bit [other cpus will start sending flush ipis],
- * and test the bit.
- * 1b3) if the bit was 0: leave_mm was called, flush the tlb.
- * 2) switch %%esp, ie current
- *

Quoted 8 lines

- * The interrupt must handle 2 special cases:
- * - cr3 is changed before %%esp, ie. it cannot use current->{active_,}mm.
- * - the cpu performs speculative tlb reads, i.e. even if the cpu only
- * runs in kernel space, the cpu could load tlb entries for user space
- * pages.
- *
- * The good news is that cpu_tlbstate is local to each cpu, no
- * write/read ordering problems.

While the new code is really well commented, it would be a good thing to have a single place where all of this including the ordering constraints are documented.

Quoted 4 lines

arch/x86/mm/tlb.c +0 -0

@@ -215,12 +200,13 @@ static void flush_tlb_func_common(const struct flush_tlb_info *f,
      VM_WARN_ON(this_cpu_read(cpu_tlbstate.ctxs[0].ctx_id) !=
                 loaded_mm->context.ctx_id);

arch/x86/mm/tlb.c static void flush_tlb_func_common(const struct flush_tlb_info *f, · line 208

Quoted 10 lines

arch/x86/mm/tlb.c +5 -3

 
-     if (this_cpu_read(cpu_tlbstate.state) != TLBSTATE_OK) {
+     if (!cpumask_test_cpu(smp_processor_id(), mm_cpumask(loaded_mm))) {
              /*
-              * leave_mm() is adequate to handle any type of flush, and
-              * we would prefer not to receive further IPIs.
+              * We're in lazy mode -- don't flush.  We can get here on
+              * remote flushes due to races and on local flushes if a
+              * kernel thread coincidentally flushes the mm it's lazily
+              * still using.

Ok. That's more informative.

Reviewed-by Thomas Gleixner <[email protected]>

I'll look at the end of the whole series and see if I can come up with something good.

x86's lazy TLB mode used to be fairly weak -- it would switch to init_mm the first time it tried to flush a lazy TLB. This meant an unnecessary CR3 write and, if the flush was remote, an unnecessary IPI.

Rewrite it entirely. When we enter lazy mode, we simply remove the cpu from mm_cpumask. This means that we need a way to figure out

whether we've missed a flush when we switch back out of lazy mode. I use the tlb_gen machinery to track whether a context is up to date.

Note to reviewers: this patch, my itself, looks a bit odd. I'm using an array of length 1 containing (ctx_id, tlb_gen) rather than just storing tlb_gen, and making it at array isn't necessary yet. I'm doing this because the next few patches add PCID support, and, with PCID, we need ctx_id, and the array will end up with a length greater than 1. Making it an array now means that there will be less churn and therefore less stress on your eyeballs.

NB: This is dubious but, AFAICT, still correct on Xen and UV. xen_exit_mmap() uses mm_cpumask() for nefarious purposes and this patch changes the way that mm_cpumask() works. This should be okay, since Xen *also* iterates all online CPUs to find all the CPUs it needs to twiddle.

The UV tlbflush code is rather dated and should be changed.

Cc Andrew Banman <[email protected]>

Cc Mike Travis <[email protected]>

Cc Dimitri Sivanich <[email protected]>

Cc Juergen Gross <[email protected]>

Cc Boris Ostrovsky <[email protected]>

Signed-off-by Andy Lutomirski <[email protected]>

---
 arch/x86/include/asm/mmu_context.h |   6 +-
 arch/x86/include/asm/tlbflush.h    |   4 -
 arch/x86/mm/init.c                 |   1 -
 arch/x86/mm/tlb.c                  | 227 +++++++++++++++++++------------------
 arch/x86/xen/mmu_pv.c              |   3 +-
 5 files changed, 119 insertions(+), 122 deletions(-)

arch/x86/include/asm/mmu_context.h +0 -0

index e5295d485899..69a4f1ee86ac 100644

arch/x86/include/asm/mmu_context.h +0 -0

 extern atomic64_t last_mm_ctx_id;

arch/x86/include/asm/tlbflush.h +0 -4

index 4f6c30d6ec39..87b13e51e867 100644
@@ -95,7 +95,6 @@ struct tlb_state {
 	 * mode even if we've already switched back to swapper_pg_dir.
 	 */
 	struct mm_struct *loaded_mm;
-	int state;
 
 	/*
 	 * Access to this CR4 shadow and to H/W CR4 is protected by

@@ -310,9 +309,6 @@ static inline void flush_tlb_page(struct vm_area_struct *vma, unsigned long a)
 void native_flush_tlb_others(const struct cpumask *cpumask,
 			     const struct flush_tlb_info *info);
 
-#define TLBSTATE_OK	1
-#define TLBSTATE_LAZY	2
-
 static inline void arch_tlbbatch_add_mm(struct arch_tlbflush_unmap_batch *batch,
 					struct mm_struct *mm)
 {

arch/x86/mm/init.c +0 -1

index 88ee942cb47d..7d6fa4676af9 100644
@@ -812,7 +812,6 @@ void __init zone_sizes_init(void)
 
 DEFINE_PER_CPU_SHARED_ALIGNED(struct tlb_state, cpu_tlbstate) = {
 	.loaded_mm = &init_mm,
-	.state = 0,
 	.cr4 = ~0UL,	/* fail hard if we screw up cr4 shadow initialization */
 };
 EXPORT_SYMBOL_GPL(cpu_tlbstate);

arch/x86/mm/tlb.c +0 -0

index 9f5ef7a5e74a..fea2b07ac7d8 100644

arch/x86/mm/tlb.c +4 -4

 	/*
-	 * NB: The scheduler will call us with prev == next when
-	 * switching from lazy TLB mode to normal mode if active_mm
-	 * isn't changing.  When this happens, there is no guarantee
-	 * that CR3 (and hence cpu_tlbstate.loaded_mm) matches next.
+	 * NB: The scheduler will call us with prev == next when switching
+	 * from lazy TLB mode to normal mode if active_mm isn't changing.
+	 * When this happens, we don't assume that CR3 (and hence
+	 * cpu_tlbstate.loaded_mm) matches next.

arch/x86/mm/tlb.c +0 -6

 	if (real_prev == next) {
-		/*
-		 * There's nothing to do: we always keep the per-mm control
-		 * regs in sync with cpu_tlbstate.loaded_mm.  Just
-		 * sanity-check mm_cpumask.
-		 */
-		if (WARN_ON_ONCE(!cpumask_test_cpu(cpu, mm_cpumask(next))))

arch/x86/mm/tlb.c +8 -0

+			/*
+			 * This gets called via leave_mm() in the idle path
+			 * where RCU functions differently.  Tracing normally
+			 * uses RCU, so we have to call the tracepoint
+			 * specially here.
+			 */
+			trace_tlb_flush_rcuidle(TLB_FLUSH_ON_TASK_SWITCH,
+						TLB_FLUSH_ALL);

arch/x86/mm/tlb.c +0 -4

-	this_cpu_write(cpu_tlbstate.loaded_mm, next);
-	this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id, next->context.ctx_id);
-	this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
-		       atomic64_read(&next->context.tlb_gen));

arch/x86/mm/tlb.c +0 -29

-	/*
-	 * Re-load page tables.
-	 *
-	 * This logic has an ordering constraint:
-	 *
-	 *  CPU 0: Write to a PTE for 'next'
-	 *  CPU 0: load bit 1 in mm_cpumask.  if nonzero, send IPI.
-	 *  CPU 1: set bit 1 in next's mm_cpumask
-	 *  CPU 1: load from the PTE that CPU 0 writes (implicit)
-	 *
-	 * We need to prevent an outcome in which CPU 1 observes
-	 * the new PTE value and CPU 0 observes bit 1 clear in
-	 * mm_cpumask.  (If that occurs, then the IPI will never
-	 * be sent, and CPU 0's TLB will contain a stale entry.)
-	 *
-	 * The bad outcome can occur if either CPU's load is
-	 * reordered before that CPU's store, so both CPUs must
-	 * execute full barriers to prevent this from happening.
-	 *
-	 * Thus, switch_mm needs a full barrier between the
-	 * store to mm_cpumask and any operation that could load
-	 * from next->pgd.  TLB fills are special and can happen
-	 * due to instruction fetches or for no reason at all,
-	 * and neither LOCK nor MFENCE orders them.
-	 * Fortunately, load_cr3() is serializing and gives the
-	 * ordering guarantee we need.
-	 */
-	load_cr3(next->pgd);
-

arch/x86/mm/tlb.c +2 -2

+		this_cpu_write(cpu_tlbstate.loaded_mm, next);
+		write_cr3(__pa(next->pgd));
 
-	/* Stop flush ipis for the previous mm */
-	WARN_ON_ONCE(!cpumask_test_cpu(cpu, mm_cpumask(real_prev)) &&

arch/x86/mm/tlb.c +0 -1

-	/* Load per-mm CR4 and LDTR state */
 	load_mm_cr4(next);
 	switch_ldt(real_prev, next);
 }

Quoted 2 lines · Andy Lutomirski

Rewrite it entirely. When we enter lazy mode, we simply remove the cpu from mm_cpumask. This means that we need a way to figure out

s/cpu/CPU/

Quoted 13 lines

whether we've missed a flush when we switch back out of lazy mode. I use the tlb_gen machinery to track whether a context is up to date.

Note to reviewers: this patch, my itself, looks a bit odd. I'm using an array of length 1 containing (ctx_id, tlb_gen) rather than just storing tlb_gen, and making it at array isn't necessary yet. I'm doing this because the next few patches add PCID support, and, with PCID, we need ctx_id, and the array will end up with a length greater than 1. Making it an array now means that there will be less churn and therefore less stress on your eyeballs.

Quoted 5 lines

NB: This is dubious but, AFAICT, still correct on Xen and UV. xen_exit_mmap() uses mm_cpumask() for nefarious purposes and this patch changes the way that mm_cpumask() works. This should be okay, since Xen *also* iterates all online CPUs to find all the CPUs it needs to twiddle.

This whole text should be under the "---" line below if we don't want it in the commit message.

Quoted 3 lines

The UV tlbflush code is rather dated and should be changed.

arch/x86/include/asm/mmu_context.h static inline void switch_ldt(struct mm_struct *prev, struct mm_struct *next) · line 131

Quoted 3 lines

arch/x86/include/asm/mmu_context.h +3 -0

+
+     if (cpumask_test_cpu(cpu, mm_cpumask(mm)))
+             cpumask_clear_cpu(cpu, mm_cpumask(mm));

It seems we haz a helper for that: cpumask_test_and_clear_cpu() which does BTR straightaway.

arch/x86/mm/tlb.c void leave_mm(int cpu) · line 49

Quoted 4 lines

arch/x86/mm/tlb.c +2 -2

-     if (this_cpu_read(cpu_tlbstate.state) == TLBSTATE_OK)
-             BUG();
+     /* Warn if we're not lazy. */
+     WARN_ON(cpumask_test_cpu(smp_processor_id(), mm_cpumask(loaded_mm)));

We don't BUG() anymore?

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 70

Quoted 9 lines

      switch_mm(NULL, &init_mm, NULL);

}

arch/x86/mm/tlb.c +1 -0

@@ -67,133 +67,118 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
 {
      unsigned cpu = smp_processor_id();
      struct mm_struct *real_prev = this_cpu_read(cpu_tlbstate.loaded_mm);
+     u64 next_tlb_gen;

Please sort function local variables declaration in a reverse christmas tree order:

        <type> longest_variable_name;
        <type> shorter_var_name;
        <type> even_shorter;
        <type> i;

Quoted 14 lines

arch/x86/mm/tlb.c +4 -4

      /*
-      * NB: The scheduler will call us with prev == next when
-      * switching from lazy TLB mode to normal mode if active_mm
-      * isn't changing.  When this happens, there is no guarantee
-      * that CR3 (and hence cpu_tlbstate.loaded_mm) matches next.
+      * NB: The scheduler will call us with prev == next when switching
+      * from lazy TLB mode to normal mode if active_mm isn't changing.
+      * When this happens, we don't assume that CR3 (and hence
+      * cpu_tlbstate.loaded_mm) matches next.
       *
       * NB: leave_mm() calls us with prev == NULL and tsk == NULL.
       */
 

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 85

Quoted 6 lines

arch/x86/mm/tlb.c +5 -1

-     this_cpu_write(cpu_tlbstate.state, TLBSTATE_OK);
+     /* We don't want flush_tlb_func_* to run concurrently with us. */
+     if (IS_ENABLED(CONFIG_PROVE_LOCKING))
+             WARN_ON_ONCE(!irqs_disabled());
+
+     VM_BUG_ON(read_cr3_pa() != __pa(real_prev->pgd));

Why do we need that check? Can that ever happen?

Quoted 8 lines

arch/x86/mm/tlb.c +0 -6

      if (real_prev == next) {
-             /*
-              * There's nothing to do: we always keep the per-mm control
-              * regs in sync with cpu_tlbstate.loaded_mm.  Just
-              * sanity-check mm_cpumask.
-              */
-             if (WARN_ON_ONCE(!cpumask_test_cpu(cpu, mm_cpumask(next))))

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 94

Quoted 10 lines

arch/x86/mm/tlb.c +7 -3

-                     cpumask_set_cpu(cpu, mm_cpumask(next));
-             return;
-     }
+             if (cpumask_test_cpu(cpu, mm_cpumask(next))) {
+                     /*
+                      * There's nothing to do: we weren't lazy, and we
+                      * aren't changing our mm.  We don't need to flush
+                      * anything, nor do we need to update CR3, CR4, or
+                      * LDTR.
+                      */

Nice comment.

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 103

Quoted 9 lines

arch/x86/mm/tlb.c +9 -0

+                     return;
+             }
+
+             /* Resume remote flushes and then read tlb_gen. */
+             cpumask_set_cpu(cpu, mm_cpumask(next));
+             next_tlb_gen = atomic64_read(&next->context.tlb_gen);
+
+             VM_BUG_ON(this_cpu_read(cpu_tlbstate.ctxs[0].ctx_id) !=
+                       next->context.ctx_id);

I guess this check should be right under the if (real_prev == next), right?

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 106

Quoted 4 lines

arch/x86/mm/tlb.c +3 -0

+
+             if (this_cpu_read(cpu_tlbstate.ctxs[0].tlb_gen) <
+                 next_tlb_gen) {

Yeah, let it stick out - that trailing '<' doesn't make it any nicer.

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 115

Quoted 9 lines

arch/x86/mm/tlb.c +9 -0

+                     /*
+                      * Ideally, we'd have a flush_tlb() variant that
+                      * takes the known CR3 value as input.  This would
+                      * be faster on Xen PV and on hypothetical CPUs
+                      * on which INVPCID is fast.
+                      */
+                     this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
+                                    next_tlb_gen);
+                     write_cr3(__pa(next->pgd));

<---- newline here.

Quoted 8 lines

arch/x86/mm/tlb.c +8 -0

+                     /*
+                      * This gets called via leave_mm() in the idle path
+                      * where RCU functions differently.  Tracing normally
+                      * uses RCU, so we have to call the tracepoint
+                      * specially here.
+                      */
+                     trace_tlb_flush_rcuidle(TLB_FLUSH_ON_TASK_SWITCH,
+                                             TLB_FLUSH_ALL);

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 129

Quoted 10 lines

arch/x86/mm/tlb.c +4 -4

+             }
 
-     if (IS_ENABLED(CONFIG_VMAP_STACK)) {
              /*
-              * If our current stack is in vmalloc space and isn't
-              * mapped in the new pgd, we'll double-fault.  Forcibly
-              * map it.
+              * We just exited lazy mode, which means that CR4 and/or LDTR
+              * may be stale.  (Changes to the required CR4 and LDTR states
+              * are not reflected in tlb_gen.)

We need that comment because... ? I mean, we do update CR4/LDTR at the end of the function.

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 138

Quoted 11 lines

arch/x86/mm/tlb.c +8 -1

               */
-             unsigned int stack_pgd_index = pgd_index(current_stack_pointer());
+     } else {
+             if (IS_ENABLED(CONFIG_VMAP_STACK)) {
+                     /*
+                      * If our current stack is in vmalloc space and isn't
+                      * mapped in the new pgd, we'll double-fault.  Forcibly
+                      * map it.
+                      */
+                     unsigned int stack_pgd_index =

Shorten that var name and make it fit into 80ish cols so that the linebreak is gone.

Quoted 9 lines

arch/x86/mm/tlb.c +7 -0

+                             pgd_index(current_stack_pointer());
+
+                     pgd_t *pgd = next->pgd + stack_pgd_index;
+
+                     if (unlikely(pgd_none(*pgd)))
+                             set_pgd(pgd, init_mm.pgd[stack_pgd_index]);
+             }
 

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 149

Quoted 4 lines

arch/x86/mm/tlb.c +3 -1

-             pgd_t *pgd = next->pgd + stack_pgd_index;
+             /* Stop remote flushes for the previous mm */
+             if (cpumask_test_cpu(cpu, mm_cpumask(real_prev)))
+                     cpumask_clear_cpu(cpu, mm_cpumask(real_prev));

cpumask_test_and_clear_cpu()

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 151

Quoted 6 lines

arch/x86/mm/tlb.c +1 -3

-             if (unlikely(pgd_none(*pgd)))
-                     set_pgd(pgd, init_mm.pgd[stack_pgd_index]);
-     }
+             WARN_ON_ONCE(cpumask_test_cpu(cpu, mm_cpumask(next)));

We warn if the next task is not lazy because...?

Quoted 11 lines

arch/x86/mm/tlb.c +5 -4

-     this_cpu_write(cpu_tlbstate.loaded_mm, next);
-     this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id, next->context.ctx_id);
-     this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
-                    atomic64_read(&next->context.tlb_gen));
+             /*
+              * Start remote flushes and then read tlb_gen.
+              */
+             cpumask_set_cpu(cpu, mm_cpumask(next));
+             next_tlb_gen = atomic64_read(&next->context.tlb_gen);
 

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 160

Quoted 4 lines

arch/x86/mm/tlb.c +2 -2

-     WARN_ON_ONCE(cpumask_test_cpu(cpu, mm_cpumask(next)));
-     cpumask_set_cpu(cpu, mm_cpumask(next));
+             VM_BUG_ON(this_cpu_read(cpu_tlbstate.ctxs[0].ctx_id) ==
+                       next->context.ctx_id);

Also put it as the first statement after the else { ?

Quoted 30 lines

arch/x86/mm/tlb.c +0 -29

-     /*
-      * Re-load page tables.
-      *
-      * This logic has an ordering constraint:
-      *
-      *  CPU 0: Write to a PTE for 'next'
-      *  CPU 0: load bit 1 in mm_cpumask.  if nonzero, send IPI.
-      *  CPU 1: set bit 1 in next's mm_cpumask
-      *  CPU 1: load from the PTE that CPU 0 writes (implicit)
-      *
-      * We need to prevent an outcome in which CPU 1 observes
-      * the new PTE value and CPU 0 observes bit 1 clear in
-      * mm_cpumask.  (If that occurs, then the IPI will never
-      * be sent, and CPU 0's TLB will contain a stale entry.)
-      *
-      * The bad outcome can occur if either CPU's load is
-      * reordered before that CPU's store, so both CPUs must
-      * execute full barriers to prevent this from happening.
-      *
-      * Thus, switch_mm needs a full barrier between the
-      * store to mm_cpumask and any operation that could load
-      * from next->pgd.  TLB fills are special and can happen
-      * due to instruction fetches or for no reason at all,
-      * and neither LOCK nor MFENCE orders them.
-      * Fortunately, load_cr3() is serializing and gives the
-      * ordering guarantee we need.
-      */
-     load_cr3(next->pgd);
-

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 165

Quoted 10 lines

arch/x86/mm/tlb.c +4 -6

-     /*
-      * This gets called via leave_mm() in the idle path where RCU
-      * functions differently.  Tracing normally uses RCU, so we have to
-      * call the tracepoint specially here.
-      */
-     trace_tlb_flush_rcuidle(TLB_FLUSH_ON_TASK_SWITCH, TLB_FLUSH_ALL);
+             this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id,
+                            next->context.ctx_id);
+             this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
+                            next_tlb_gen);

Yeah, just let all those three stick out.

Also, no:

                if (this_cpu_read(cpu_tlbstate.ctxs[0].tlb_gen) <
                    next_tlb_gen) {

check?

Quoted 6 lines

arch/x86/mm/tlb.c +2 -2

+             this_cpu_write(cpu_tlbstate.loaded_mm, next);
+             write_cr3(__pa(next->pgd));
 
-     /* Stop flush ipis for the previous mm */
-     WARN_ON_ONCE(!cpumask_test_cpu(cpu, mm_cpumask(real_prev)) &&

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 176

Quoted 10 lines

arch/x86/mm/tlb.c +8 -2

-                  real_prev != &init_mm);
-     cpumask_clear_cpu(cpu, mm_cpumask(real_prev));
+             /*
+              * This gets called via leave_mm() in the idle path where RCU
+              * functions differently.  Tracing normally uses RCU, so we
+              * have to call the tracepoint specially here.
+              */
+             trace_tlb_flush_rcuidle(TLB_FLUSH_ON_TASK_SWITCH,
+                                     TLB_FLUSH_ALL);
+     }

That's repeated as in the if-branch above. Move it out I guess.

I figured that some future reader of this patch might actually want to see this text, though.

Quoted 1 line

Quoted 1 line

The UV tlbflush code is rather dated and should be changed.

And I'd definitely like the UV maintainers to notice this part, now or in the future :) I don't want to personally touch the UV code with a ten-foot pole, but it really should be updated by someone who has a chance of getting it right and being able to test it.

Quoted 6 lines

arch/x86/include/asm/mmu_context.h static inline void switch_ldt(struct mm_struct *prev, struct mm_struct *next) · line 131

Quoted 3 lines

arch/x86/include/asm/mmu_context.h +3 -0

+
+     if (cpumask_test_cpu(cpu, mm_cpumask(mm)))
+             cpumask_clear_cpu(cpu, mm_cpumask(mm));

It seems we haz a helper for that: cpumask_test_and_clear_cpu() which does BTR straightaway.

Yeah, but I'm doing this for performance. I think that all the various one-line helpers do a LOCKed op right away, and I think it's faster to see if we can avoid the LOCKed op by trying an ordinary read first.

OTOH, maybe this is misguided -- if the cacheline lives somewhere else and we do end up needing to update it, we'll end up first sharing it and then making it exclusive, which increases the amount of cache coherency traffic, so maybe I'm optimizing for the wrong thing. What do you think?

It did in one particular buggy incarnation. It would also trigger if, say, suspend/resume corrupts CR3. Admittedly this is unlikely, but I'd rather catch it. Once PCID is on, corruption seems a bit less farfetched -- this assertion will catch anyone who accidentally does write_cr3(read_cr3_pa()).

Quoted 5 lines · Borislav Petkov

Quoted 2 lines · Andy Lutomirski

Rewrite it entirely. When we enter lazy mode, we simply remove the cpu from mm_cpumask. This means that we need a way to figure out

s/cpu/CPU/

Done.

Quoted 13 lines

Quoted 13 lines

whether we've missed a flush when we switch back out of lazy mode. I use the tlb_gen machinery to track whether a context is up to date.

Note to reviewers: this patch, my itself, looks a bit odd. I'm using an array of length 1 containing (ctx_id, tlb_gen) rather than just storing tlb_gen, and making it at array isn't necessary yet. I'm doing this because the next few patches add PCID support, and, with PCID, we need ctx_id, and the array will end up with a length greater than 1. Making it an array now means that there will be less churn and therefore less stress on your eyeballs.

Quoted 8 lines

Quoted 5 lines

NB: This is dubious but, AFAICT, still correct on Xen and UV. xen_exit_mmap() uses mm_cpumask() for nefarious purposes and this patch changes the way that mm_cpumask() works. This should be okay, since Xen *also* iterates all online CPUs to find all the CPUs it needs to twiddle.

This whole text should be under the "---" line below if we don't want it in the commit message.

I figured that some future reader of this patch might actually want to see this text, though.

Quoted 2 lines

Quoted 2 lines

Quoted 1 line

Quoted 1 line

The UV tlbflush code is rather dated and should be changed.

And I'd definitely like the UV maintainers to notice this part, now or in the future :) I don't want to personally touch the UV code with a ten-foot pole, but it really should be updated by someone who has a chance of getting it right and being able to test it.

Quoted 7 lines

arch/x86/include/asm/mmu_context.h static inline void switch_ldt(struct mm_struct *prev, struct mm_struct *next) · line 131

Quoted 4 lines

arch/x86/include/asm/mmu_context.h +3 -0

+
+     if (cpumask_test_cpu(cpu, mm_cpumask(mm)))
+             cpumask_clear_cpu(cpu, mm_cpumask(mm));

It seems we haz a helper for that: cpumask_test_and_clear_cpu() which does BTR straightaway.

Yeah, but I'm doing this for performance. I think that all the various one-line helpers do a LOCKed op right away, and I think it's faster to see if we can avoid the LOCKed op by trying an ordinary read first. OTOH, maybe this is misguided -- if the cacheline lives somewhere else and we do end up needing to update it, we'll end up first sharing it and then making it exclusive, which increases the amount of cache coherency traffic, so maybe I'm optimizing for the wrong thing. What do you think?

Quoted 6 lines

arch/x86/mm/tlb.c void leave_mm(int cpu) · line 49

Quoted 4 lines

arch/x86/mm/tlb.c +2 -2

-     if (this_cpu_read(cpu_tlbstate.state) == TLBSTATE_OK)
-             BUG();
+     /* Warn if we're not lazy. */
+     WARN_ON(cpumask_test_cpu(smp_processor_id(), mm_cpumask(loaded_mm)));

We don't BUG() anymore?

We could. But, when the whole patch series is applied, the only caller left is a somewhat dubious Xen optimization, and if we blindly continue executing, I think the worst that happens is that we OOPS later or that we get segfaults when we shouldn't get segfaults.

Quoted 32 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 70

Quoted 9 lines

      switch_mm(NULL, &init_mm, NULL);

}

arch/x86/mm/tlb.c +1 -0

@@ -67,133 +67,118 @@ void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next,
 {
      unsigned cpu = smp_processor_id();
      struct mm_struct *real_prev = this_cpu_read(cpu_tlbstate.loaded_mm);
+     u64 next_tlb_gen;

Please sort function local variables declaration in a reverse christmas tree order:

        <type> longest_variable_name;
        <type> shorter_var_name;
        <type> even_shorter;
        <type> i;

Quoted 14 lines

arch/x86/mm/tlb.c +4 -4

      /*
-      * NB: The scheduler will call us with prev == next when
-      * switching from lazy TLB mode to normal mode if active_mm
-      * isn't changing.  When this happens, there is no guarantee
-      * that CR3 (and hence cpu_tlbstate.loaded_mm) matches next.
+      * NB: The scheduler will call us with prev == next when switching
+      * from lazy TLB mode to normal mode if active_mm isn't changing.
+      * When this happens, we don't assume that CR3 (and hence
+      * cpu_tlbstate.loaded_mm) matches next.
       *
       * NB: leave_mm() calls us with prev == NULL and tsk == NULL.
       */
 

Quoted 8 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 85

Quoted 6 lines

arch/x86/mm/tlb.c +5 -1

-     this_cpu_write(cpu_tlbstate.state, TLBSTATE_OK);
+     /* We don't want flush_tlb_func_* to run concurrently with us. */
+     if (IS_ENABLED(CONFIG_PROVE_LOCKING))
+             WARN_ON_ONCE(!irqs_disabled());
+
+     VM_BUG_ON(read_cr3_pa() != __pa(real_prev->pgd));

Why do we need that check? Can that ever happen?

It did in one particular buggy incarnation. It would also trigger if, say, suspend/resume corrupts CR3. Admittedly this is unlikely, but I'd rather catch it. Once PCID is on, corruption seems a bit less farfetched -- this assertion will catch anyone who accidentally does write_cr3(read_cr3_pa()).

Quoted 21 lines

Quoted 8 lines

arch/x86/mm/tlb.c +0 -6

      if (real_prev == next) {
-             /*
-              * There's nothing to do: we always keep the per-mm control
-              * regs in sync with cpu_tlbstate.loaded_mm.  Just
-              * sanity-check mm_cpumask.
-              */
-             if (WARN_ON_ONCE(!cpumask_test_cpu(cpu, mm_cpumask(next))))

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 94

Quoted 10 lines

arch/x86/mm/tlb.c +7 -3

-                     cpumask_set_cpu(cpu, mm_cpumask(next));
-             return;
-     }
+             if (cpumask_test_cpu(cpu, mm_cpumask(next))) {
+                     /*
+                      * There's nothing to do: we weren't lazy, and we
+                      * aren't changing our mm.  We don't need to flush
+                      * anything, nor do we need to update CR3, CR4, or
+                      * LDTR.
+                      */

Nice comment.

Quoted 11 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 103

Quoted 9 lines

arch/x86/mm/tlb.c +9 -0

+                     return;
+             }
+
+             /* Resume remote flushes and then read tlb_gen. */
+             cpumask_set_cpu(cpu, mm_cpumask(next));
+             next_tlb_gen = atomic64_read(&next->context.tlb_gen);
+
+             VM_BUG_ON(this_cpu_read(cpu_tlbstate.ctxs[0].ctx_id) !=
+                       next->context.ctx_id);

I guess this check should be right under the if (real_prev == next), right?

Moved.

Quoted 38 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 106

Quoted 4 lines

arch/x86/mm/tlb.c +3 -0

+
+             if (this_cpu_read(cpu_tlbstate.ctxs[0].tlb_gen) <
+                 next_tlb_gen) {

Yeah, let it stick out - that trailing '<' doesn't make it any nicer.

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 115

Quoted 9 lines

arch/x86/mm/tlb.c +9 -0

+                     /*
+                      * Ideally, we'd have a flush_tlb() variant that
+                      * takes the known CR3 value as input.  This would
+                      * be faster on Xen PV and on hypothetical CPUs
+                      * on which INVPCID is fast.
+                      */
+                     this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
+                                    next_tlb_gen);
+                     write_cr3(__pa(next->pgd));

<---- newline here.

Quoted 8 lines

arch/x86/mm/tlb.c +8 -0

+                     /*
+                      * This gets called via leave_mm() in the idle path
+                      * where RCU functions differently.  Tracing normally
+                      * uses RCU, so we have to call the tracepoint
+                      * specially here.
+                      */
+                     trace_tlb_flush_rcuidle(TLB_FLUSH_ON_TASK_SWITCH,
+                                             TLB_FLUSH_ALL);

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 129

Quoted 10 lines

arch/x86/mm/tlb.c +4 -4

+             }
 
-     if (IS_ENABLED(CONFIG_VMAP_STACK)) {
              /*
-              * If our current stack is in vmalloc space and isn't
-              * mapped in the new pgd, we'll double-fault.  Forcibly
-              * map it.
+              * We just exited lazy mode, which means that CR4 and/or LDTR
+              * may be stale.  (Changes to the required CR4 and LDTR states
+              * are not reflected in tlb_gen.)

Quoted 2 lines

We need that comment because... ? I mean, we do update CR4/LDTR at the end of the function.

I'm trying to explain to the potentially confused reader why we fall through and update CR4 and LDTR even if we decided not to update the TLB.

Quoted 12 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 138

Quoted 11 lines

arch/x86/mm/tlb.c +8 -1

               */
-             unsigned int stack_pgd_index = pgd_index(current_stack_pointer());
+     } else {
+             if (IS_ENABLED(CONFIG_VMAP_STACK)) {
+                     /*
+                      * If our current stack is in vmalloc space and isn't
+                      * mapped in the new pgd, we'll double-fault.  Forcibly
+                      * map it.
+                      */
+                     unsigned int stack_pgd_index =

Quoted 2 lines

Shorten that var name and make it fit into 80ish cols so that the linebreak is gone.

Done.

Quoted 9 lines

Quoted 9 lines

arch/x86/mm/tlb.c +7 -0

+                             pgd_index(current_stack_pointer());
+
+                     pgd_t *pgd = next->pgd + stack_pgd_index;
+
+                     if (unlikely(pgd_none(*pgd)))
+                             set_pgd(pgd, init_mm.pgd[stack_pgd_index]);
+             }
 

Quoted 6 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 149

Quoted 4 lines

arch/x86/mm/tlb.c +3 -1

-             pgd_t *pgd = next->pgd + stack_pgd_index;
+             /* Stop remote flushes for the previous mm */
+             if (cpumask_test_cpu(cpu, mm_cpumask(real_prev)))
+                     cpumask_clear_cpu(cpu, mm_cpumask(real_prev));

cpumask_test_and_clear_cpu()

Same as before. I can change this, but it'll have different performance characteristics. This one here will optimize the case where we go lazy and then switch away.

Quoted 8 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 151

Quoted 6 lines

arch/x86/mm/tlb.c +1 -3

-             if (unlikely(pgd_none(*pgd)))
-                     set_pgd(pgd, init_mm.pgd[stack_pgd_index]);
-     }
+             WARN_ON_ONCE(cpumask_test_cpu(cpu, mm_cpumask(next)));

We warn if the next task is not lazy because...?

The next task had better not be in mm_cpumask(), but I agree that this is minor. I'll change it to VM_WARN_ON_ONCE.

Quoted 11 lines

Quoted 11 lines

arch/x86/mm/tlb.c +5 -4

-     this_cpu_write(cpu_tlbstate.loaded_mm, next);
-     this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id, next->context.ctx_id);
-     this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
-                    atomic64_read(&next->context.tlb_gen));
+             /*
+              * Start remote flushes and then read tlb_gen.
+              */
+             cpumask_set_cpu(cpu, mm_cpumask(next));
+             next_tlb_gen = atomic64_read(&next->context.tlb_gen);
 

Quoted 6 lines

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 160

Quoted 4 lines

arch/x86/mm/tlb.c +2 -2

-     WARN_ON_ONCE(cpumask_test_cpu(cpu, mm_cpumask(next)));
-     cpumask_set_cpu(cpu, mm_cpumask(next));
+             VM_BUG_ON(this_cpu_read(cpu_tlbstate.ctxs[0].ctx_id) ==
+                       next->context.ctx_id);

Also put it as the first statement after the else { ?

Done.

Quoted 41 lines

Quoted 30 lines

arch/x86/mm/tlb.c +0 -29

-     /*
-      * Re-load page tables.
-      *
-      * This logic has an ordering constraint:
-      *
-      *  CPU 0: Write to a PTE for 'next'
-      *  CPU 0: load bit 1 in mm_cpumask.  if nonzero, send IPI.
-      *  CPU 1: set bit 1 in next's mm_cpumask
-      *  CPU 1: load from the PTE that CPU 0 writes (implicit)
-      *
-      * We need to prevent an outcome in which CPU 1 observes
-      * the new PTE value and CPU 0 observes bit 1 clear in
-      * mm_cpumask.  (If that occurs, then the IPI will never
-      * be sent, and CPU 0's TLB will contain a stale entry.)
-      *
-      * The bad outcome can occur if either CPU's load is
-      * reordered before that CPU's store, so both CPUs must
-      * execute full barriers to prevent this from happening.
-      *
-      * Thus, switch_mm needs a full barrier between the
-      * store to mm_cpumask and any operation that could load
-      * from next->pgd.  TLB fills are special and can happen
-      * due to instruction fetches or for no reason at all,
-      * and neither LOCK nor MFENCE orders them.
-      * Fortunately, load_cr3() is serializing and gives the
-      * ordering guarantee we need.
-      */
-     load_cr3(next->pgd);
-

arch/x86/mm/tlb.c void switch_mm_irqs_off(struct mm_struct *prev, struct mm_struct *next, · line 165

Quoted 10 lines

arch/x86/mm/tlb.c +4 -6

-     /*
-      * This gets called via leave_mm() in the idle path where RCU
-      * functions differently.  Tracing normally uses RCU, so we have to
-      * call the tracepoint specially here.
-      */
-     trace_tlb_flush_rcuidle(TLB_FLUSH_ON_TASK_SWITCH, TLB_FLUSH_ALL);
+             this_cpu_write(cpu_tlbstate.ctxs[0].ctx_id,
+                            next->context.ctx_id);
+             this_cpu_write(cpu_tlbstate.ctxs[0].tlb_gen,
+                            next_tlb_gen);

Quoted 8 lines

Yeah, just let all those three stick out.

Also, no:

                if (this_cpu_read(cpu_tlbstate.ctxs[0].tlb_gen) <
                    next_tlb_gen) {

check?

What would the check do? Without PCID, we don't have a choice as to whether we flush.

Quoted 17 lines

Quoted 16 lines

arch/x86/mm/tlb.c +10 -4

+             this_cpu_write(cpu_tlbstate.loaded_mm, next);
+             write_cr3(__pa(next->pgd));
 
-     /* Stop flush ipis for the previous mm */
-     WARN_ON_ONCE(!cpumask_test_cpu(cpu, mm_cpumask(real_prev)) &&
-                  real_prev != &init_mm);
-     cpumask_clear_cpu(cpu, mm_cpumask(real_prev));
+             /*
+              * This gets called via leave_mm() in the idle path where RCU
+              * functions differently.  Tracing normally uses RCU, so we
+              * have to call the tracepoint specially here.
+              */
+             trace_tlb_flush_rcuidle(TLB_FLUSH_ON_TASK_SWITCH,
+                                     TLB_FLUSH_ALL);
+     }