Re: [PATCH v13 07/14] x86/mm: global ASID allocation helper functions

[Borislav Petkov <bp@alien8.de>]

On Sun, Feb 23, 2025 at 02:48:57PM -0500, Rik van Riel wrote:
> Subject: Re: [PATCH v13 07/14] x86/mm: global ASID allocation helper functions

This subject needs a verb. I.e.,

"Add global ..."

> Functions to manage global ASID space. 

Ditto.

> Multithreaded processes that
> are simultaneously active on 4 or more CPUs can get a global ASID,
> resulting in the same PCID being used for that process on every CPU.
> 
> This in turn will allow the kernel to use hardware-assisted TLB flushing
> through AMD INVLPGB or Intel RAR for these processes.
> 
> Helper functions split out by request.

No need for that sentence.

> Signed-off-by: Rik van Riel <riel@surriel.com>
> Reviewed-by: Nadav Amit <nadav.amit@gmail.com>
> Tested-by: Manali Shukla <Manali.Shukla@amd.com>
> Tested-by: Brendan Jackman <jackmanb@google.com>
> Tested-by: Michael Kelley <mhklinux@outlook.com>
> ---
>  arch/x86/include/asm/mmu.h      |  11 +++
>  arch/x86/include/asm/tlbflush.h |  43 ++++++++++
>  arch/x86/mm/tlb.c               | 144 +++++++++++++++++++++++++++++++-
>  3 files changed, 195 insertions(+), 3 deletions(-)

arch/x86/mm/tlb.c:378:6: warning: no previous prototype for ‘destroy_context_free_global_asid’ [-Wmissing-prototypes]
  378 | void destroy_context_free_global_asid(struct mm_struct *mm)
      |      ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
arch/x86/mm/tlb.c:355:13: warning: ‘use_global_asid’ defined but not used [-Wunused-function]
  355 | static void use_global_asid(struct mm_struct *mm)
      |             ^~~~~~~~~~~~~~~
arch/x86/mm/tlb.c:327:13: warning: ‘mm_active_cpus_exceeds’ defined but not used [-Wunused-function]
  327 | static bool mm_active_cpus_exceeds(struct mm_struct *mm, int threshold)
      |             ^~~~~~~~~~~~~~~~~~~~~~

If those functions are going to remain global they better get a proper prefix
like "tlb_".

And add the functions with their first use - no need to pre-add them.

> diff --git a/arch/x86/include/asm/mmu.h b/arch/x86/include/asm/mmu.h
> index 3b496cdcb74b..edb5942d4829 100644
> --- a/arch/x86/include/asm/mmu.h
> +++ b/arch/x86/include/asm/mmu.h
> @@ -69,6 +69,17 @@ typedef struct {
>  	u16 pkey_allocation_map;
>  	s16 execute_only_pkey;
>  #endif
> +
> +#ifdef CONFIG_X86_BROADCAST_TLB_FLUSH
> +	/*
> +	 * The global ASID will be a non-zero value when the process has
> +	 * the same ASID across all CPUs, allowing it to make use of
> +	 * hardware-assisted remote TLB invalidation like AMD INVLPGB.
> +	 */
> +	u16 global_asid;
> +	/* The process is transitioning to a new global ASID number. */
> +	bool asid_transition;
> +#endif
>  } mm_context_t;
>  
>  #define INIT_MM_CONTEXT(mm)						\
> diff --git a/arch/x86/include/asm/tlbflush.h b/arch/x86/include/asm/tlbflush.h
> index 09463a2fb05f..83f1da2f1e4a 100644
> --- a/arch/x86/include/asm/tlbflush.h
> +++ b/arch/x86/include/asm/tlbflush.h
> @@ -6,6 +6,7 @@
>  #include <linux/mmu_notifier.h>
>  #include <linux/sched.h>
>  
> +#include <asm/barrier.h>
>  #include <asm/processor.h>
>  #include <asm/cpufeature.h>
>  #include <asm/special_insns.h>
> @@ -234,6 +235,48 @@ void flush_tlb_one_kernel(unsigned long addr);
>  void flush_tlb_multi(const struct cpumask *cpumask,
>  		      const struct flush_tlb_info *info);
>  
> +static inline bool is_dyn_asid(u16 asid)
> +{
> +	return asid < TLB_NR_DYN_ASIDS;
> +}
> +
> +#ifdef CONFIG_X86_BROADCAST_TLB_FLUSH
> +static inline u16 mm_global_asid(struct mm_struct *mm)
> +{
> +	u16 asid;
> +
> +	if (!cpu_feature_enabled(X86_FEATURE_INVLPGB))
> +		return 0;
> +
> +	asid = smp_load_acquire(&mm->context.global_asid);
> +
> +	/* mm->context.global_asid is either 0, or a global ASID */
> +	VM_WARN_ON_ONCE(asid && is_dyn_asid(asid));
> +
> +	return asid;
> +}
> +
> +static inline void assign_mm_global_asid(struct mm_struct *mm, u16 asid)

