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

Borislav Petkov
2025-02-24T14:16:33Z

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.