[PATCH 0/2] gpu/buddy: make range-biased allocation independent of live blocks

[PATCH 0/2] gpu/buddy: make range-biased allocation independent of live blocks

[William Palacek]

A range-restricted allocation walks every split block in the range, so
its cost grows linearly with the number of live allocations there.
amdgpu restricts every VRAM allocation to its memory partition under
NPS2/NPS4, so on MI300X in CPX/NPS4 KFD allocation latency grows with
every live BO, under a VRAM manager lock shared by the compute
partitions.

Patch 1 looks up the lowest free block inside the range in the
offset-sorted free trees and restricts the existing walk to that
block's ancestors. Patch 2 adds a KUnit test for the NPS4 pattern
(range-biased, cleared, 64 KiB): about 1 s before, 2 ms after.

This is based on amd-staging-drm-next. drm-misc-next no longer has
separate clear and dirty free trees, so the lookup there needs a
different form and will be handled separately.

William Palacek (2):
  gpu/buddy: find range-biased blocks in the free trees
  gpu/tests/buddy: add range-biased allocation performance test

 drivers/gpu/buddy.c                | 36 ++++++++++++++++++++++++++++++
 drivers/gpu/tests/gpu_buddy_test.c | 25 +++++++++++++++++++++
 2 files changed, 61 insertions(+)

-- 
2.43.0

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

[William Palacek]

__gpu_buddy_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 already sorted by offset. For each order from the
requested one up, look up the lowest free block at or above start,
preferred tree first. If it lies wholly inside the range, run the
existing walk restricted to that block's extent: it only visits the
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.

Behaviour notes:
- the smallest order is tried first, lowest offset within an order;
- a free block of the other clear state inside the range is preferred
  over an edge-straddling block of the requested state, as the existing
  fallback pass already accepts either state;
- when no free block lies wholly inside the range (range nearly full),
  the linear walk still runs.

Filling 1 GiB of a 4 GiB manager with 64 KiB range-biased allocations
that request cleared memory takes about 1 s and 268M block visits
before, 2 ms and 410K visits after.
On MI300X CPX/NPS4, 8000 live 2 MiB hipMalloc allocations: median
latency of the last 800 over the first 800 is 6.2x before, 1.04x after.
vLLM piecewise graph capture of a 70B model with tensor parallelism 64:
57 s for the second graph rising to 113 s for the 17th before, a flat
5-6 s per graph after.

Signed-off-by: William Palacek <William.Palacek@amd.com>
---
 drivers/gpu/buddy.c | 36 ++++++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/drivers/gpu/buddy.c b/drivers/gpu/buddy.c
index a5553fcec28c..3a1fbafe42dc 100644
--- a/drivers/gpu/buddy.c
+++ b/drivers/gpu/buddy.c
@@ -120,6 +120,12 @@ rbtree_last_free_block(struct rb_root *root)
 	return rbtree_get_free_block(rb_last(root));
 }
 