mm_assign_global_asid()

> +{
> +	/*
> +	 * Notably flush_tlb_mm_range() -> broadcast_tlb_flush() ->
> +	 * finish_asid_transition() needs to observe asid_transition = true
> +	 * once it observes global_asid.
> +	 */
> +	mm->context.asid_transition = true;
> +	smp_store_release(&mm->context.global_asid, asid);
> +}
> +#else
> +static inline u16 mm_global_asid(struct mm_struct *mm)
> +{
> +	return 0;
> +}
> +
> +static inline void assign_mm_global_asid(struct mm_struct *mm, u16 asid)
> +{
> +}
> +#endif
> +
>  #ifdef CONFIG_PARAVIRT
>  #include <asm/paravirt.h>
>  #endif
> diff --git a/arch/x86/mm/tlb.c b/arch/x86/mm/tlb.c
> index 16839651f67f..405630479b90 100644
> --- a/arch/x86/mm/tlb.c
> +++ b/arch/x86/mm/tlb.c
> @@ -74,13 +74,15 @@
>   * use different names for each of them:
>   *
>   * ASID  - [0, TLB_NR_DYN_ASIDS-1]
> - *         the canonical identifier for an mm
> + *         the canonical identifier for an mm, dynamically allocated on each CPU
> + *         [TLB_NR_DYN_ASIDS, MAX_ASID_AVAILABLE-1]
> + *         the canonical, global identifier for an mm, identical across all CPUs
>   *
> - * kPCID - [1, TLB_NR_DYN_ASIDS]
> + * kPCID - [1, MAX_ASID_AVAILABLE]
>   *         the value we write into the PCID part of CR3; corresponds to the
>   *         ASID+1, because PCID 0 is special.
>   *
> - * uPCID - [2048 + 1, 2048 + TLB_NR_DYN_ASIDS]
> + * uPCID - [2048 + 1, 2048 + MAX_ASID_AVAILABLE]
>   *         for KPTI each mm has two address spaces and thus needs two
>   *         PCID values, but we can still do with a single ASID denomination
>   *         for each mm. Corresponds to kPCID + 2048.
> @@ -251,6 +253,142 @@ static void choose_new_asid(struct mm_struct *next, u64 next_tlb_gen,
>  	*need_flush = true;
>  }
>  
> +/*
> + * Global ASIDs are allocated for multi-threaded processes that are
> + * active on multiple CPUs simultaneously, giving each of those
> + * processes the same PCIDs on every CPU, for use with hardware-assisted

"the same PCID" or "the same PCIDs"?

Does a multithreaded process get one or more than one PCIDs? I'd expect only
one ofc.

> + * TLB shootdown on remote CPUs, like AMD INVLPGB or Intel RAR.
> + *
> + * These global ASIDs are held for the lifetime of the process.
> + */
> +static DEFINE_RAW_SPINLOCK(global_asid_lock);
> +static u16 last_global_asid = MAX_ASID_AVAILABLE;
> +static DECLARE_BITMAP(global_asid_used, MAX_ASID_AVAILABLE);
> +static DECLARE_BITMAP(global_asid_freed, MAX_ASID_AVAILABLE);
> +static int global_asid_available = MAX_ASID_AVAILABLE - TLB_NR_DYN_ASIDS - 1;
> +
> +/*
> + * When the search for a free ASID in the global ASID space reaches
> + * MAX_ASID_AVAILABLE, a global TLB flush guarantees that previously
> + * freed global ASIDs are safe to re-use.
> + *
> + * This way the global flush only needs to happen at ASID rollover
> + * time, and not at ASID allocation time.
> + */
> +static void reset_global_asid_space(void)
> +{
> +	lockdep_assert_held(&global_asid_lock);
> +
> +	invlpgb_flush_all_nonglobals();
> +
> +	/*
> +	 * The TLB flush above makes it safe to re-use the previously
> +	 * freed global ASIDs.
> +	 */
> +	bitmap_andnot(global_asid_used, global_asid_used,
> +			global_asid_freed, MAX_ASID_AVAILABLE);
> +	bitmap_clear(global_asid_freed, 0, MAX_ASID_AVAILABLE);
> +
> +	/* Restart the search from the start of global ASID space. */
> +	last_global_asid = TLB_NR_DYN_ASIDS;
> +}
> +
> +static u16 allocate_global_asid(void)
> +{
> +	u16 asid;
> +
> +	lockdep_assert_held(&global_asid_lock);
> +
> +	/* The previous allocation hit the edge of available address space */
> +	if (last_global_asid >= MAX_ASID_AVAILABLE - 1)
> +		reset_global_asid_space();
> +
> +	asid = find_next_zero_bit(global_asid_used, MAX_ASID_AVAILABLE, last_global_asid);
> +
> +	if (asid >= MAX_ASID_AVAILABLE) {

	if (asid >= MAX_ASID_AVAILABLE && !global_asid_available)

> +		/* This should never happen. */
> +		VM_WARN_ONCE(1, "Unable to allocate global ASID despite %d available\n",
> +				global_asid_available);
> +		return 0;
> +	}
> +
> +	/* Claim this global ASID. */
> +	__set_bit(asid, global_asid_used);
> +	last_global_asid = asid;
> +	global_asid_available--;
> +	return asid;
> +}
> +
> +/*
> + * Check whether a process is currently active on more than "threshold" CPUs.

@threshold - then it is clear that you mean the function argument.

> + * This is a cheap estimation on whether or not it may make sense to assign
> + * a global ASID to this process, and use broadcast TLB invalidation.
> + */
> +static bool mm_active_cpus_exceeds(struct mm_struct *mm, int threshold)
> +{
> +	int count = 0;
> +	int cpu;

Is this function supposed to hold some sort of a lock?

> +
> +	/* This quick check should eliminate most single threaded programs. */
> +	if (cpumask_weight(mm_cpumask(mm)) <= threshold)
> +		return false;
> +
> +	/* Slower check to make sure. */
> +	for_each_cpu(cpu, mm_cpumask(mm)) {

Needs CPU hotplug protection?

> +		/* Skip the CPUs that aren't really running this process. */
> +		if (per_cpu(cpu_tlbstate.loaded_mm, cpu) != mm)
> +			continue;
> +
> +		if (per_cpu(cpu_tlbstate_shared.is_lazy, cpu))
> +			continue;
> +
> +		if (++count > threshold)
> +			return true;
> +	}
> +	return false;
> +}
> +
> +/*
> + * Assign a global ASID to the current process, protecting against
> + * races between multiple threads in the process.
> + */
> +static void use_global_asid(struct mm_struct *mm)
> +{
> +	u16 asid;
> +
> +	guard(raw_spinlock_irqsave)(&global_asid_lock);
> +
> +	/* This process is already using broadcast TLB invalidation. */
> +	if (mm_global_asid(mm))
> +		return;
> +
> +	/* The last global ASID was consumed while waiting for the lock. */
> +	if (!global_asid_available) {
> +		VM_WARN_ONCE(1, "Ran out of global ASIDs\n");

And? Why do we want to warn here? We need to reset global ASIDs anyway.

> +		return;
> +	}
> +
> +	asid = allocate_global_asid();
> +	if (!asid)
> +		return;
> +
> +	assign_mm_global_asid(mm, asid);
> +}
> +
> +void destroy_context_free_global_asid(struct mm_struct *mm)

That name is a mouthful. mm_free_global_asid() is just fine.

-- 
Regards/Gruss,
    Boris.

https://people.kernel.org/tglx/notes-about-netiquette