[PATCH v2 1/2] gpu/buddy: find range-biased blocks in the free trees

[William Palacek <William.Palacek@amd.com>]

__alloc_range_bias() walks the block tree from the roots and visits
every split block overlapping [start, end) until it finds a suitable
free block, so a range-restricted allocation costs time linear in the
number of live allocations in the range.

amdgpu pays this on every VRAM allocation when memory is partitioned
(NPS2/NPS4): each BO, including the page-table BOs created when memory
is mapped, is restricted to its partition's range, and KFD requests
cleared VRAM. On MI300X in CPX/NPS4 the compute partitions that share a
memory partition also share the VRAM manager lock, so the walk is
serialised across processes. hipMalloc latency grows linearly with the
number of live allocations, and vLLM piecewise graph capture slows down
with every captured graph.

The free trees are sorted by offset. Before walking, look up the highest
free block that lies wholly inside the range in each free tree from the
requested order up, and narrow the walk to the highest of them: it then
only visits that block's ancestors and keeps the existing split and
error handling. The full-range walk remains for space that is only
available in blocks straddling the range edges.

Placement changes in two cases: a free block straddling the upper range
edge is no longer used ahead of a lower block wholly inside the range,
and for a clear request without GPU_BUDDY_TOPDOWN_ALLOCATION the clear
steering only applies within the chosen block, so a lower clear block
is not preferred over a higher one that is not clear. Otherwise blocks
are placed top-down as before, including with
GPU_BUDDY_TOPDOWN_ALLOCATION.

Filling 1 GiB of a 4 GiB manager with 64 KiB range-biased allocations
that request cleared memory takes 1190 ms before, 3 ms after (KUnit,
UML). On MI300X in CPX/NPS4, with 8000 live 2 MiB hipMalloc allocations
on one partition, the median latency of the last 800 over the first 800
is 1.60x and 4.45x before (two runs), 0.98x to 1.03x after (four runs).

Cc: Matthew Auld <matthew.auld@intel.com>
Cc: Arunpravin Paneer Selvam <Arunpravin.PaneerSelvam@amd.com>
Signed-off-by: William Palacek <William.Palacek@amd.com>
---
 drivers/gpu/buddy.c | 35 ++++++++++++++++++++++++++++++++++-
 1 file changed, 34 insertions(+), 1 deletion(-)

diff --git a/drivers/gpu/buddy.c b/drivers/gpu/buddy.c
index 28cff97d57dc8..8652ae35d323c 100644
--- a/drivers/gpu/buddy.c
+++ b/drivers/gpu/buddy.c
@@ -665,6 +665,26 @@ rbtree_last_clear_free_block(struct rb_root *root,
 	return block;
 }
 
+/* Return the highest free block with lo <= offset <= hi. */
+static struct gpu_buddy_block *
+rbtree_last_free_block_in(struct rb_root *root, u64 lo, u64 hi)
+{
+	struct gpu_buddy_block *block, *last = NULL;
+	struct rb_node *node = root->rb_node;
+
+	while (node) {
+		block = rbtree_get_free_block(node);
+		if (gpu_buddy_block_offset(block) > hi) {
+			node = node->rb_left;
+		} else {
+			last = block;
+			node = node->rb_right;
+		}
+	}
+
+	return last && gpu_buddy_block_offset(last) >= lo ? last : NULL;
+}
+
 static void rbtree_insert(struct gpu_buddy *mm,
 			  struct gpu_buddy_block *block)
 {
@@ -1247,11 +1267,24 @@ __alloc_range_bias(struct gpu_buddy *mm,
 		   unsigned long flags)
 {
 	u64 req_size = mm->chunk_size << order;
-	struct gpu_buddy_block *block;
+	struct gpu_buddy_block *block = NULL, *tmp;
+	unsigned int o;
 	LIST_HEAD(dfs);
 	int err;
 	int i;
 
+	/* Narrow the walk to the highest free block wholly in range (top-down). */
+	for (o = order; o <= mm->max_order && start + (mm->chunk_size << o) <= end; o++) {
+		tmp = rbtree_last_free_block_in(&mm->free_tree[o], start,
+						end - (mm->chunk_size << o));
+		if (tmp && (!block || gpu_buddy_block_offset(tmp) > gpu_buddy_block_offset(block)))
+			block = tmp;
+	}
+	if (block) {
+		start = gpu_buddy_block_offset(block);
+		end = start + gpu_buddy_block_size(mm, block);
+	}
+
 	end = end - 1;
 
 	/*
-- 
2.43.0