+/* rb_find_first() with this returns the lowest free block at or above *start. */
+static int rbtree_offset_below_cmp(const void *start, const struct rb_node *node)
+{
+	return gpu_buddy_block_offset(rbtree_get_free_block(node)) < *(const u64 *)start;
+}
+
 static bool rbtree_is_empty(struct rb_root *root)
 {
 	return RB_EMPTY_ROOT(root);
@@ -846,9 +852,39 @@ __gpu_buddy_alloc_range_bias(struct gpu_buddy *mm,
 			     unsigned int order,
 			     unsigned long flags)
 {
+	enum gpu_buddy_free_tree tree = (flags & GPU_BUDDY_CLEAR_ALLOCATION) ?
+		GPU_BUDDY_CLEAR_TREE : GPU_BUDDY_DIRTY_TREE;
 	struct gpu_buddy_block *block;
 	bool fallback = false;
+	u64 bstart, bend;
+	unsigned int i;
+	int pass;
+
+	/*
+	 * Walking [start, end) visits every split block in it, linear in the
+	 * live allocations. Instead take the lowest free block wholly inside
+	 * the range from the offset-sorted free trees, preferred tree first;
+	 * the walk restricted to that block only visits its ancestors.
+	 */
+	for (pass = 0; pass < GPU_BUDDY_MAX_FREE_TREES; pass++) {
+		for (i = order; i <= mm->max_order; i++) {
+			block = rbtree_get_free_block(rb_find_first(&start,
+								    &mm->free_trees[tree][i],
+								    rbtree_offset_below_cmp));
+			if (!block)
+				continue;
+
+			bstart = gpu_buddy_block_offset(block);
+			bend = bstart + gpu_buddy_block_size(mm, block);
+			if (bend <= end)
+				return __alloc_range_bias(mm, bstart, bend, order,
+							  flags, true);
+		}
+		tree = (tree == GPU_BUDDY_CLEAR_TREE) ?
+			GPU_BUDDY_DIRTY_TREE : GPU_BUDDY_CLEAR_TREE;
+	}
 
+	/* Only space in blocks straddling the range edges is left. */
 	block = __alloc_range_bias(mm, start, end, order,
 				   flags, fallback);
 	if (IS_ERR(block))
-- 
2.43.0

[PATCH 2/2] gpu/tests/buddy: add range-biased allocation performance test

[William Palacek]

Fill a 1 GiB range of a 4 GiB manager with 64 KiB range-biased
allocations that request cleared memory, as amdgpu does for KFD VRAM
on NPS4 partitions, and report the time taken. Before the previous
change this took about 1 s because every allocation walked all live
blocks in the range; it now takes about 2 ms.

Signed-off-by: William Palacek <William.Palacek@amd.com>
---
 drivers/gpu/tests/gpu_buddy_test.c | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/drivers/gpu/tests/gpu_buddy_test.c b/drivers/gpu/tests/gpu_buddy_test.c
index 4d288f58fb5c..10b4fab2407c 100644
--- a/drivers/gpu/tests/gpu_buddy_test.c
+++ b/drivers/gpu/tests/gpu_buddy_test.c
@@ -289,6 +289,30 @@ static void gpu_test_buddy_fragmentation_performance(struct kunit *test)
 	gpu_buddy_fini(&mm);
 }
 
+static void gpu_test_buddy_range_bias_performance(struct kunit *test)
+{
+	LIST_HEAD(allocated);
+	struct gpu_buddy mm;
+	ktime_t t0;
+	u64 i;
+
+	/*
+	 * Fill one NPS4-like quarter with 64 KiB range-biased allocations
+	 * that ask for cleared memory, as amdgpu does for KFD VRAM.
+	 */
+	KUNIT_ASSERT_FALSE(test, gpu_buddy_init(&mm, SZ_4G, SZ_4K));
+	t0 = ktime_get();
+	for (i = 0; i < SZ_1G / SZ_64K; i++)
+		KUNIT_ASSERT_FALSE(test, gpu_buddy_alloc_blocks(&mm, SZ_1G, SZ_2G, SZ_64K,
+								SZ_64K, &allocated,
+								GPU_BUDDY_RANGE_ALLOCATION |
+								GPU_BUDDY_CLEAR_ALLOCATION));
+	kunit_info(test, "Range-biased fill took %lld ms\n",
+		   ktime_ms_delta(ktime_get(), t0));
+	gpu_buddy_free_list(&mm, &allocated, 0);
+	gpu_buddy_fini(&mm);
+}
+
 static void gpu_test_buddy_dirty_tracker_performance(struct kunit *test)
 {
 	struct gpu_buddy_block *block, *tmp;
@@ -1612,6 +1636,7 @@ static struct kunit_case gpu_buddy_tests[] = {
 	KUNIT_CASE(gpu_test_buddy_alloc_range_bias),
 	KUNIT_CASE_SLOW(gpu_test_buddy_fragmentation_performance),
 	KUNIT_CASE_SLOW(gpu_test_buddy_dirty_tracker_performance),
+	KUNIT_CASE_SLOW(gpu_test_buddy_range_bias_performance),
 	KUNIT_CASE(gpu_test_buddy_alloc_exceeds_max_order),
 	KUNIT_CASE(gpu_test_buddy_offset_aligned_allocation),
 	KUNIT_CASE(gpu_test_buddy_subtree_offset_alignment_stress),
-- 
2.43.0