Re: [PATCH v4 1/4] mm/mmu_notifier: Allow two-pass struct mmu_interval_notifiers

[Maarten Lankhorst]

Hey,


Den 2026-03-05 kl. 10:39, skrev Thomas Hellström:
> GPU use-cases for mmu_interval_notifiers with hmm often involve
> starting a gpu operation and then waiting for it to complete.
> These operations are typically context preemption or TLB flushing.
>
> With single-pass notifiers per GPU this doesn't scale in
> multi-gpu scenarios. In those scenarios we'd want to first start
> preemption- or TLB flushing on all GPUs and as a second pass wait
> for them to complete.
>
> One can do this on per-driver basis multiplexing per-driver
> notifiers but that would mean sharing the notifier "user" lock
> across all GPUs and that doesn't scale well either, so adding support
> for multi-pass in the core appears to be the right choice.
>
> Implement two-pass capability in the mmu_interval_notifier. Use a
> linked list for the final passes to minimize the impact for
> use-cases that don't need the multi-pass functionality by avoiding
> a second interval tree walk, and to be able to easily pass data
> between the two passes.
>
> v1:
> - Restrict to two passes (Jason Gunthorpe)
> - Improve on documentation (Jason Gunthorpe)
> - Improve on function naming (Alistair Popple)
> v2:
> - Include the invalidate_finish() callback in the
>   struct mmu_interval_notifier_ops.
> - Update documentation (GitHub Copilot:claude-sonnet-4.6)
> - Use lockless list for list management.
> v3:
> - Update kerneldoc for the struct mmu_interval_notifier_finish::list member
>   (Matthew Brost)
> - Add a WARN_ON_ONCE() checking for NULL invalidate_finish() op if
>   if invalidate_start() is non-NULL. (Matthew Brost)
> v4:
> - Addressed documentation review comments by David Hildenbrand.
>
> Cc: Matthew Brost <matthew.brost@xxxxxxxxx>
> Cc: Christian König <christian.koenig@xxxxxxx>
> Cc: David Hildenbrand <david@xxxxxxxxxx>
> Cc: Lorenzo Stoakes <lorenzo.stoakes@xxxxxxxxxx>
> Cc: Liam R. Howlett <Liam.Howlett@xxxxxxxxxx>
> Cc: Vlastimil Babka <vbabka@xxxxxxxxxx>
> Cc: Mike Rapoport <rppt@xxxxxxxxxx>
> Cc: Suren Baghdasaryan <surenb@xxxxxxxxxx>
> Cc: Michal Hocko <mhocko@xxxxxxxx>
> Cc: Jason Gunthorpe <jgg@xxxxxxxx>
> Cc: Andrew Morton <akpm@xxxxxxxxxxxxxxxxxxxx>
> Cc: Simona Vetter <simona.vetter@xxxxxxxx>
> Cc: Dave Airlie <airlied@xxxxxxxxx>
> Cc: Alistair Popple <apopple@xxxxxxxxxx>
> Cc: <dri-devel@xxxxxxxxxxxxxxxxxxxxx>
> Cc: <linux-mm@xxxxxxxxx>
> Cc: <linux-kernel@xxxxxxxxxxxxxxx>
>
> Assisted-by: GitHub Copilot:claude-sonnet-4.6 # Documentation only.
> Signed-off-by: Thomas Hellström <thomas.hellstrom@xxxxxxxxxxxxxxx>
> ---
>  include/linux/mmu_notifier.h | 42 +++++++++++++++++++++++
>  mm/mmu_notifier.c            | 65 +++++++++++++++++++++++++++++++-----
>  2 files changed, 98 insertions(+), 9 deletions(-)
>
> diff --git a/include/linux/mmu_notifier.h b/include/linux/mmu_notifier.h
> index 07a2bbaf86e9..dcdfdf1e0b39 100644
> --- a/include/linux/mmu_notifier.h
> +++ b/include/linux/mmu_notifier.h
> @@ -233,16 +233,58 @@ struct mmu_notifier {
>  	unsigned int users;
>  };
>  
> +/**
> + * struct mmu_interval_notifier_finish - mmu_interval_notifier two-pass abstraction
> + * @link: Lockless list link for the notifiers pending pass list
> + * @notifier: The mmu_interval_notifier for which the finish pass is called.
> + *
> + * Allocate, typically using GFP_NOWAIT in the interval notifier's start pass.
> + * Note that with a large number of notifiers implementing two passes,
> + * allocation with GFP_NOWAIT will become increasingly likely to fail, so consider
> + * implementing a small pool instead of using kmalloc() allocations.
> + *
> + * If the implementation needs to pass data between the start and the finish passes,
> + * the recommended way is to embed struct mmu_interval_notifier_finish into a larger
> + * structure that also contains the data needed to be shared. Keep in mind that
> + * a notifier callback can be invoked in parallel, and each invocation needs its
> + * own struct mmu_interval_notifier_finish.
> + *
> + * If allocation fails, then the &mmu_interval_notifier_ops->invalidate_start op
> + * needs to implements the full notifier functionality. Please refer to its
> + * documentation.
> + */
> +struct mmu_interval_notifier_finish {
> +	struct llist_node link;
> +	struct mmu_interval_notifier *notifier;
> +};
> +
>  /**
>   * struct mmu_interval_notifier_ops
>   * @invalidate: Upon return the caller must stop using any SPTEs within this
>   *              range. This function can sleep. Return false only if sleeping
>   *              was required but mmu_notifier_range_blockable(range) is false.
> + * @invalidate_start: Similar to @invalidate, but intended for two-pass notifier
> + *                    callbacks where the call to @invalidate_start is the first
> + *                    pass and any struct mmu_interval_notifier_finish pointer
> + *                    returned in the @finish parameter describes the finish pass.
> + *                    If *@finish is %NULL on return, then no final pass will be
> + *                    called, and @invalidate_start needs to implement the full
> + *                    notifier, behaving like @invalidate. The value of *@finish
> + *                    is guaranteed to be %NULL at function entry.
> + * @invalidate_finish: Called as the second pass for any notifier that returned
> + *                     a non-NULL *@finish from @invalidate_start. The @finish
> + *                     pointer passed here is the same one returned by
> + *                     @invalidate_start.
>   */
>  struct mmu_interval_notifier_ops {
>  	bool (*invalidate)(struct mmu_interval_notifier *interval_sub,
>  			   const struct mmu_notifier_range *range,
>  			   unsigned long cur_seq);
> +	bool (*invalidate_start)(struct mmu_interval_notifier *interval_sub,
> +				 const struct mmu_notifier_range *range,
> +				 unsigned long cur_seq,
> +				 struct mmu_interval_notifier_finish **finish);
> +	void (*invalidate_finish)(struct mmu_interval_notifier_finish *finish);
>  };
>  
>  struct mmu_interval_notifier {
> diff --git a/mm/mmu_notifier.c b/mm/mmu_notifier.c
> index a6cdf3674bdc..4d8a64ce8eda 100644
> --- a/mm/mmu_notifier.c
> +++ b/mm/mmu_notifier.c
> @@ -260,6 +260,15 @@ mmu_interval_read_begin(struct mmu_interval_notifier *interval_sub)
>  }
>  EXPORT_SYMBOL_GPL(mmu_interval_read_begin);
>  
> +static void mn_itree_finish_pass(struct llist_head *finish_passes)
> +{
> +	struct llist_node *first = llist_reverse_order(__llist_del_all(finish_passes));
> +	struct mmu_interval_notifier_finish *f, *next;
> +
> +	llist_for_each_entry_safe(f, next, first, link)
> +		f->notifier->ops->invalidate_finish(f);
> +}
> +
>  static void mn_itree_release(struct mmu_notifier_subscriptions *subscriptions,
>  			     struct mm_struct *mm)
>  {
> @@ -271,6 +280,7 @@ static void mn_itree_release(struct mmu_notifier_subscriptions *subscriptions,
>  		.end = ULONG_MAX,
>  	};
>  	struct mmu_interval_notifier *interval_sub;
> +	LLIST_HEAD(finish_passes);
>  	unsigned long cur_seq;
>  	bool ret;
>  
> @@ -278,11 +288,27 @@ static void mn_itree_release(struct mmu_notifier_subscriptions *subscriptions,
>  		     mn_itree_inv_start_range(subscriptions, &range, &cur_seq);
>  	     interval_sub;
>  	     interval_sub = mn_itree_inv_next(interval_sub, &range)) {
> -		ret = interval_sub->ops->invalidate(interval_sub, &range,
> -						    cur_seq);
> +		if (interval_sub->ops->invalidate_start) {
> +			struct mmu_interval_notifier_finish *finish = NULL;
> +
> +			ret = interval_sub->ops->invalidate_start(interval_sub,
> +								  &range,
> +								  cur_seq,
> +								  &finish);
> +			if (ret && finish) {
> +				finish->notifier = interval_sub;
> +				__llist_add(&finish->link, &finish_passes);
> +			}
Should we warn if !ret && finish?

Anyway, looks good either way.

Reviewed-by: Maarten Lankhorst <maarten.lankhorst@xxxxxxxxxxxxxxx>


> +		} else {
> +			ret = interval_sub->ops->invalidate(interval_sub,
> +							    &range,
> +							    cur_seq);
> +		}
>  		WARN_ON(!ret);
>  	}
>  
> +	mn_itree_finish_pass(&finish_passes);
>  	mn_itree_inv_end(subscriptions);
>  }
>  
> @@ -430,7 +456,9 @@ static int mn_itree_invalidate(struct mmu_notifier_subscriptions *subscriptions,
>  			       const struct mmu_notifier_range *range)
>  {
>  	struct mmu_interval_notifier *interval_sub;
> +	LLIST_HEAD(finish_passes);
>  	unsigned long cur_seq;
> +	int err = 0;
>  
>  	for (interval_sub =
>  		     mn_itree_inv_start_range(subscriptions, range, &cur_seq);
> @@ -438,23 +466,41 @@ static int mn_itree_invalidate(struct mmu_notifier_subscriptions *subscriptions,
>  	     interval_sub = mn_itree_inv_next(interval_sub, range)) {
>  		bool ret;
>  
> -		ret = interval_sub->ops->invalidate(interval_sub, range,
> -						    cur_seq);
> +		if (interval_sub->ops->invalidate_start) {
> +			struct mmu_interval_notifier_finish *finish = NULL;
> +
> +			ret = interval_sub->ops->invalidate_start(interval_sub,
> +								  range,
> +								  cur_seq,
> +								  &finish);
> +			if (ret && finish) {
> +				finish->notifier = interval_sub;
> +				__llist_add(&finish->link, &finish_passes);
> +			}
> +
> +		} else {
> +			ret = interval_sub->ops->invalidate(interval_sub,
> +							    range,
> +							    cur_seq);
> +		}
>  		if (!ret) {
>  			if (WARN_ON(mmu_notifier_range_blockable(range)))
>  				continue;
> -			goto out_would_block;
> +			err = -EAGAIN;
> +			break;
>  		}
>  	}
> -	return 0;
>  
> -out_would_block:
> +	mn_itree_finish_pass(&finish_passes);
> +
>  	/*
>  	 * On -EAGAIN the non-blocking caller is not allowed to call
>  	 * invalidate_range_end()
>  	 */
> -	mn_itree_inv_end(subscriptions);
> -	return -EAGAIN;
> +	if (err)
> +		mn_itree_inv_end(subscriptions);
> +
> +	return err;
>  }
>  
>  static int mn_hlist_invalidate_range_start(
> @@ -976,6 +1022,7 @@ int mmu_interval_notifier_insert(struct mmu_interval_notifier *interval_sub,
>  	struct mmu_notifier_subscriptions *subscriptions;
>  	int ret;
>  
> +	WARN_ON_ONCE(ops->invalidate_start && !ops->invalidate_finish);
>  	might_lock(&mm->mmap_lock);
>  
>  	subscriptions = smp_load_acquire(&mm->notifier_subscriptions);