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

["Thomas Hellström" <thomas.hellstrom@linux.intel.com>]

Hi, Matt,

On Mon, 2026-03-02 at 11:48 -0800, Matthew Brost wrote:

Thanks for reviewing.

> On Mon, Mar 02, 2026 at 05:32:45PM +0100, Thomas Hellström wrote:
> > 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.
> > 
> > Cc: Jason Gunthorpe <jgg@ziepe.ca>
> 
> I thought Jason had given a RB on previous revs - did you drop
> because
> enough has changed?

Yes. In particular the inclusion of 
invalidate_finish() in the ops, although IIRC that
was actually suggested by Jason.

>  
> > Cc: Andrew Morton <akpm@linux-foundation.org>
> > Cc: Simona Vetter <simona.vetter@ffwll.ch>
> > Cc: Dave Airlie <airlied@gmail.com>
> > Cc: Alistair Popple <apopple@nvidia.com>
> > Cc: <dri-devel@lists.freedesktop.org>
> > Cc: <linux-mm@kvack.org>
> > Cc: <linux-kernel@vger.kernel.org>
> > 
> > Assisted-by: GitHub Copilot:claude-sonnet-4.6 # Documentation only.
> > Signed-off-by: Thomas Hellström <thomas.hellstrom@linux.intel.com>
> > ---
> >  include/linux/mmu_notifier.h | 38 +++++++++++++++++++++
> >  mm/mmu_notifier.c            | 64 +++++++++++++++++++++++++++++++-
> > ----
> >  2 files changed, 93 insertions(+), 9 deletions(-)
> > 
> > diff --git a/include/linux/mmu_notifier.h
> > b/include/linux/mmu_notifier.h
> > index 07a2bbaf86e9..de0e742ea808 100644
> > --- a/include/linux/mmu_notifier.h
> > +++ b/include/linux/mmu_notifier.h
> > @@ -233,16 +233,54 @@ struct mmu_notifier {
> >  	unsigned int users;
> >  };
> >  
> > +/**
> > + * struct mmu_interval_notifier_finish - mmu_interval_notifier
> > two-pass abstraction
> > + * @link: List link for the notifiers pending pass list
> 
> Lockless list?

Sure, Can add that.

> 
> > + * @notifier: The mmu_interval_notifier for which the finish pass
> > is called.
> > + *
> > + * Allocate, typically using GFP_NOWAIT in the interval notifier's
> > first pass.
> > + * If allocation fails (which is not unlikely under memory
> > pressure), fall back
> > + * to single-pass operation. 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 two
> > 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.
> > + */
> > +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 final pass.
> > + *                    If @finish is %NULL on return, then no final
> > pass will be
> > + *                    called.
> > + * @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);
> 
> Should we complain somewhere if a caller registers a notifier with
> invalidate_start set but not invalidate_finish?

Good idea. I'll update.

Thanks,
Thomas


> 
> Matt
> 
> >  };
> >  
> >  struct mmu_interval_notifier {
> > diff --git a/mm/mmu_notifier.c b/mm/mmu_notifier.c
> > index a6cdf3674bdc..38acd5ef8eb0 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);
> > +			}
> > +
> > +		} 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(
> > -- 
> > 2.53.0
> >