mm/mmu_notifier: Allow two-pass struct mmu_interval_notifiers

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@intel.com>
Cc: Christian König <christian.koenig@amd.com>
Cc: David Hildenbrand <david@kernel.org>
Cc: Lorenzo Stoakes <lorenzo.stoakes@oracle.com>
Cc: Liam R. Howlett <Liam.Howlett@oracle.com>
Cc: Vlastimil Babka <vbabka@kernel.org>
Cc: Mike Rapoport <rppt@kernel.org>
Cc: Suren Baghdasaryan <surenb@google.com>
Cc: Michal Hocko <mhocko@suse.com>
Cc: Jason Gunthorpe <jgg@ziepe.ca>
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>
Acked-by: David Hildenbrand (Arm) <david@kernel.org>
Reviewed-by: Maarten Lankhorst <maarten.lankhorst@linux.intel.com>
Link: https://patch.msgid.link/20260305093909.43623-2-thomas.hellstrom@linux.intel.com

diff --git a/include/linux/mmu_notifier.h b/include/linux/mmu_notifier.h
index d1094c2d5fb61e5e39dcabb48072964f1f64aee3..b60673a8e0bbbd55eabfc376dc3100ebb0763997 100644 (file)
--- 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 8e0125dc05227909558aaaa9faac3fba00e52b0c..33023dbbd76d8899e8d3e5f545afabfb4fefae3f 100644 (file)
--- 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(
@@ -977,6 +1023,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);