[PATCH 2/3] rbtree: update augmented data on the way down in rb_add_augmented_cached()

[Yiwei Lin <s921975628@gmail.com>]

rb_add_augmented_cached() links the new node and then calls
->propagate() from the parent up to the root to bring the augmented
data of the ancestors up to date.  It marks that as "suboptimal": every
augmented user that open-codes its insertion (interval trees, vmalloc,
drm_mm, gpu buddy, ...) instead folds the new node's value into each
ancestor while walking down, which touches no sibling and needs no
second pass.  ->propagate() has to read both children of every
ancestor until it finds one whose value did not change.

Add a ->merge() callback to struct rb_augment_callbacks that folds the
augmented data of a node being inserted into one of its future
ancestors, and use it on the descent.

The new node's augmented data must describe the node alone before the
call. That was already required, since ->propagate() read it, and both
callers set it up that way.

Measured with rbtree_test's cached augmented test switched to
rb_add_augmented_cached() (next patch) on a Raspberry Pi 4 Model B
(Cortex-A72, 1.8 GHz), nnodes cached insert+delete, median of 5 runs,
arch timer ticks (54 MHz), for reference also the open-coded descent
used by augmented test 1. This brings the helper to the level of
the open-coded descent:

    nnodes    open-coded      before       after
       100           468         486         475  (-2.3%)
     10000        182772      188644      185810  (-1.5%)
    100000       5551237     5700029     5518725  (-3.2%)

Signed-off-by: Yiwei Lin <s921975628@gmail.com>
Assisted-by: LLM
---
 Documentation/core-api/rbtree.rst | 27 ++++++++++++++++++---
 include/linux/rbtree_augmented.h  | 40 ++++++++++++++++++++++++++-----
 kernel/sched/fair.c               | 13 +++++++++-
 3 files changed, 70 insertions(+), 10 deletions(-)

diff --git a/Documentation/core-api/rbtree.rst b/Documentation/core-api/rbtree.rst
index cce80e19087b4..210045f128d65 100644
--- a/Documentation/core-api/rbtree.rst
+++ b/Documentation/core-api/rbtree.rst
@@ -257,8 +257,14 @@ When erasing a node, the user must call rb_erase_augmented() instead of
 rb_erase(). rb_erase_augmented() calls back into user provided functions
 to update the augmented information on affected subtrees.
 
-In both cases, the callbacks are provided through struct rb_augment_callbacks.
-3 callbacks must be defined:
+Alternatively, rb_add_augmented_cached() performs the whole insertion into
+a cached tree: it walks down to the insertion point, merges the new node's
+augmented information into every node on that path, links the node and
+rebalances.  The new node's augmented information must already describe
+the node alone when it is called.
+
+In all cases, the callbacks are provided through struct rb_augment_callbacks.
+3 callbacks must be defined, plus a fourth one for rb_add_augmented_cached():
 
 - A propagation callback, which updates the augmented value for a given
   node and its ancestors, up to a given stop point (or NULL to update
@@ -271,6 +277,10 @@ In both cases, the callbacks are provided through struct rb_augment_callbacks.
   subtree to a newly assigned subtree root AND recomputes the augmented
   information for the former subtree root.
 
+- A merge callback, which folds the augmented value of a node being inserted
+  into the augmented value of one of its future ancestors.  It is only used
+  by rb_add_augmented_cached().
+
 The compiled code for rb_erase_augmented() may inline the propagation and
 copy callbacks, which results in a large function, so each augmented rbtree
 user should have a single rb_erase_augmented() call site in order to limit
@@ -395,8 +405,19 @@ Insertion/removal are defined using the following augmented callbacks::
 	old->__subtree_last = compute_subtree_last(old);
   }
 
