Re: [PATCH v6 2/5] lazy tlb: allow lazy tlb mm refcounting to be configurable

[Nadav Amit <nadav.amit@gmail.com>]

On 1/18/23 10:00 AM, Nicholas Piggin wrote:
> Add CONFIG_MMU_TLB_REFCOUNT which enables refcounting of the lazy tlb mm
> when it is context switched. This can be disabled by architectures that
> don't require this refcounting if they clean up lazy tlb mms when the
> last refcount is dropped. Currently this is always enabled, which is
> what existing code does, so the patch is effectively a no-op.
> 
> Rename rq->prev_mm to rq->prev_lazy_mm, because that's what it is.
> 
> Signed-off-by: Nicholas Piggin <npiggin@gmail.com>
> ---
>   Documentation/mm/active_mm.rst |  6 ++++++
>   arch/Kconfig                   | 17 +++++++++++++++++
>   include/linux/sched/mm.h       | 18 +++++++++++++++---
>   kernel/sched/core.c            | 22 ++++++++++++++++++----
>   kernel/sched/sched.h           |  4 +++-
>   5 files changed, 59 insertions(+), 8 deletions(-)
> 
> diff --git a/Documentation/mm/active_mm.rst b/Documentation/mm/active_mm.rst
> index 6f8269c284ed..2b0d08332400 100644
> --- a/Documentation/mm/active_mm.rst
> +++ b/Documentation/mm/active_mm.rst
> @@ -4,6 +4,12 @@
>   Active MM
>   =========
>   
> +Note, the mm_count refcount may no longer include the "lazy" users
> +(running tasks with ->active_mm == mm && ->mm == NULL) on kernels
> +with CONFIG_MMU_LAZY_TLB_REFCOUNT=n. Taking and releasing these lazy
> +references must be done with mmgrab_lazy_tlb() and mmdrop_lazy_tlb()
> +helpers which abstracts this config option.
> +
>   ::
>   
>    List:       linux-kernel
> diff --git a/arch/Kconfig b/arch/Kconfig
> index 12e3ddabac9d..b07d36f08fea 100644
> --- a/arch/Kconfig
> +++ b/arch/Kconfig
> @@ -465,6 +465,23 @@ config ARCH_WANT_IRQS_OFF_ACTIVATE_MM
>   	  irqs disabled over activate_mm. Architectures that do IPI based TLB
>   	  shootdowns should enable this.
>   
> +# Use normal mm refcounting for MMU_LAZY_TLB kernel thread references.
> +# MMU_LAZY_TLB_REFCOUNT=n can improve the scalability of context switching
> +# to/from kernel threads when the same mm is running on a lot of CPUs (a large
> +# multi-threaded application), by reducing contention on the mm refcount.
> +#
> +# This can be disabled if the architecture ensures no CPUs are using an mm as a
> +# "lazy tlb" beyond its final refcount (i.e., by the time __mmdrop frees the mm
> +# or its kernel page tables). This could be arranged by arch_exit_mmap(), or
> +# final exit(2) TLB flush, for example.
> +#
> +# To implement this, an arch *must*:
> +# Ensure the _lazy_tlb variants of mmgrab/mmdrop are used when dropping the
> +# lazy reference of a kthread's ->active_mm (non-arch code has been converted
> +# already).
> +config MMU_LAZY_TLB_REFCOUNT
> +	def_bool y
> +
>   config ARCH_HAVE_NMI_SAFE_CMPXCHG
>   	bool
>   
> diff --git a/include/linux/sched/mm.h b/include/linux/sched/mm.h
> index 5376caf6fcf3..68bbe8d90c2e 100644
> --- a/include/linux/sched/mm.h
> +++ b/include/linux/sched/mm.h
> @@ -82,17 +82,29 @@ static inline void mmdrop_sched(struct mm_struct *mm)
>   /* Helpers for lazy TLB mm refcounting */
>   static inline void mmgrab_lazy_tlb(struct mm_struct *mm)
>   {
> -	mmgrab(mm);
> +	if (IS_ENABLED(CONFIG_MMU_LAZY_TLB_REFCOUNT))
> +		mmgrab(mm);
>   }
>   
>   static inline void mmdrop_lazy_tlb(struct mm_struct *mm)
>   {
> -	mmdrop(mm);
> +	if (IS_ENABLED(CONFIG_MMU_LAZY_TLB_REFCOUNT)) {
> +		mmdrop(mm);
> +	} else {
> +		/*
> +		 * mmdrop_lazy_tlb must provide a full memory barrier, see the
> +		 * membarrier comment finish_task_switch which relies on this.
> +		 */
> +		smp_mb();
> +	}
>   }

Considering the fact that mmdrop_lazy_tlb() replaced mmdrop() in various 
locations in which smp_mb() was not required, this comment might be 
confusing. IOW, for the cases in most cases where mmdrop_lazy_tlb() 
replaced mmdrop(), this comment was irrelevant, and therefore it now 
becomes confusing.

I am not sure the include the smp_mb() here instead of "open-coding" it 
helps.

>   
>   static inline void mmdrop_lazy_tlb_sched(struct mm_struct *mm)
>   {
> -	mmdrop_sched(mm);
> +	if (IS_ENABLED(CONFIG_MMU_LAZY_TLB_REFCOUNT))
> +		mmdrop_sched(mm);
> +	else
> +		smp_mb(); // see above
>   }

Wrong style of comment.