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

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

[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 narrows the walk to the highest free block wholly inside the
range, found in the offset-sorted free trees. Patch 2 adds a KUnit test
for the NPS4 pattern (range-biased, cleared, top-down, 64 KiB) that
also checks placement: 1190 ms before, 3 ms after.

v2:
- rebased onto drm-misc-next (single free tree) (Arun)
- take the highest contained block, keeping the walk's top-down
  placement, including for GPU_BUDDY_TOPDOWN_ALLOCATION (Arun)
- the test checks top-down placement
- placement differs from the full walk only for blocks straddling the
  upper range edge and for clear steering across blocks (see patch 1)

v1: https://patchwork.freedesktop.org/series/176957/

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                | 35 +++++++++++++++++++++++++++++-
 drivers/gpu/tests/gpu_buddy_test.c | 27 +++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)


base-commit: 6f7098738fff3ec7393f924343188cb8e6ebc5d2
-- 
2.43.0

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

[William Palacek]

__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

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

[William Palacek]

Fill a 1 GiB range of a 4 GiB manager top-down with 64 KiB range-biased
allocations that request cleared memory, as amdgpu does for KFD VRAM on
NPS4 partitions, check that each one is placed directly below the
previous one, and report the time taken: 1190 ms with the full-range
walk, 3 ms with the free tree lookup.

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/tests/gpu_buddy_test.c | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/drivers/gpu/tests/gpu_buddy_test.c b/drivers/gpu/tests/gpu_buddy_test.c
index 2c6f1da1a6d1c..ef4176b4baa39 100644
--- a/drivers/gpu/tests/gpu_buddy_test.c
+++ b/drivers/gpu/tests/gpu_buddy_test.c
@@ -283,6 +283,32 @@ 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)
+{
+	struct gpu_buddy_block *block;
+	LIST_HEAD(allocated);
+	struct gpu_buddy mm;
+	ktime_t t0;
+	u64 i;
+
+	/* Fill one NPS4-like quarter top-down with cleared 64 KiB allocations. */
+	KUNIT_ASSERT_FALSE(test, gpu_buddy_init(&mm, SZ_4G, SZ_4K));
+	t0 = ktime_get();
+	for (i = 1; 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 |
+								GPU_BUDDY_TOPDOWN_ALLOCATION));
+		block = list_last_entry(&allocated, struct gpu_buddy_block, link);
+		KUNIT_ASSERT_EQ(test, gpu_buddy_block_offset(block), SZ_2G - i * SZ_64K);
+	}
+	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;
@@ -1617,6 +1643,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