+  static void augment_merge(struct rb_node *rb, struct rb_node *rb_new)
+  {
+	struct interval_tree_node *node =
+		rb_entry(rb, struct interval_tree_node, rb);
+	struct interval_tree_node *new =
+		rb_entry(rb_new, struct interval_tree_node, rb);
+
+	if (node->__subtree_last < new->__subtree_last)
+		node->__subtree_last = new->__subtree_last;
+  }
+
   static const struct rb_augment_callbacks augment_callbacks = {
-	augment_propagate, augment_copy, augment_rotate
+	augment_propagate, augment_copy, augment_rotate, augment_merge
   };
 
   void interval_tree_insert(struct interval_tree_node *node,
diff --git a/include/linux/rbtree_augmented.h b/include/linux/rbtree_augmented.h
index d2fa1c41bfd2b..6a1b78f6cb1e2 100644
--- a/include/linux/rbtree_augmented.h
+++ b/include/linux/rbtree_augmented.h
@@ -28,6 +28,7 @@ struct rb_augment_callbacks {
 	void (*propagate)(struct rb_node *node, struct rb_node *stop);
 	void (*copy)(struct rb_node *old, struct rb_node *new);
 	void (*rotate)(struct rb_node *old, struct rb_node *new);
+	void (*merge)(struct rb_node *node, struct rb_node *new);
 };
 
 extern void __rb_insert_augmented(struct rb_node *node, struct rb_root *root,
@@ -60,6 +61,12 @@ rb_insert_augmented_cached(struct rb_node *node,
 	rb_insert_augmented(node, &root->rb_root, augment);
 }
 
+/*
+ * Insert @node into the leftmost cached augmented tree @tree.
+ *
+ * The augmented data of @node must already describe @node alone; it is
+ * merged into every ancestor on the way down through augment->merge().
+ */
 static __always_inline struct rb_node *
 rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree,
 			bool (*less)(struct rb_node *, const struct rb_node *),
@@ -71,6 +78,7 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree,
 
 	while (*link) {
 		parent = *link;
+		augment->merge(parent, node);
 		if (less(node, parent)) {
 			link = &parent->rb_left;
 		} else {
@@ -80,7 +88,6 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree,
 	}
 
 	rb_link_node(node, parent, link);
-	augment->propagate(parent, NULL); /* suboptimal */
 	rb_insert_augmented_cached(node, tree, leftmost, augment);
 
 	return leftmost ? node : NULL;
@@ -95,10 +102,13 @@ rb_add_augmented_cached(struct rb_node *node, struct rb_root_cached *tree,
  * RBFIELD:     name of struct rb_node field within RBSTRUCT
  * RBCOPY:	name of function that copies the RBAUGMENTED datas
  * RBCOMPUTE:   name of function that recomputes the RBAUGMENTED datas
+ * RBMERGE:     name of function that merges a new node's RBAUGMENTED datas
+ *              into an ancestor's, used on the insertion path
  */
 
 #define RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME,			\
-			     RBSTRUCT, RBFIELD, RBCOPY, RBCOMPUTE)	\
+			     RBSTRUCT, RBFIELD, RBCOPY, RBCOMPUTE,	\
+			     RBMERGE)					\
 static inline void							\
 RBNAME ## _propagate(struct rb_node *rb, struct rb_node *stop)		\
 {									\
@@ -124,10 +134,18 @@ RBNAME ## _rotate(struct rb_node *rb_old, struct rb_node *rb_new)	\
 	RBCOPY(new, old);						\
 	RBCOMPUTE(old, false);						\
 }									\
+static inline void							\
+RBNAME ## _merge(struct rb_node *rb, struct rb_node *rb_new)		\
+{									\
+	RBSTRUCT *node = rb_entry(rb, RBSTRUCT, RBFIELD);		\
+	RBSTRUCT *new = rb_entry(rb_new, RBSTRUCT, RBFIELD);		\
+	RBMERGE(node, new);						\
+}									\
 RBSTATIC const struct rb_augment_callbacks RBNAME = {			\
 	.propagate = RBNAME ## _propagate,				\
 	.copy = RBNAME ## _copy,					\
-	.rotate = RBNAME ## _rotate					\
+	.rotate = RBNAME ## _rotate,					\
+	.merge = RBNAME ## _merge					\
 };
 
 /*
@@ -139,17 +157,21 @@ RBSTATIC const struct rb_augment_callbacks RBNAME = {			\
  * RBFIELD:     name of struct rb_node field within RBSTRUCT
  * RBAUGMENTED: name of field within RBSTRUCT holding data for subtree
  * RBCOMPUTE:   name of function that recomputes the RBAUGMENTED data
+ * RBMERGE:     name of function that merges a new node's RBAUGMENTED data
+ *              into an ancestor's, used on the insertion path
  */
 
 #define RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,				\
-			     RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE)	\
+			     RBSTRUCT, RBFIELD, RBAUGMENTED, RBCOMPUTE,	\
+			     RBMERGE)					\
 static inline void							\
 RBNAME ## _copy_single(RBSTRUCT *new, RBSTRUCT *old)			\
 {									\
 	new->RBAUGMENTED = old->RBAUGMENTED;				\
 }									\
 RB_DECLARE_CALLBACKS_MULTI(RBSTATIC, RBNAME,				\
-		     RBSTRUCT, RBFIELD, RBNAME ## _copy_single, RBCOMPUTE)
+		     RBSTRUCT, RBFIELD, RBNAME ## _copy_single,		\
+		     RBCOMPUTE, RBMERGE)
 
 /*
  * Template for declaring augmented rbtree callbacks,
@@ -185,8 +207,14 @@ static inline bool RBNAME ## _compute_max(RBSTRUCT *node, bool exit)	      \
 	node->RBAUGMENTED = max;					      \
 	return false;							      \
 }									      \
+static inline void RBNAME ## _merge_max(RBSTRUCT *node, RBSTRUCT *new)	      \
+{									      \
+	if (node->RBAUGMENTED < new->RBAUGMENTED)			      \
+		node->RBAUGMENTED = new->RBAUGMENTED;			      \
+}									      \
 RB_DECLARE_CALLBACKS(RBSTATIC, RBNAME,					      \
-		     RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max)
+		     RBSTRUCT, RBFIELD, RBAUGMENTED, RBNAME ## _compute_max,  \
+		     RBNAME ## _merge_max)
 
 
 #define	RB_RED		0
diff --git a/kernel/sched/fair.c b/kernel/sched/fair.c
index 7455a83a6a990..60db4624b9897 100644
--- a/kernel/sched/fair.c
+++ b/kernel/sched/fair.c
@@ -1066,9 +1066,20 @@ static inline bool min_vruntime_update(struct sched_entity *se, bool exit)
 	       se->max_slice == old_max_slice;
 }
 
+/*
+ * Fold @new's subtree data into @se, for each @se on @new's insertion path.
+ */
+static inline void
+min_vruntime_merge(struct sched_entity *se, struct sched_entity *new)
+{
+	__min_vruntime_update(se, &new->run_node);
+	__min_slice_update(se, &new->run_node);
+	__max_slice_update(se, &new->run_node);
+}
 
 RB_DECLARE_CALLBACKS_MULTI(static, min_vruntime_cb, struct sched_entity,
-		     run_node, min_vruntime_copy, min_vruntime_update);
+		     run_node, min_vruntime_copy, min_vruntime_update,
+		     min_vruntime_merge);
 
 /*
  * Enqueue an entity into the rb-tree:
-- 
2.34.1