[PATCH 0/3] x86/mm: Cleanups

[PATCH 0/3] x86/mm: Cleanups

[Peter Zijlstra]

Hi!

When reviewing fea4e317f9e7 ("x86/mm: Eliminate window where TLB flushes may be
inadvertently skipped") various other issues were raised. None of these are
critical, so here goes.

[PATCH 1/3] x86/mm: Unexport tlb_state_shared

[Peter Zijlstra]

Never export data; modules have no business being able to change tlb
state.

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 arch/x86/hyperv/mmu.c           |    9 ++-------
 arch/x86/include/asm/tlbflush.h |    2 ++
 arch/x86/mm/tlb.c               |    7 ++++++-
 3 files changed, 10 insertions(+), 8 deletions(-)

--- a/arch/x86/hyperv/mmu.c
+++ b/arch/x86/hyperv/mmu.c
@@ -51,11 +51,6 @@ static inline int fill_gva_list(u64 gva_
 	return gva_n - offset;
 }
 
-static bool cpu_is_lazy(int cpu)
-{
-	return per_cpu(cpu_tlbstate_shared.is_lazy, cpu);
-}
-
 static void hyperv_flush_tlb_multi(const struct cpumask *cpus,
 				   const struct flush_tlb_info *info)
 {
@@ -113,7 +108,7 @@ static void hyperv_flush_tlb_multi(const
 			goto do_ex_hypercall;
 
 		for_each_cpu(cpu, cpus) {
-			if (do_lazy && cpu_is_lazy(cpu))
+			if (do_lazy && cpu_tlbstate_is_lazy(cpu))
 				continue;
 			vcpu = hv_cpu_number_to_vp_number(cpu);
 			if (vcpu == VP_INVAL) {
@@ -198,7 +193,7 @@ static u64 hyperv_flush_tlb_others_ex(co
 
 	flush->hv_vp_set.format = HV_GENERIC_SET_SPARSE_4K;
 	nr_bank = cpumask_to_vpset_skip(&flush->hv_vp_set, cpus,
-			info->freed_tables ? NULL : cpu_is_lazy);
+			info->freed_tables ? NULL : cpu_tlbstate_is_lazy);
 	if (nr_bank < 0)
 		return HV_STATUS_INVALID_PARAMETER;
 
--- a/arch/x86/include/asm/tlbflush.h
+++ b/arch/x86/include/asm/tlbflush.h
@@ -172,6 +172,8 @@ struct tlb_state_shared {
 };
 DECLARE_PER_CPU_SHARED_ALIGNED(struct tlb_state_shared, cpu_tlbstate_shared);
 
+bool cpu_tlbstate_is_lazy(int cpu);
+
 bool nmi_uaccess_okay(void);
 #define nmi_uaccess_okay nmi_uaccess_okay
 
--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -1322,7 +1322,12 @@ static bool should_trim_cpumask(struct m
 }
 
 DEFINE_PER_CPU_SHARED_ALIGNED(struct tlb_state_shared, cpu_tlbstate_shared);
-EXPORT_PER_CPU_SYMBOL(cpu_tlbstate_shared);
+
+bool cpu_tlbstate_is_lazy(int cpu)
+{
+	return per_cpu(cpu_tlbstate_shared.is_lazy, cpu);
+}
+EXPORT_SYMBOL_GPL(cpu_tlbstate_is_lazy);
 
 STATIC_NOPV void native_flush_tlb_multi(const struct cpumask *cpumask,
 					 const struct flush_tlb_info *info)

[PATCH 2/3] x86/mm: Avoid repeated this_cpu_*() ops in switch_mm_irqs_off()

[Peter Zijlstra]

Aside from generating slightly better code for not having to use %fs
prefixed ops, the real purpose is to clarify code by switching some to
smp_store_release() later on.

Notably, this_cpu_{read,write}() imply {READ,WRITE}_ONCE().

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 arch/x86/mm/tlb.c |   21 +++++++++++----------
 1 file changed, 11 insertions(+), 10 deletions(-)

--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -51,7 +51,7 @@
 
 /*
  * Bits to mangle the TIF_SPEC_* state into the mm pointer which is
- * stored in cpu_tlb_state.last_user_mm_spec.
+ * stored in cpu_tlbstate.last_user_mm_spec.
  */
 #define LAST_USER_MM_IBPB	0x1UL
 #define LAST_USER_MM_L1D_FLUSH	0x2UL
@@ -782,8 +782,9 @@ static inline void cr4_update_pce_mm(str
 void switch_mm_irqs_off(struct mm_struct *unused, struct mm_struct *next,
 			struct task_struct *tsk)
 {
-	struct mm_struct *prev = this_cpu_read(cpu_tlbstate.loaded_mm);
-	u16 prev_asid = this_cpu_read(cpu_tlbstate.loaded_mm_asid);
+	struct tlb_state *this_tlbstate = this_cpu_ptr(&cpu_tlbstate);
+	struct mm_struct *prev = READ_ONCE(this_tlbstate->loaded_mm);
+	u16 prev_asid = READ_ONCE(this_tlbstate->loaded_mm_asid);
 	bool was_lazy = this_cpu_read(cpu_tlbstate_shared.is_lazy);
 	unsigned cpu = smp_processor_id();
 	unsigned long new_lam;
@@ -840,7 +841,7 @@ void switch_mm_irqs_off(struct mm_struct
 	if (prev == next) {
 		/* Not actually switching mm's */
 		VM_WARN_ON(is_dyn_asid(prev_asid) &&
-			   this_cpu_read(cpu_tlbstate.ctxs[prev_asid].ctx_id) !=
+			   READ_ONCE(this_tlbstate->ctxs[prev_asid].ctx_id) !=
 			   next->context.ctx_id);
 
 		/*
@@ -888,7 +889,7 @@ void switch_mm_irqs_off(struct mm_struct
 		 */
 		smp_mb();
 		next_tlb_gen = atomic64_read(&next->context.tlb_gen);
-		if (this_cpu_read(cpu_tlbstate.ctxs[prev_asid].tlb_gen) ==
+		if (READ_ONCE(this_tlbstate->ctxs[prev_asid].tlb_gen) ==
 				next_tlb_gen)
 			return;
 
@@ -910,7 +911,7 @@ void switch_mm_irqs_off(struct mm_struct
 		 * and others are sensitive to the window where mm_cpumask(),
 		 * CR3 and cpu_tlbstate.loaded_mm are not all in sync.
 		 */
-		this_cpu_write(cpu_tlbstate.loaded_mm, LOADED_MM_SWITCHING);
+		WRITE_ONCE(this_tlbstate->loaded_mm, LOADED_MM_SWITCHING);
 		barrier();
 
 		/* Start receiving IPIs and then read tlb_gen (and LAM below) */
@@ -925,8 +926,8 @@ void switch_mm_irqs_off(struct mm_struct
 	new_lam = mm_lam_cr3_mask(next);
 	if (ns.need_flush) {
 		VM_WARN_ON_ONCE(is_global_asid(ns.asid));
-		this_cpu_write(cpu_tlbstate.ctxs[ns.asid].ctx_id, next->context.ctx_id);
-		this_cpu_write(cpu_tlbstate.ctxs[ns.asid].tlb_gen, next_tlb_gen);
+		WRITE_ONCE(this_tlbstate->ctxs[ns.asid].ctx_id, next->context.ctx_id);
+		WRITE_ONCE(this_tlbstate->ctxs[ns.asid].tlb_gen, next_tlb_gen);
 		load_new_mm_cr3(next->pgd, ns.asid, new_lam, true);
 
 		trace_tlb_flush(TLB_FLUSH_ON_TASK_SWITCH, TLB_FLUSH_ALL);
@@ -940,8 +941,8 @@ void switch_mm_irqs_off(struct mm_struct
 	/* Make sure we write CR3 before loaded_mm. */
 	barrier();
 
-	this_cpu_write(cpu_tlbstate.loaded_mm, next);
-	this_cpu_write(cpu_tlbstate.loaded_mm_asid, ns.asid);
+	WRITE_ONCE(this_tlbstate->loaded_mm, next);
+	WRITE_ONCE(this_tlbstate->loaded_mm_asid, ns.asid);
 	cpu_tlbstate_update_lam(new_lam, mm_untag_mask(next));
 
 	if (next != prev) {

[PATCH 3/3] x86/mm: Clarify should_flush_tlb() ordering

[Peter Zijlstra]

The ordering in should_flush_tlb() is entirely non-obvious and is only
correct because x86 is TSO. Clarify the situation by replacing two
WRITE_ONCE()s with smp_store_release(), which on x86 is cosmetic.

Additionally, clarify the comment on should_flush_tlb().

Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
---
 arch/x86/mm/tlb.c |   30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

--- a/arch/x86/mm/tlb.c
+++ b/arch/x86/mm/tlb.c
@@ -910,8 +910,10 @@ void switch_mm_irqs_off(struct mm_struct
 		 * Indicate that CR3 is about to change. nmi_uaccess_okay()
 		 * and others are sensitive to the window where mm_cpumask(),
 		 * CR3 and cpu_tlbstate.loaded_mm are not all in sync.
+		 *
+		 * Also, see should_flush_tlb().
 		 */
-		WRITE_ONCE(this_tlbstate->loaded_mm, LOADED_MM_SWITCHING);
+		smp_store_release(&this_tlbstate->loaded_mm, LOADED_MM_SWITCHING);
 		barrier();
 
 		/* Start receiving IPIs and then read tlb_gen (and LAM below) */
@@ -938,10 +940,11 @@ void switch_mm_irqs_off(struct mm_struct
 		trace_tlb_flush(TLB_FLUSH_ON_TASK_SWITCH, 0);
 	}
 
-	/* Make sure we write CR3 before loaded_mm. */
-	barrier();
-
-	WRITE_ONCE(this_tlbstate->loaded_mm, next);
+	/*
+	 * Make sure we write CR3 before loaded_mm.
+	 * See nmi_uaccess_okay() and should_flush_tlb().
+	 */
+	smp_store_release(&this_tlbstate->loaded_mm, next);
 	WRITE_ONCE(this_tlbstate->loaded_mm_asid, ns.asid);
 	cpu_tlbstate_update_lam(new_lam, mm_untag_mask(next));
 
@@ -1280,9 +1283,20 @@ static bool should_flush_tlb(int cpu, vo
 	struct flush_tlb_info *info = data;
 
 	/*
-	 * Order the 'loaded_mm' and 'is_lazy' against their
-	 * write ordering in switch_mm_irqs_off(). Ensure
-	 * 'is_lazy' is at least as new as 'loaded_mm'.
+	 * switch_mm_irqs_off()				should_flush_tlb()
+	 *   WRITE_ONCE(is_lazy, false);		  loaded_mm = READ_ONCE(loaded_mm);
+	 *   smp_store_release(loaded_mm, SWITCHING);     smp_rmb();
+	 *   mov-cr3
+	 *   smp_store_release(loaded_mm, next)
+	 *                                                if (READ_ONCE(is_lazy))
+	 *                                                  return false;
+	 *
+	 * Where smp_rmb() matches against either smp_store_release() to
+	 * ensure that if we observe loaded_mm to be either SWITCHING or next
+	 * we must also observe is_lazy == false.
+	 *
+	 * If this were not so, it would be possible to falsely return false
+	 * and miss sending an invalidation IPI.
 	 */
 	smp_rmb();
 

RE: [PATCH 1/3] x86/mm: Unexport tlb_state_shared

[Michael Kelley]

From: Peter Zijlstra <peterz@infradead.org> Sent: Tuesday, May 20, 2025 3:56 AM
> 
> Never export data; modules have no business being able to change tlb
> state.
> 
> Signed-off-by: Peter Zijlstra (Intel) <peterz@infradead.org>
> ---
>  arch/x86/hyperv/mmu.c           |    9 ++-------
>  arch/x86/include/asm/tlbflush.h |    2 ++
>  arch/x86/mm/tlb.c               |    7 ++++++-
>  3 files changed, 10 insertions(+), 8 deletions(-)
> 
> --- a/arch/x86/hyperv/mmu.c
> +++ b/arch/x86/hyperv/mmu.c
> @@ -51,11 +51,6 @@ static inline int fill_gva_list(u64 gva_
>  	return gva_n - offset;
>  }
> 
> -static bool cpu_is_lazy(int cpu)
> -{
> -	return per_cpu(cpu_tlbstate_shared.is_lazy, cpu);
> -}
> -
>  static void hyperv_flush_tlb_multi(const struct cpumask *cpus,
>  				   const struct flush_tlb_info *info)
>  {
> @@ -113,7 +108,7 @@ static void hyperv_flush_tlb_multi(const
>  			goto do_ex_hypercall;
> 
>  		for_each_cpu(cpu, cpus) {
> -			if (do_lazy && cpu_is_lazy(cpu))
> +			if (do_lazy && cpu_tlbstate_is_lazy(cpu))
>  				continue;
>  			vcpu = hv_cpu_number_to_vp_number(cpu);
>  			if (vcpu == VP_INVAL) {
> @@ -198,7 +193,7 @@ static u64 hyperv_flush_tlb_others_ex(co
> 
>  	flush->hv_vp_set.format = HV_GENERIC_SET_SPARSE_4K;
>  	nr_bank = cpumask_to_vpset_skip(&flush->hv_vp_set, cpus,
> -			info->freed_tables ? NULL : cpu_is_lazy);
> +			info->freed_tables ? NULL : cpu_tlbstate_is_lazy);
>  	if (nr_bank < 0)
>  		return HV_STATUS_INVALID_PARAMETER;
> 
> --- a/arch/x86/include/asm/tlbflush.h
> +++ b/arch/x86/include/asm/tlbflush.h
> @@ -172,6 +172,8 @@ struct tlb_state_shared {
>  };
>  DECLARE_PER_CPU_SHARED_ALIGNED(struct tlb_state_shared, cpu_tlbstate_shared);
> 
> +bool cpu_tlbstate_is_lazy(int cpu);
> +
>  bool nmi_uaccess_okay(void);
>  #define nmi_uaccess_okay nmi_uaccess_okay
> 
> --- a/arch/x86/mm/tlb.c
> +++ b/arch/x86/mm/tlb.c
> @@ -1322,7 +1322,12 @@ static bool should_trim_cpumask(struct m
>  }
> 
>  DEFINE_PER_CPU_SHARED_ALIGNED(struct tlb_state_shared, cpu_tlbstate_shared);
> -EXPORT_PER_CPU_SYMBOL(cpu_tlbstate_shared);

FWIW, this EXPORT wasn't there for Hyper-V code. The EXPORT was
added in commit 2f4305b19fe6a, and it's not clear why. That commit
was 2 years before the Hyper-V MMU code started checking the lazy
flag.

> +
> +bool cpu_tlbstate_is_lazy(int cpu)
> +{
> +	return per_cpu(cpu_tlbstate_shared.is_lazy, cpu);
> +}
> +EXPORT_SYMBOL_GPL(cpu_tlbstate_is_lazy);

This EXPORT isn't needed for Hyper-V. The Hyper-V MMU code is
never built as a module, even if CONFIG_HYPERV=m. And I can't
see any other reason the EXPORT would be needed.

In any case,

Reviewed-by: Michael Kelley <mhklinux@outlook.com>

> 
>  STATIC_NOPV void native_flush_tlb_multi(const struct cpumask *cpumask,
>  					 const struct flush_tlb_info *info)
> 
> 

Re: [PATCH 1/3] x86/mm: Unexport tlb_state_shared

[Dave Hansen]

On 5/20/25 03:55, Peter Zijlstra wrote:
> Never export data; modules have no business being able to change tlb
> state.

Acked-by: Dave Hansen <dave.hansen@linux.intel.com>

Re: [PATCH 2/3] x86/mm: Avoid repeated this_cpu_*() ops in switch_mm_irqs_off()

[Dave Hansen]

On 5/20/25 03:55, Peter Zijlstra wrote:
> Aside from generating slightly better code for not having to use %fs
> prefixed ops, the real purpose is to clarify code by switching some to
> smp_store_release() later on.
> 
> Notably, this_cpu_{read,write}() imply {READ,WRITE}_ONCE().

Acked-by: Dave Hansen <dave.hansen@linux.intel.com>

Re: [PATCH 3/3] x86/mm: Clarify should_flush_tlb() ordering

[Dave Hansen]

On 5/20/25 03:55, Peter Zijlstra wrote:
> The ordering in should_flush_tlb() is entirely non-obvious and is only
> correct because x86 is TSO. Clarify the situation by replacing two
> WRITE_ONCE()s with smp_store_release(), which on x86 is cosmetic.
> 
> Additionally, clarify the comment on should_flush_tlb().

Thanks for clarifying the ordering in those comments. It's much
appreciated by us mere mortals!

Oh, and there are a few we's that snuck into the comments. But whoever
applies this can fix them up.

Acked-by: Dave Hansen <dave.hansen@linux.intel.com>