Re: [PATCH v7 1/2] gpu/buddy: replace dual-tree/force_merge with decoupled dirty tracker

Arunpravin Paneer Selvam <[email protected]> Thu, 6 Aug 2026 10:20:43 +0530
Newsgroups org.freedesktop.lists.amd-gfx,org.freedesktop.lists.dri-devel,org.freedesktop.lists.intel-gfx,org.freedesktop.lists.intel-xe
Message-ID <[email protected]>
Adding Anand, as he is working directly with the customer who is waiting 
for this fix to be upstreamed.

@Matthew just checking whether you have had a chance to evaluate the 
changes on your hardware.
If you notice any regressions or issues, please let me know.

Regards,
Arun.

On 8/3/2026 1:54 PM, Arunpravin Paneer Selvam wrote:
> Hi Matthew,
>
> I have posted v7 with the Sashiko's review feedback addressed along 
> with few additional changes.
> Could you review the patches ?
>
> Regards,
> Arun.
>
> On 7/31/2026 12:37 PM, Arunpravin Paneer Selvam wrote:
>> The current buddy allocator maintains separate clear_tree[] and
>> dirty_tree[] rbtrees per order, preventing coalescing between cleared
>> and dirty buddies. Under mixed workloads, this creates a merge barrier:
>> adjacent buddies frequently end up split across trees, forcing reliance
>> on __force_merge() during allocation.
>>
>> __force_merge() performs an O(N x max_order) scan under the VRAM manager
>> lock, leading to allocation stalls and failures for large contiguous
>> requests even when sufficient total free memory is available.
>>
>> Solution
>>
>> Replace the dual-tree design with:
>> - A single free_tree[order] rbtree for dirty and mixed free blocks
>>    (fully cleared free blocks float outside this tree)
>> - A lightweight out-of-band dirty tracker (gpu_dirty_tracker)
>>
>> Fully cleared free blocks are tracked outside the buddy trees using an
>> augmented interval rbtree, enabling O(log E) lookup of the largest
>> cleared extents.
>>
>> Buddy coalescing is now unconditional in __gpu_buddy_free(), regardless
>> of clear/dirty state. This removes the merge barrier and eliminates the
>> need for __force_merge().
>>
>> Benefits
>>
>> - Correct high-order allocations after mixed clear/dirty workloads
>> - Elimination of O(N x max_order) merge cost from the allocation path
>> - O(log E) cleared-extent lookup replacing O(N) scans
>> - Predictable allocation latency under fragmentation
>> - Reduced complexity with a single tree per order
>>
>> Test:
>> dEQP-VK.memory.allocation.basic.size_8KiB.reverse.count_4000
>>
>> Below data is from /sys/kernel/debug/dri/1/amdgpu_vram_mm:
>>
>> Base (dual-tree), before VKCTS test:
>>    order- 6 free:   6 MiB,  blocks: 26
>>    order- 5 free:   1 MiB,  blocks: 15
>>    order- 4 free: 960 KiB,  blocks: 15
>>    order- 3 free:   5 MiB,  blocks: 171
>>    order- 2 free:   2 MiB,  blocks: 176
>>    order- 1 free:   1 MiB,  blocks: 165
>>    order- 0 free:  16 KiB,  blocks: 4
>>
>> Base (dual-tree), after VKCTS test:
>>    order- 6 free: 768 KiB,  blocks: 3
>>    order- 5 free: 499 MiB,  blocks: 3999
>>    order- 4 free: 250 MiB,  blocks: 4001
>>    order- 3 free: 129 MiB,  blocks: 4157
>>    order- 2 free:  65 MiB,  blocks: 4161
>>    order- 1 free:  63 MiB,  blocks: 8138
>>    order- 0 free:  20 KiB,  blocks: 5
>>
>> Dirty tracker, before VKCTS test:
>>    order- 6 free:   4 MiB,  blocks: 19
>>    order- 5 free:   2 MiB,  blocks: 18
>>    order- 4 free: 704 KiB,  blocks: 11
>>    order- 3 free:   5 MiB,  blocks: 168
>>    order- 2 free:   2 MiB,  blocks: 174
>>    order- 1 free:   1 MiB,  blocks: 167
>>    order- 0 free:  32 KiB,  blocks: 8
>>
>> Dirty tracker, after VKCTS test:
>>    order- 6 free:   4 MiB,  blocks: 19
>>    order- 5 free:   2 MiB,  blocks: 18
>>    order- 4 free: 704 KiB,  blocks: 11
>>    order- 3 free:   5 MiB,  blocks: 168
>>    order- 2 free:   2 MiB,  blocks: 174
>>    order- 1 free:   1 MiB,  blocks: 167
>>    order- 0 free:  28 KiB,  blocks: 7
>>
>> v2:
>>   - Code-style cleanup and minor refactoring
>>   - Renamed locals for clarity
>>
>> v3:
>>   - Keep cleared blocks inside free_tree[] instead of floating them.
>>   - Add subtree_has_dirty rbtree augment for O(log N) dirty-first walk.
>>
>> v4:
>>   - Fixed checkpatch warnings.
>>   - Optimized gpu_buddy_reset_clear() to a single post-order walk that
>>     flips block headers and recomputes the rbtree augment in one pass.
>>   - Propagate subtree_max_size top-down in insert_extent() so ancestors
>>     are not left with stale values on no-rotation inserts. (sashiko)
>>   - Drop the whole extent in gpu_dirty_tracker_mark_dirty() when the
>>     inside-split allocation fails, avoiding a stale clear claim. 
>> (sashiko)
>>   - Make gpu_dirty_tracker_find() alignment-aware and fall back to the
>>     dirty tree on steered failure to avoid spurious -ENOSPC. (sashiko)
>>
>> v5:
>>   - Track dirty extents instead of cleared ones: steer dirty allocs onto
>>     tracked dirty windows and pick clear allocs via a free-tree augment,
>>     avoiding clear-memory wastage by keeping cleared free blocks 
>> untouched
>>     during dirty allocation.
>>
>> v6:
>>   - Make __alloc_range_bias() return the highest/right-most address by
>>     default, establishing top-down as the intended placement for
>>     range-biased allocations.
>>   - Honour GPU_BUDDY_CLEAR_ALLOCATION in __alloc_range_bias() by 
>> steering
>>     the descent towards clear subtrees for non-top-down clear
>>     requests. (sashiko)
>>   - Skip dirty-tracker steering for offset-aligned requests so they keep
>>     their min_block_size alignment. (sashiko)
>>   - sashiko reported that the __GFP_NOFAIL dirty-extent allocations on
>>     the free path could deadlock during memory reclaim, since that is a
>>     GFP_KERNEL allocation on the free path; move to a per-tracker
>>     mempool so extent nodes are guaranteed without __GFP_NOFAIL.
>>     (sashiko)
>>   - Derive each free block's clear/dirty class from the blocks already
>>     in hand on split, free, alloc, trim and init instead of querying the
>>     dirty tracker, removing the tracker lookups from the hot paths.
>>
>> v7:
>>   - Preserve mixed-block clear state in __gpu_buddy_free() when a mixed
>>     split child is re-merged after an undone split. (sashiko)
>>   - Prefer a fully-clear block over a mixed one of the same order via a
>>     single ordered clear-state max augment on free_tree[].
>>
>> Assisted-by: Claude:claude-opus-4-8
>> Cc: Matthew Auld <[email protected]>
>> Cc: Christian König <[email protected]>
>> Signed-off-by: Arunpravin Paneer Selvam 
>> <[email protected]>
>> ---
>>   drivers/gpu/buddy.c                | 1318 ++++++++++++++++++++--------
>>   drivers/gpu/tests/gpu_buddy_test.c |   32 +-
>>   include/linux/gpu_buddy.h          |   97 +-
>>   3 files changed, 1034 insertions(+), 413 deletions(-)
>>
>> diff --git a/drivers/gpu/buddy.c b/drivers/gpu/buddy.c
>> index dc81fe0301ce..a5d68cd8b78d 100644
>> --- a/drivers/gpu/buddy.c
>> +++ b/drivers/gpu/buddy.c
>> @@ -8,6 +8,7 @@
>>   #include <linux/kmemleak.h>
>>   #include <linux/module.h>
>>   #include <linux/sizes.h>
>> +#include <linux/slab.h>
>>     #include <linux/gpu_buddy.h>
>>   @@ -34,6 +35,441 @@
>>   #endif
>>     static struct kmem_cache *slab_blocks;
>> +static struct kmem_cache *slab_extents;
>> +
>> +/*
>> + * A single reserved extent suffices. Every allocation uses GFP_KERNEL
>> + * from sleepable context, so the underlying slab alloc almost always
>> + * succeeds via reclaim; the reserve only backstops the rare case where
>> + * it still returns NULL (e.g. the current task is an OOM victim),
>> + * guaranteeing a non-NULL extent without __GFP_NOFAIL. Because each
>> + * alloc can independently wait for reclaim, the reserve need not scale
>> + * with the number of extents added in one locked section (e.g. by
>> + * gpu_buddy_reset_clear()).
>> + */
>> +#define GPU_DIRTY_EXTENT_POOL_MIN 1
>> +
>> +/*
>> + * Dirty tracker
>> + * -------------
>> + *
>> + * The dirty tracker maintains an augmented interval rbtree of 
>> contiguous
>> + * dirty address ranges, decoupled from the buddy free trees.
>> + * Each node covers a maximal coalesced run; adjacent extents are 
>> merged
>> + * on insertion so the tree always holds the smallest possible 
>> number of
>> + * extents.  The augmentation field @subtree_max_size lets the 
>> allocator
>> + * locate the largest dirty extent in O(log E).
>> + *
>> + * Free trees (mm->free_tree[])
>> + * ----------------------------
>> + *
>> + * Per-order augmented rbtrees of FREE buddy blocks, keyed by offset.
>> + * Every node carries:
>> + *   - subtree_max_alignment: largest natural alignment in the subtree,
>> + *     used by aligned/range allocations to skip unsuitable subtrees in
>> + *     O(log N).
>> + *   - subtree_block_state: the highest clear class (DIRTY < MIXED < 
>> CLEAR)
>> + *     of any block in the subtree, maintained as a max augment. A 
>> value of
>> + *     >= MIXED means a clear-or-mixed block exists; == CLEAR means a
>> + *     fully-clear block exists.
>> + *
>> + * Block classes
>> + * -------------
>> + *
>> + * Each FREE block falls into one of three classes, determined in
>> + * mark_free() by querying the dirty tracker for the block's range:
>> + *
>> + *   clear   -- HEADER_CLEAR set; no dirty extent overlaps the range.
>> + *   mixed   -- HEADER_CLEAR unset; range has both dirty and clear 
>> bytes.
>> + *   dirty   -- HEADER_CLEAR unset; range is fully dirty.
>> + *
>> + * Clear allocation
>> + * ----------------
>> + *
>> + * A clear (CLEAR_ALLOCATION) request prefers clear -> mixed -> dirty.
>> + * Climbing from the requested order up to max_order, 
>> rbtree_last_clear_free_block()
>> + * returns, in one O(log N) descent per order, the right-most 
>> clear-or-mixed block
>> + * (fully-clear preferred over mixed) at the lowest order that has 
>> one. Only if no
>> + * clear-or-mixed block exists at any order >= the requested one 
>> does it fall back
>> + * to a dirty block.
>> + *
>> + * Clear state is reported to the driver per whole block via 
>> HEADER_CLEAR, so a
>> + * fully-clear block of the requested order lets the driver skip the 
>> clear pass.
>> + *
>> + * The effective selection order therefore depends on the driver's
>> + * free policy:
>> + *
>> + *   1) If the driver never clears freed blocks, no free block ever 
>> holds
>> + *      clear bytes, so a clear request always falls back to a dirty 
>> block.
>> + *   2) If the driver clears every freed block, cleared ranges 
>> accumulate at
>> + *      the high end of the address space, so picking the right-most 
>> block
>> + *      yields clear -> mixed -> dirty.
>> + *   3) If the driver clears freed blocks selectively, fully-clear 
>> blocks are
>> + *      still preferred over mixed ones at the same order, and the 
>> right-most
>> + *      candidate wins, giving a clear -> mixed -> dirty order.
>> + */
>> +
>> +static u64 extent_size(struct gpu_dirty_extent *dirty_extent)
>> +{
>> +    return dirty_extent->end - dirty_extent->start;
>> +}
>> +
>> +RB_DECLARE_CALLBACKS_MAX(static, gpu_dirty_augment_cb,
>> +             struct gpu_dirty_extent, rb,
>> +             u64, subtree_max_size,
>> +             extent_size)
>> +
>> +static struct gpu_dirty_extent *extent_alloc(struct 
>> gpu_dirty_tracker *dirty_tracker)
>> +{
>> +    /*
>> +     * The void free/reset paths must record an extent and cannot 
>> handle
>> +     * failure, so the mempool reserve guarantees a non-NULL return
>> +     * without __GFP_NOFAIL. GFP_KERNEL is safe under the buddy 
>> lock: no
>> +     * driver frees buddy blocks from a shrinker, so reclaim cannot
>> +     * recurse into the lock we hold.
>> +     */
>> +    return mempool_alloc(dirty_tracker->extent_pool, GFP_KERNEL);
>> +}
>> +
>> +static void extent_free(struct gpu_dirty_tracker *dirty_tracker,
>> +            struct gpu_dirty_extent *dirty_extent)
>> +{
>> +    mempool_free(dirty_extent, dirty_tracker->extent_pool);
>> +}
>> +
>> +/* Return the rightmost extent whose start is strictly below 
>> @offset. */
>> +static struct gpu_dirty_extent *
>> +prev_extent(struct gpu_dirty_tracker *dirty_tracker, u64 offset)
>> +{
>> +    struct rb_node *rb = dirty_tracker->root.rb_node;
>> +    struct gpu_dirty_extent *dirty_extent = NULL;
>> +
>> +    while (rb) {
>> +        struct gpu_dirty_extent *tmp_extent =
>> +            rb_entry(rb, struct gpu_dirty_extent, rb);
>> +
>> +        if (tmp_extent->start < offset) {
>> +            dirty_extent = tmp_extent;
>> +            rb = rb->rb_right;
>> +        } else {
>> +            rb = rb->rb_left;
>> +        }
>> +    }
>> +
>> +    return dirty_extent;
>> +}
>> +
>> +/* Return the leftmost extent whose start is at or above @offset. */
>> +static struct gpu_dirty_extent *
>> +next_extent(struct gpu_dirty_tracker *dirty_tracker, u64 offset)
>> +{
>> +    struct rb_node *rb = dirty_tracker->root.rb_node;
>> +    struct gpu_dirty_extent *dirty_extent = NULL;
>> +
>> +    while (rb) {
>> +        struct gpu_dirty_extent *tmp_extent =
>> +            rb_entry(rb, struct gpu_dirty_extent, rb);
>> +
>> +        if (tmp_extent->start >= offset) {
>> +            dirty_extent = tmp_extent;
>> +            rb = rb->rb_left;
>> +        } else {
>> +            rb = rb->rb_right;
>> +        }
>> +    }
>> +
>> +    return dirty_extent;
>> +}
>> +
>> +static void insert_extent(struct gpu_dirty_tracker *dirty_tracker,
>> +              struct gpu_dirty_extent *dirty_extent)
>> +{
>> +    struct rb_node **link = &dirty_tracker->root.rb_node;
>> +    struct rb_node *parent = NULL;
>> +    u64 size = extent_size(dirty_extent);
>> +
>> +    while (*link) {
>> +        struct gpu_dirty_extent *tmp_extent;
>> +
>> +        parent = *link;
>> +        tmp_extent = rb_entry(parent, struct gpu_dirty_extent, rb);
>> +
>> +        if (tmp_extent->subtree_max_size < size)
>> +            tmp_extent->subtree_max_size = size;
>> +
>> +        if (dirty_extent->start < tmp_extent->start)
>> +            link = &parent->rb_left;
>> +        else
>> +            link = &parent->rb_right;
>> +    }
>> +
>> +    dirty_extent->subtree_max_size = size;
>> +    rb_link_node(&dirty_extent->rb, parent, link);
>> +    rb_insert_augmented(&dirty_extent->rb, &dirty_tracker->root, 
>> &gpu_dirty_augment_cb);
>> +}
>> +
>> +static void remove_extent(struct gpu_dirty_tracker *dirty_tracker,
>> +              struct gpu_dirty_extent *dirty_extent)
>> +{
>> +    rb_erase_augmented(&dirty_extent->rb, &dirty_tracker->root, 
>> &gpu_dirty_augment_cb);
>> +    RB_CLEAR_NODE(&dirty_extent->rb);
>> +}
>> +
>> +static int gpu_dirty_tracker_init(struct gpu_dirty_tracker 
>> *dirty_tracker)
>> +{
>> +    dirty_tracker->root = RB_ROOT;
>> +    dirty_tracker->total_dirty = 0;
>> +
>> +    dirty_tracker->extent_pool =
>> +        mempool_create_slab_pool(GPU_DIRTY_EXTENT_POOL_MIN, 
>> slab_extents);
>> +    if (!dirty_tracker->extent_pool)
>> +        return -ENOMEM;
>> +
>> +    return 0;
>> +}
>> +
>> +static void gpu_dirty_tracker_empty(struct gpu_dirty_tracker 
>> *dirty_tracker)
>> +{
>> +    struct rb_node *rb;
>> +
>> +    while ((rb = rb_first(&dirty_tracker->root))) {
>> +        struct gpu_dirty_extent *dirty_extent =
>> +            rb_entry(rb, struct gpu_dirty_extent, rb);
>> +
>> +        remove_extent(dirty_tracker, dirty_extent);
>> +        extent_free(dirty_tracker, dirty_extent);
>> +    }
>> +
>> +    dirty_tracker->total_dirty = 0;
>> +}
>> +
>> +static void gpu_dirty_tracker_fini(struct gpu_dirty_tracker 
>> *dirty_tracker)
>> +{
>> +    gpu_dirty_tracker_empty(dirty_tracker);
>> +    mempool_destroy(dirty_tracker->extent_pool);
>> +    dirty_tracker->extent_pool = NULL;
>> +}
>> +
>> +/*
>> + * Mark the range [start, start + size] as dirty. Merge with the 
>> neighbour on
>> + * each side if they are contiguous, so the tree never holds two 
>> adjacent ranges.
>> + */
>> +static void gpu_dirty_tracker_mark_dirty(struct gpu_dirty_tracker 
>> *dirty_tracker,
>> +                     u64 start, u64 size)
>> +{
>> +    struct gpu_dirty_extent *left, *right, *dirty_extent;
>> +    u64 end = start + size;
>> +
>> +    if (!size)
>> +        return;
>> +
>> +    /* Find contiguous neighbours, if any. */
>> +    left = prev_extent(dirty_tracker, start);
>> +    if (left && left->end != start)
>> +        left = NULL;
>> +
>> +    right = next_extent(dirty_tracker, end);
>> +    if (right && right->start != end)
>> +        right = NULL;
>> +
>> +    if (left && right) {
>> +        /* Merge left + new + right into a single extent. */
>> +        remove_extent(dirty_tracker, left);
>> +        remove_extent(dirty_tracker, right);
>> +        left->end = right->end;
>> +        extent_free(dirty_tracker, right);
>> +        insert_extent(dirty_tracker, left);
>> +    } else if (left) {
>> +        /* Extend left neighbour rightwards. */
>> +        remove_extent(dirty_tracker, left);
>> +        left->end = end;
>> +        insert_extent(dirty_tracker, left);
>> +    } else if (right) {
>> +        /* Extend right neighbour leftwards. */
>> +        remove_extent(dirty_tracker, right);
>> +        right->start = start;
>> +        insert_extent(dirty_tracker, right);
>> +    } else {
>> +        /* Standalone extent. */
>> +        dirty_extent = extent_alloc(dirty_tracker);
>> +        dirty_extent->start = start;
>> +        dirty_extent->end   = end;
>> +        insert_extent(dirty_tracker, dirty_extent);
>> +    }
>> +
>> +    dirty_tracker->total_dirty += size;
>> +}
>> +
>> +/*
>> + * Remove the range [start, start + size] from the dirty tracker. 
>> Punch the
>> + * range out of every overlapping dirty extent, splitting one extent 
>> in two if
>> + * the removed range falls strictly inside it.
>> + */
>> +static void gpu_dirty_tracker_remove_range(struct gpu_dirty_tracker 
>> *dirty_tracker,
>> +                       u64 start, u64 size)
>> +{
>> +    struct gpu_dirty_extent *dirty_extent, *next;
>> +    u64 end = start + size;
>> +
>> +    if (!size)
>> +        return;
>> +
>> +    dirty_extent = prev_extent(dirty_tracker, start + 1);
>> +    if (!dirty_extent)
>> +        dirty_extent = next_extent(dirty_tracker, start);
>> +
>> +    while (dirty_extent && dirty_extent->start < end) {
>> +        struct rb_node *next_node = rb_next(&dirty_extent->rb);
>> +        u64 extent_start = dirty_extent->start;
>> +        u64 extent_end = dirty_extent->end;
>> +
>> +        if (next_node)
>> +            next = rb_entry(next_node, struct gpu_dirty_extent, rb);
>> +        else
>> +            next = NULL;
>> +
>> +        /* Skip a non-overlapping neighbour returned by 
>> prev_extent(). */
>> +        if (extent_end <= start) {
>> +            dirty_extent = next;
>> +            continue;
>> +        }
>> +
>> +        if (extent_start < start && extent_end > end) {
>> +            /*
>> +             * Removed range lies strictly inside this dirty extent:
>> +             * split it into the dirty left and right halves.
>> +             */
>> +            struct gpu_dirty_extent *right = 
>> extent_alloc(dirty_tracker);
>> +
>> +            remove_extent(dirty_tracker, dirty_extent);
>> +
>> +            dirty_extent->end = start;
>> +            right->start = end;
>> +            right->end   = extent_end;
>> +
>> +            insert_extent(dirty_tracker, dirty_extent);
>> +            insert_extent(dirty_tracker, right);
>> +
>> +            dirty_tracker->total_dirty -= size;
>> +        } else if (extent_start >= start && extent_end <= end) {
>> +            /* Extent fully covered: drop it. */
>> +            remove_extent(dirty_tracker, dirty_extent);
>> +            extent_free(dirty_tracker, dirty_extent);
>> +
>> +            dirty_tracker->total_dirty -= (extent_end - extent_start);
>> +        } else if (extent_start < start) {
>> +            /* Extent overlaps from the left: trim its right end. */
>> +            remove_extent(dirty_tracker, dirty_extent);
>> +            dirty_extent->end = start;
>> +            insert_extent(dirty_tracker, dirty_extent);
>> +
>> +            dirty_tracker->total_dirty -= (extent_end - start);
>> +        } else {
>> +            /* Extent overlaps from the right: trim its left end. */
>> +            remove_extent(dirty_tracker, dirty_extent);
>> +            dirty_extent->start = end;
>> +            insert_extent(dirty_tracker, dirty_extent);
>> +
>> +            dirty_tracker->total_dirty -= (end - extent_start);
>> +        }
>> +
>> +        dirty_extent = next;
>> +    }
>> +}
>> +
>> +static enum gpu_block_state
>> +gpu_dirty_range_state(struct gpu_dirty_tracker *dirty_tracker,
>> +              u64 start, u64 size)
>> +{
>> +    struct gpu_dirty_extent *dirty_extent;
>> +    u64 end = start + size;
>> +
>> +    dirty_extent = prev_extent(dirty_tracker, start + 1);
>> +    if (dirty_extent) {
>> +        if (dirty_extent->start <= start && dirty_extent->end >= end)
>> +            return GPU_BLOCK_DIRTY;
>> +        if (dirty_extent->start < end && dirty_extent->end > start)
>> +            return GPU_BLOCK_MIXED;
>> +    }
>> +
>> +    dirty_extent = next_extent(dirty_tracker, start);
>> +    if (dirty_extent && dirty_extent->start < end)
>> +        return GPU_BLOCK_MIXED;
>> +
>> +    return GPU_BLOCK_CLEAR;
>> +}
>> +
>> +static struct rb_node *
>> +dirty_tracker_descend_right(struct rb_node *node, u64 min_size)
>> +{
>> +    while (node->rb_right) {
>> +        struct gpu_dirty_extent *tmp_extent;
>> +
>> +        tmp_extent = rb_entry(node->rb_right, struct 
>> gpu_dirty_extent, rb);
>> +
>> +        if (tmp_extent->subtree_max_size < min_size)
>> +            break;
>> +        node = node->rb_right;
>> +    }
>> +
>> +    return node;
>> +}
>> +
>> +static struct gpu_dirty_extent *
>> +gpu_dirty_tracker_find(struct gpu_dirty_tracker *dirty_tracker,
>> +               u64 min_size, u64 *aligned_start_out)
>> +{
>> +    struct rb_node *rb = dirty_tracker->root.rb_node;
>> +    struct gpu_dirty_extent *root_extent;
>> +    struct rb_node *parent;
>> +
>> +    if (!min_size || !is_power_of_2(min_size))
>> +        return NULL;
>> +
>> +    if (!rb)
>> +        return NULL;
>> +
>> +    root_extent = rb_entry(rb, struct gpu_dirty_extent, rb);
>> +    if (root_extent->subtree_max_size < min_size)
>> +        return NULL;
>> +
>> +    rb = dirty_tracker_descend_right(rb, min_size);
>> +
>> +    while (rb) {
>> +        struct gpu_dirty_extent *dirty_extent;
>> +        u64 aligned_start;
>> +
>> +        dirty_extent = rb_entry(rb, struct gpu_dirty_extent, rb);
>> +        aligned_start = ALIGN(dirty_extent->start, min_size);
>> +
>> +        /* Check if a min_size block fits after the alignment skip. */
>> +        if (aligned_start <= dirty_extent->end &&
>> +            dirty_extent->end - aligned_start >= min_size) {
>> +            *aligned_start_out = aligned_start;
>> +            return dirty_extent;
>> +        }
>> +
>> +        if (rb->rb_left) {
>> +            struct gpu_dirty_extent *tmp_extent;
>> +
>> +            tmp_extent = rb_entry(rb->rb_left, struct 
>> gpu_dirty_extent, rb);
>> +            if (tmp_extent->subtree_max_size >= min_size) {
>> +                rb = dirty_tracker_descend_right(rb->rb_left, 
>> min_size);
>> +                continue;
>> +            }
>> +        }
>> +
>> +        /* Walk up until we exit a node via its right child. */
>> +        parent = rb_parent(rb);
>> +        while (parent && parent->rb_right != rb) {
>> +            rb = parent;
>> +            parent = rb_parent(rb);
>> +        }
>> +        rb = parent;
>> +    }
>> +
>> +    return NULL;
>> +}
>>     static unsigned int
>>   gpu_buddy_block_state(struct gpu_buddy_block *block)
>> @@ -67,10 +503,97 @@ static unsigned int 
>> gpu_buddy_block_offset_alignment(struct gpu_buddy_block *blo
>>       return __ffs64(offset);
>>   }
>>   -RB_DECLARE_CALLBACKS_MAX(static, gpu_buddy_augment_cb,
>> -             struct gpu_buddy_block, rb,
>> -             unsigned int, subtree_max_alignment,
>> -             gpu_buddy_block_offset_alignment);
>> +static inline enum gpu_block_state
>> +gpu_block_cached_state(struct gpu_buddy_block *block)
>> +{
>> +    if (gpu_buddy_block_is_clear(block))
>> +        return GPU_BLOCK_CLEAR;
>> +    if (block->has_clear)
>> +        return GPU_BLOCK_MIXED;
>> +    return GPU_BLOCK_DIRTY;
>> +}
>> +
>> +static inline void gpu_buddy_augment_compute(struct gpu_buddy_block 
>> *block)
>> +{
>> +    enum gpu_block_state block_state;
>> +    struct gpu_buddy_block *right;
>> +    struct gpu_buddy_block *left;
>> +    unsigned int max_align;
>> +
>> +    max_align = gpu_buddy_block_offset_alignment(block);
>> +    block_state = gpu_block_cached_state(block);
>> +
>> +    left = rb_entry_safe(block->rb.rb_left, struct gpu_buddy_block, 
>> rb);
>> +    if (left) {
>> +        if (left->subtree_max_alignment > max_align)
>> +            max_align = left->subtree_max_alignment;
>> +
>> +        block_state = max(block_state, left->subtree_block_state);
>> +    }
>> +
>> +    right = rb_entry_safe(block->rb.rb_right, struct 
>> gpu_buddy_block, rb);
>> +    if (right) {
>> +        if (right->subtree_max_alignment > max_align)
>> +            max_align = right->subtree_max_alignment;
>> +
>> +        block_state = max(block_state, right->subtree_block_state);
>> +    }
>> +
>> +    block->subtree_max_alignment = max_align;
>> +    block->subtree_block_state = block_state;
>> +}
>> +
>> +static void gpu_buddy_augment_propagate(struct rb_node *rb, struct 
>> rb_node *stop)
>> +{
>> +    while (rb != stop) {
>> +        struct gpu_buddy_block *block;
>> +        unsigned int old_align;
>> +        enum gpu_block_state old_block_state;
>> +
>> +        block = rb_entry(rb, struct gpu_buddy_block, rb);
>> +        old_align = block->subtree_max_alignment;
>> +        old_block_state = block->subtree_block_state;
>> +
>> +        gpu_buddy_augment_compute(block);
>> +        if (block->subtree_max_alignment == old_align &&
>> +            block->subtree_block_state == old_block_state)
>> +            break;
>> +
>> +        rb = rb_parent(&block->rb);
>> +    }
>> +}
>> +
>> +static void gpu_buddy_augment_copy(struct rb_node *rb_old, struct 
>> rb_node *rb_new)
>> +{
>> +    struct gpu_buddy_block *old;
>> +    struct gpu_buddy_block *new;
>> +
>> +    old = rb_entry(rb_old, struct gpu_buddy_block, rb);
>> +    new = rb_entry(rb_new, struct gpu_buddy_block, rb);
>> +
>> +    new->subtree_max_alignment = old->subtree_max_alignment;
>> +    new->subtree_block_state = old->subtree_block_state;
>> +}
>> +
>> +static void gpu_buddy_augment_rotate(struct rb_node *rb_old, struct 
>> rb_node *rb_new)
>> +{
>> +    struct gpu_buddy_block *old;
>> +    struct gpu_buddy_block *new;
>> +
>> +    old = rb_entry(rb_old, struct gpu_buddy_block, rb);
>> +    new = rb_entry(rb_new, struct gpu_buddy_block, rb);
>> +
>> +    new->subtree_max_alignment = old->subtree_max_alignment;
>> +    new->subtree_block_state = old->subtree_block_state;
>> +
>> +    gpu_buddy_augment_compute(old);
>> +}
>> +
>> +static const struct rb_augment_callbacks gpu_buddy_augment_cb = {
>> +    .propagate = gpu_buddy_augment_propagate,
>> +    .copy      = gpu_buddy_augment_copy,
>> +    .rotate    = gpu_buddy_augment_rotate,
>> +};
>>     static struct gpu_buddy_block *gpu_block_alloc(struct gpu_buddy *mm,
>>                              struct gpu_buddy_block *parent,
>> @@ -81,6 +604,10 @@ static struct gpu_buddy_block 
>> *gpu_block_alloc(struct gpu_buddy *mm,
>>         BUG_ON(order > GPU_BUDDY_MAX_ORDER);
>>   +    /*
>> +     * GFP_KERNEL is safe under the buddy lock: no consumer runs a
>> +     * shrinker that re-enters it during direct reclaim.
>> +     */
>>       block = kmem_cache_zalloc(slab_blocks, GFP_KERNEL);
>>       if (!block)
>>           return NULL;
>> @@ -101,13 +628,6 @@ static void gpu_block_free(struct gpu_buddy *mm,
>>       kmem_cache_free(slab_blocks, block);
>>   }
>>   -static enum gpu_buddy_free_tree
>> -get_block_tree(struct gpu_buddy_block *block)
>> -{
>> -    return gpu_buddy_block_is_clear(block) ?
>> -           GPU_BUDDY_CLEAR_TREE : GPU_BUDDY_DIRTY_TREE;
>> -}
>> -
>>   static struct gpu_buddy_block *
>>   rbtree_get_free_block(const struct rb_node *node)
>>   {
>> @@ -120,24 +640,64 @@ rbtree_last_free_block(struct rb_root *root)
>>       return rbtree_get_free_block(rb_last(root));
>>   }
>>   -static bool rbtree_is_empty(struct rb_root *root)
>> +static struct gpu_buddy_block *
>> +rbtree_last_clear_free_block(struct rb_root *root,
>> +                 enum gpu_block_state min_block_state)
>>   {
>> -    return RB_EMPTY_ROOT(root);
>> +    struct rb_node *node = root->rb_node;
>> +    struct gpu_buddy_block *block = NULL;
>> +    struct gpu_buddy_block *root_block;
>> +    enum gpu_block_state target_state;
>> +
>> +    root_block = rbtree_get_free_block(node);
>> +    if (!root_block || root_block->subtree_block_state < 
>> min_block_state)
>> +        return NULL;
>> +
>> +    target_state = root_block->subtree_block_state;
>> +
>> +    while (node) {
>> +        struct gpu_buddy_block *right_block;
>> +        struct gpu_buddy_block *node_block;
>> +
>> +        node_block = rbtree_get_free_block(node);
>> +        right_block = rbtree_get_free_block(node->rb_right);
>> +
>> +        if (right_block && right_block->subtree_block_state >= 
>> target_state) {
>> +            node = node->rb_right;
>> +            continue;
>> +        }
>> +
>> +        if (gpu_block_cached_state(node_block) == target_state) {
>> +            block = node_block;
>> +            break;
>> +        }
>> +
>> +        node = node->rb_left;
>> +    }
>> +
>> +    return block;
>> +}
>> +
>> +static inline void gpu_buddy_sync_clear_avail(struct gpu_buddy *mm)
>> +{
>> +    mm->clear_avail = mm->avail - mm->dirty.total_dirty;
>>   }
>>     static void rbtree_insert(struct gpu_buddy *mm,
>> -              struct gpu_buddy_block *block,
>> -              enum gpu_buddy_free_tree tree)
>> +              struct gpu_buddy_block *block)
>>   {
>>       struct rb_node **link, *parent = NULL;
>> -    unsigned int block_alignment, order;
>> +    enum gpu_block_state block_state;
>>       struct gpu_buddy_block *node;
>> +    unsigned int block_alignment;
>>       struct rb_root *root;
>> +    unsigned int order;
>>         order = gpu_buddy_block_order(block);
>>       block_alignment = gpu_buddy_block_offset_alignment(block);
>> +    block_state = gpu_block_cached_state(block);
>>   -    root = &mm->free_trees[tree][order];
>> +    root = &mm->free_tree[order];
>>       link = &root->rb_node;
>>         while (*link) {
>> @@ -147,10 +707,12 @@ static void rbtree_insert(struct gpu_buddy *mm,
>>            * Manual augmentation update during insertion traversal. 
>> Required
>>            * because rb_insert_augmented() only calls rotate callback 
>> during
>>            * rotations. This ensures all ancestors on the insertion 
>> path have
>> -         * correct subtree_max_alignment values.
>> +         * correct subtree_max_alignment / subtree_block_state values.
>>            */
>>           if (node->subtree_max_alignment < block_alignment)
>>               node->subtree_max_alignment = block_alignment;
>> +        if (node->subtree_block_state < block_state)
>> +            node->subtree_block_state = block_state;
>>             if (gpu_buddy_block_offset(block) < 
>> gpu_buddy_block_offset(node))
>>               link = &parent->rb_left;
>> @@ -159,6 +721,7 @@ static void rbtree_insert(struct gpu_buddy *mm,
>>       }
>>         block->subtree_max_alignment = block_alignment;
>> +    block->subtree_block_state = block_state;
>>       rb_link_node(&block->rb, parent, link);
>>       rb_insert_augmented(&block->rb, root, &gpu_buddy_augment_cb);
>>   }
>> @@ -167,26 +730,11 @@ static void rbtree_remove(struct gpu_buddy *mm,
>>                 struct gpu_buddy_block *block)
>>   {
>>       unsigned int order = gpu_buddy_block_order(block);
>> -    enum gpu_buddy_free_tree tree;
>> -    struct rb_root *root;
>>   -    tree = get_block_tree(block);
>> -    root = &mm->free_trees[tree][order];
>> -
>> -    rb_erase_augmented(&block->rb, root, &gpu_buddy_augment_cb);
>> +    rb_erase_augmented(&block->rb, &mm->free_tree[order], 
>> &gpu_buddy_augment_cb);
>>       RB_CLEAR_NODE(&block->rb);
>>   }
>>   -static void clear_reset(struct gpu_buddy_block *block)
>> -{
>> -    block->header &= ~GPU_BUDDY_HEADER_CLEAR;
>> -}
>> -
>> -static void mark_cleared(struct gpu_buddy_block *block)
>> -{
>> -    block->header |= GPU_BUDDY_HEADER_CLEAR;
>> -}
>> -
>>   static void mark_allocated(struct gpu_buddy *mm,
>>                  struct gpu_buddy_block *block)
>>   {
>> @@ -199,21 +747,36 @@ static void mark_allocated(struct gpu_buddy *mm,
>>       rbtree_remove(mm, block);
>>   }
>>   -static void mark_free(struct gpu_buddy *mm,
>> -              struct gpu_buddy_block *block)
>> +static void __mark_free(struct gpu_buddy *mm,
>> +            struct gpu_buddy_block *block,
>> +            enum gpu_block_state block_state)
>>   {
>> -    enum gpu_buddy_free_tree tree;
>> -
>>       if (gpu_buddy_block_is_allocated(block))
>> mm->used_scoreboard[gpu_buddy_block_order(block)]--;
>>         block->header &= ~GPU_BUDDY_HEADER_STATE;
>>       block->header |= GPU_BUDDY_FREE;
>>   +    block->header &= ~GPU_BUDDY_HEADER_CLEAR;
>> +
>> +    block->has_clear = (block_state != GPU_BLOCK_DIRTY);
>> +    if (block_state == GPU_BLOCK_CLEAR)
>> +        block->header |= GPU_BUDDY_HEADER_CLEAR;
>> +
>>       mm->free_scoreboard[gpu_buddy_block_order(block)]++;
>>   -    tree = get_block_tree(block);
>> -    rbtree_insert(mm, block, tree);
>> +    rbtree_insert(mm, block);
>> +}
>> +
>> +static void mark_free(struct gpu_buddy *mm,
>> +              struct gpu_buddy_block *block)
>> +{
>> +    enum gpu_block_state block_state;
>> +
>> +    block_state = gpu_dirty_range_state(&mm->dirty,
>> +                        gpu_buddy_block_offset(block),
>> +                        gpu_buddy_block_size(mm, block));
>> +    __mark_free(mm, block, block_state);
>>   }
>>     static void mark_split(struct gpu_buddy *mm,
>> @@ -253,37 +816,31 @@ __get_buddy(struct gpu_buddy_block *block)
>>   }
>>     static unsigned int __gpu_buddy_free(struct gpu_buddy *mm,
>> -                     struct gpu_buddy_block *block,
>> -                     bool force_merge)
>> +                     struct gpu_buddy_block *block)
>>   {
>> +    enum gpu_block_state block_state;
>>       struct gpu_buddy_block *parent;
>>       unsigned int order;
>>   -    while ((parent = block->parent)) {
>> -        struct gpu_buddy_block *buddy;
>> +    block_state = gpu_block_cached_state(block);
>>   -        buddy = __get_buddy(block);
>> +    while ((parent = block->parent)) {
>> +        struct gpu_buddy_block *buddy = __get_buddy(block);
>>             if (!gpu_buddy_block_is_free(buddy))
>>               break;
>>   -        if (!force_merge) {
>> -            /*
>> -             * Check the block and its buddy clear state and exit
>> -             * the loop if they both have the dissimilar state.
>> -             */
>> -            if (gpu_buddy_block_is_clear(block) !=
>> -                gpu_buddy_block_is_clear(buddy))
>> -                break;
>> +        if (block_state != GPU_BLOCK_MIXED) {
>> +            enum gpu_block_state buddy_state;
>>   -            if (gpu_buddy_block_is_clear(block))
>> -                mark_cleared(parent);
>> +            buddy_state = gpu_block_cached_state(buddy);
>> +
>> +            if (buddy_state != block_state)
>> +                block_state = GPU_BLOCK_MIXED;
>>           }
>>             rbtree_remove(mm, buddy);
>> mm->free_scoreboard[gpu_buddy_block_order(buddy)]--;
>> -        if (force_merge && gpu_buddy_block_is_clear(buddy))
>> -            mm->clear_avail -= gpu_buddy_block_size(mm, buddy);
>>             if (gpu_buddy_block_is_allocated(block))
>> mm->used_scoreboard[gpu_buddy_block_order(block)]--;
>> @@ -295,74 +852,11 @@ static unsigned int __gpu_buddy_free(struct 
>> gpu_buddy *mm,
>>       }
>>         order = gpu_buddy_block_order(block);
>> -    mark_free(mm, block);
>> +    __mark_free(mm, block, block_state);
>>         return order;
>>   }
>>   -static int __force_merge(struct gpu_buddy *mm,
>> -             u64 start,
>> -             u64 end,
>> -             unsigned int min_order)
>> -{
>> -    unsigned int tree, order;
>> -    int i;
>> -
>> -    if (!min_order)
>> -        return -ENOMEM;
>> -
>> -    if (min_order > mm->max_order)
>> -        return -EINVAL;
>> -
>> -    for_each_free_tree(tree) {
>> -        for (i = min_order - 1; i >= 0; i--) {
>> -            struct rb_node *iter = rb_last(&mm->free_trees[tree][i]);
>> -
>> -            while (iter) {
>> -                struct gpu_buddy_block *block, *buddy;
>> -                u64 block_start, block_end;
>> -
>> -                block = rbtree_get_free_block(iter);
>> -                iter = rb_prev(iter);
>> -
>> -                if (!block || !block->parent)
>> -                    continue;
>> -
>> -                block_start = gpu_buddy_block_offset(block);
>> -                block_end = block_start + gpu_buddy_block_size(mm, 
>> block) - 1;
>> -
>> -                if (!contains(start, end, block_start, block_end))
>> -                    continue;
>> -
>> -                buddy = __get_buddy(block);
>> -                if (!gpu_buddy_block_is_free(buddy))
>> -                    continue;
>> -
>> - gpu_buddy_assert(gpu_buddy_block_is_clear(block) !=
>> -                         gpu_buddy_block_is_clear(buddy));
>> -
>> -                /*
>> -                 * Advance to the next node when the current node is 
>> the buddy,
>> -                 * as freeing the block will also remove its buddy 
>> from the tree.
>> -                 */
>> -                if (iter == &buddy->rb)
>> -                    iter = rb_prev(iter);
>> -
>> -                rbtree_remove(mm, block);
>> - mm->free_scoreboard[gpu_buddy_block_order(block)]--;
>> -                if (gpu_buddy_block_is_clear(block))
>> -                    mm->clear_avail -= gpu_buddy_block_size(mm, block);
>> -
>> -                order = __gpu_buddy_free(mm, block, true);
>> -                if (order >= min_order)
>> -                    return 0;
>> -            }
>> -        }
>> -    }
>> -
>> -    return -ENOMEM;
>> -}
>> -
>>   /**
>>    * gpu_buddy_init - init memory manager
>>    *
>> @@ -377,7 +871,7 @@ static int __force_merge(struct gpu_buddy *mm,
>>    */
>>   int gpu_buddy_init(struct gpu_buddy *mm, u64 size, u64 chunk_size)
>>   {
>> -    unsigned int i, j, root_count = 0;
>> +    unsigned int root_count = 0;
>>       u64 offset = 0;
>>         if (size < chunk_size)
>> @@ -411,22 +905,14 @@ int gpu_buddy_init(struct gpu_buddy *mm, u64 
>> size, u64 chunk_size)
>>       if (!mm->used_scoreboard)
>>           goto out_free_free_scoreboard;
>>   -    mm->free_trees = kmalloc_array(GPU_BUDDY_MAX_FREE_TREES,
>> -                       sizeof(*mm->free_trees),
>> -                       GFP_KERNEL);
>> -    if (!mm->free_trees)
>> +    mm->free_tree = kcalloc(mm->max_order + 1,
>> +                sizeof(struct rb_root),
>> +                GFP_KERNEL);
>> +    if (!mm->free_tree)
>>           goto out_free_used_scoreboard;
>>   -    for_each_free_tree(i) {
>> -        mm->free_trees[i] = kmalloc_array(mm->max_order + 1,
>> -                          sizeof(struct rb_root),
>> -                          GFP_KERNEL);
>> -        if (!mm->free_trees[i])
>> -            goto out_free_tree;
>> -
>> -        for (j = 0; j <= mm->max_order; ++j)
>> -            mm->free_trees[i][j] = RB_ROOT;
>> -    }
>> +    if (gpu_dirty_tracker_init(&mm->dirty))
>> +        goto out_free_tree;
>>         mm->n_roots = hweight64(size);
>>   @@ -452,7 +938,8 @@ int gpu_buddy_init(struct gpu_buddy *mm, u64 
>> size, u64 chunk_size)
>>           if (!root)
>>               goto out_free_roots;
>>   -        mark_free(mm, root);
>> +        gpu_dirty_tracker_mark_dirty(&mm->dirty, offset, root_size);
>> +        __mark_free(mm, root, GPU_BLOCK_DIRTY);
>>             BUG_ON(root_count > mm->max_order);
>>           BUG_ON(gpu_buddy_block_size(mm, root) < chunk_size);
>> @@ -474,9 +961,8 @@ int gpu_buddy_init(struct gpu_buddy *mm, u64 
>> size, u64 chunk_size)
>>           gpu_block_free(mm, mm->roots[root_count]);
>>       kfree(mm->roots);
>>   out_free_tree:
>> -    while (i--)
>> -        kfree(mm->free_trees[i]);
>> -    kfree(mm->free_trees);
>> +    gpu_dirty_tracker_fini(&mm->dirty);
>> +    kfree(mm->free_tree);
>>   out_free_used_scoreboard:
>>       kfree(mm->used_scoreboard);
>>   out_free_free_scoreboard:
>> @@ -494,7 +980,7 @@ EXPORT_SYMBOL(gpu_buddy_init);
>>    */
>>   void gpu_buddy_fini(struct gpu_buddy *mm)
>>   {
>> -    u64 root_size, size, start;
>> +    u64 root_size, size;
>>       unsigned int order;
>>       int i;
>>   @@ -502,14 +988,10 @@ void gpu_buddy_fini(struct gpu_buddy *mm)
>>         for (i = 0; i < mm->n_roots; ++i) {
>>           order = ilog2(size) - ilog2(mm->chunk_size);
>> -        start = gpu_buddy_block_offset(mm->roots[i]);
>> -        __force_merge(mm, start, start + size, order);
>> +        root_size = mm->chunk_size << order;
>> gpu_buddy_assert(gpu_buddy_block_is_free(mm->roots[i]));
>> -
>>           gpu_block_free(mm, mm->roots[i]);
>> -
>> -        root_size = mm->chunk_size << order;
>>           size -= root_size;
>>       }
>>   @@ -518,9 +1000,8 @@ void gpu_buddy_fini(struct gpu_buddy *mm)
>>       for (i = 0; i <= mm->max_order; ++i)
>>           gpu_buddy_assert(!mm->used_scoreboard[i]);
>>   -    for_each_free_tree(i)
>> -        kfree(mm->free_trees[i]);
>> -    kfree(mm->free_trees);
>> +    gpu_dirty_tracker_fini(&mm->dirty);
>> +    kfree(mm->free_tree);
>>       kfree(mm->roots);
>>       kfree(mm->free_scoreboard);
>>       kfree(mm->used_scoreboard);
>> @@ -532,6 +1013,7 @@ static int split_block(struct gpu_buddy *mm,
>>   {
>>       unsigned int block_order = gpu_buddy_block_order(block) - 1;
>>       u64 offset = gpu_buddy_block_offset(block);
>> +    enum gpu_block_state parent_state;
>>         BUG_ON(!gpu_buddy_block_is_free(block));
>>       BUG_ON(!gpu_buddy_block_order(block));
>> @@ -547,17 +1029,18 @@ static int split_block(struct gpu_buddy *mm,
>>           return -ENOMEM;
>>       }
>>   +    parent_state = gpu_block_cached_state(block);
>> +
>>       mark_split(mm, block);
>>   -    if (gpu_buddy_block_is_clear(block)) {
>> -        mark_cleared(block->left);
>> -        mark_cleared(block->right);
>> -        clear_reset(block);
>> +    if (parent_state == GPU_BLOCK_MIXED) {
>> +        mark_free(mm, block->left);
>> +        mark_free(mm, block->right);
>> +    } else {
>> +        __mark_free(mm, block->left, parent_state);
>> +        __mark_free(mm, block->right, parent_state);
>>       }
>>   -    mark_free(mm, block->left);
>> -    mark_free(mm, block->right);
>> -
>>       return 0;
>>   }
>>   @@ -572,42 +1055,39 @@ static int split_block(struct gpu_buddy *mm,
>>    */
>>   void gpu_buddy_reset_clear(struct gpu_buddy *mm, bool is_clear)
>>   {
>> -    enum gpu_buddy_free_tree src_tree, dst_tree;
>> -    u64 root_size, size, start;
>> -    unsigned int order;
>> -    int i;
>> +    unsigned int i;
>>         gpu_buddy_driver_lock_held(mm);
>> -    size = mm->size;
>> -    for (i = 0; i < mm->n_roots; ++i) {
>> -        order = ilog2(size) - ilog2(mm->chunk_size);
>> -        start = gpu_buddy_block_offset(mm->roots[i]);
>> -        __force_merge(mm, start, start + size, order);
>>   -        root_size = mm->chunk_size << order;
>> -        size -= root_size;
>> -    }
>> -
>> -    src_tree = is_clear ? GPU_BUDDY_DIRTY_TREE : GPU_BUDDY_CLEAR_TREE;
>> -    dst_tree = is_clear ? GPU_BUDDY_CLEAR_TREE : GPU_BUDDY_DIRTY_TREE;
>> +    gpu_dirty_tracker_empty(&mm->dirty);
>>         for (i = 0; i <= mm->max_order; ++i) {
>> -        struct rb_root *root = &mm->free_trees[src_tree][i];
>>           struct gpu_buddy_block *block, *tmp;
>>   -        rbtree_postorder_for_each_entry_safe(block, tmp, root, rb) {
>> -            rbtree_remove(mm, block);
>> +        rbtree_postorder_for_each_entry_safe(block, tmp,
>> +                             &mm->free_tree[i], rb) {
>>               if (is_clear) {
>> -                mark_cleared(block);
>> -                mm->clear_avail += gpu_buddy_block_size(mm, block);
>> +                if (!gpu_buddy_block_is_clear(block))
>> +                    block->header |= GPU_BUDDY_HEADER_CLEAR;
>> +                block->has_clear = true;
>> +            } else if (gpu_buddy_block_is_clear(block)) {
>> +                block->header &= ~GPU_BUDDY_HEADER_CLEAR;
>> +                block->has_clear = false;
>> +                gpu_dirty_tracker_mark_dirty(&mm->dirty,
>> +                                 gpu_buddy_block_offset(block),
>> +                                 gpu_buddy_block_size(mm, block));
>>               } else {
>> -                clear_reset(block);
>> -                mm->clear_avail -= gpu_buddy_block_size(mm, block);
>> +                block->has_clear = false;
>> +                gpu_dirty_tracker_mark_dirty(&mm->dirty,
>> +                                 gpu_buddy_block_offset(block),
>> +                                 gpu_buddy_block_size(mm, block));
>>               }
>>   -            rbtree_insert(mm, block, dst_tree);
>> +            gpu_buddy_augment_compute(block);
>>           }
>>       }
>> +
>> +    gpu_buddy_sync_clear_avail(mm);
>>   }
>>   EXPORT_SYMBOL(gpu_buddy_reset_clear);
>>   @@ -620,13 +1100,18 @@ EXPORT_SYMBOL(gpu_buddy_reset_clear);
>>   void gpu_buddy_free_block(struct gpu_buddy *mm,
>>                 struct gpu_buddy_block *block)
>>   {
>> +    u64 size = gpu_buddy_block_size(mm, block);
>> +    u64 offset = gpu_buddy_block_offset(block);
>> +
>>       gpu_buddy_driver_lock_held(mm);
>>       BUG_ON(!gpu_buddy_block_is_allocated(block));
>> -    mm->avail += gpu_buddy_block_size(mm, block);
>> -    if (gpu_buddy_block_is_clear(block))
>> -        mm->clear_avail += gpu_buddy_block_size(mm, block);
>>   -    __gpu_buddy_free(mm, block, false);
>> +    mm->avail += size;
>> +    if (!gpu_buddy_block_is_clear(block))
>> +        gpu_dirty_tracker_mark_dirty(&mm->dirty, offset, size);
>> +
>> +    gpu_buddy_sync_clear_avail(mm);
>> +    __gpu_buddy_free(mm, block);
>>   }
>>   EXPORT_SYMBOL(gpu_buddy_free_block);
>>   @@ -641,9 +1126,9 @@ static void __gpu_buddy_free_list(struct 
>> gpu_buddy *mm,
>>         list_for_each_entry_safe(block, on, objects, link) {
>>           if (mark_clear)
>> -            mark_cleared(block);
>> +            block->header |= GPU_BUDDY_HEADER_CLEAR;
>>           else if (mark_dirty)
>> -            clear_reset(block);
>> +            block->header &= ~GPU_BUDDY_HEADER_CLEAR;
>>           gpu_buddy_free_block(mm, block);
>>           cond_resched();
>>       }
>> @@ -679,13 +1164,6 @@ void gpu_buddy_free_list(struct gpu_buddy *mm,
>>   }
>>   EXPORT_SYMBOL(gpu_buddy_free_list);
>>   -static bool block_incompatible(struct gpu_buddy_block *block, 
>> unsigned int flags)
>> -{
>> -    bool needs_clear = flags & GPU_BUDDY_CLEAR_ALLOCATION;
>> -
>> -    return needs_clear != gpu_buddy_block_is_clear(block);
>> -}
>> -
>>   static void __gpu_buddy_undo_splits(struct gpu_buddy *mm,
>>                       struct gpu_buddy_block *block)
>>   {
>> @@ -696,7 +1174,7 @@ static void __gpu_buddy_undo_splits(struct 
>> gpu_buddy *mm,
>>            gpu_buddy_block_is_free(buddy))) {
>>           rbtree_remove(mm, block);
>> mm->free_scoreboard[gpu_buddy_block_order(block)]--;
>> -        __gpu_buddy_free(mm, block, false);
>> +        __gpu_buddy_free(mm, block);
>>       }
>>   }
>>   @@ -704,8 +1182,7 @@ static struct gpu_buddy_block *
>>   __alloc_range_bias(struct gpu_buddy *mm,
>>              u64 start, u64 end,
>>              unsigned int order,
>> -           unsigned long flags,
>> -           bool fallback)
>> +           unsigned long flags)
>>   {
>>       u64 req_size = mm->chunk_size << order;
>>       struct gpu_buddy_block *block;
>> @@ -715,7 +1192,15 @@ __alloc_range_bias(struct gpu_buddy *mm,
>>         end = end - 1;
>>   -    for (i = 0; i < mm->n_roots; ++i)
>> +    /*
>> +     * This range-constrained search hands back the highest/right-most
>> +     * address that satisfies the request: the roots are seeded 
>> high-to-low
>> +     * and the right (higher-address) child is descended first, making
>> +     * top-down the default placement here. A non-top-down clear 
>> request is
>> +     * the only exception, where the descent is biased towards clear or
>> +     * clear-containing subtrees to satisfy the clear preference.
>> +     */
>> +    for (i = mm->n_roots - 1; i >= 0; --i)
>>           list_add_tail(&mm->roots[i]->tmp_link, &dfs);
>>         do {
>> @@ -751,9 +1236,6 @@ __alloc_range_bias(struct gpu_buddy *mm,
>>                   continue;
>>           }
>>   -        if (!fallback && block_incompatible(block, flags))
>> -            continue;
>> -
>>           if (contains(start, end, block_start, block_end) &&
>>               order == gpu_buddy_block_order(block)) {
>>               /*
>> @@ -771,8 +1253,38 @@ __alloc_range_bias(struct gpu_buddy *mm,
>>                   goto err_undo;
>>           }
>>   -        list_add(&block->right->tmp_link, &dfs);
>> -        list_add(&block->left->tmp_link, &dfs);
>> +        /*
>> +         * Top-down is a strict address-placement policy, so when it is
>> +         * requested we ignore clear steering and simply descend the
>> +         * right (higher-address) child first. Only a non-top-down 
>> clear
>> +         * request biases the descent towards clear/has_clear subtrees.
>> +         */
>> +        if ((flags & GPU_BUDDY_CLEAR_ALLOCATION) &&
>> +            !(flags & GPU_BUDDY_TOPDOWN_ALLOCATION)) {
>> +            struct gpu_buddy_block *prefer;
>> +
>> +            if (gpu_buddy_block_is_clear(block->right))
>> +                prefer = block->right;
>> +            else if (gpu_buddy_block_is_clear(block->left))
>> +                prefer = block->left;
>> +            else if (block->right->has_clear)
>> +                prefer = block->right;
>> +            else if (block->left->has_clear)
>> +                prefer = block->left;
>> +            else
>> +                prefer = block->right;
>> +
>> +            if (prefer == block->right) {
>> +                list_add(&block->left->tmp_link, &dfs);
>> +                list_add(&block->right->tmp_link, &dfs);
>> +            } else {
>> +                list_add(&block->right->tmp_link, &dfs);
>> +                list_add(&block->left->tmp_link, &dfs);
>> +            }
>> +        } else {
>> +            list_add(&block->left->tmp_link, &dfs);
>> +            list_add(&block->right->tmp_link, &dfs);
>> +        }
>>       } while (1);
>>         return ERR_PTR(-ENOSPC);
>> @@ -787,48 +1299,32 @@ __alloc_range_bias(struct gpu_buddy *mm,
>>       return ERR_PTR(err);
>>   }
>>   -static struct gpu_buddy_block *
>> -__gpu_buddy_alloc_range_bias(struct gpu_buddy *mm,
>> -                 u64 start, u64 end,
>> -                 unsigned int order,
>> -                 unsigned long flags)
>> -{
>> -    struct gpu_buddy_block *block;
>> -    bool fallback = false;
>> -
>> -    block = __alloc_range_bias(mm, start, end, order,
>> -                   flags, fallback);
>> -    if (IS_ERR(block))
>> -        return __alloc_range_bias(mm, start, end, order,
>> -                      flags, !fallback);
>> -
>> -    return block;
>> -}
>> -
>> +/* Return the highest-address free block of at least @order. */
>>   static struct gpu_buddy_block *
>>   get_maxblock(struct gpu_buddy *mm,
>> -         unsigned int order,
>> -         enum gpu_buddy_free_tree tree)
>> +         unsigned int order)
>>   {
>> -    struct gpu_buddy_block *max_block = NULL, *block = NULL;
>> -    struct rb_root *root;
>> +    struct gpu_buddy_block *max_block;
>> +    struct gpu_buddy_block *block;
>>       unsigned int i;
>>   +    /*
>> +     * Top-down allocation is a strict address-placement policy: the 
>> block
>> +     * is chosen purely by offset, regardless of its clear/dirty state.
>> +     * Clear state is re-derived from the dirty tracker once the 
>> allocation
>> +     * completes, and the driver is responsible for issuing the 
>> clear pass
>> +     * if a clear region is required.
>> +     */
>> +    max_block = NULL;
>> +
>>       for (i = order; i <= mm->max_order; ++i) {
>> -        root = &mm->free_trees[tree][i];
>> -        block = rbtree_last_free_block(root);
>> +        block = rbtree_last_free_block(&mm->free_tree[i]);
>>           if (!block)
>>               continue;
>>   -        if (!max_block) {
>> -            max_block = block;
>> -            continue;
>> -        }
>> -
>> -        if (gpu_buddy_block_offset(block) >
>> -            gpu_buddy_block_offset(max_block)) {
>> +        if (!max_block ||
>> +            gpu_buddy_block_offset(block) > 
>> gpu_buddy_block_offset(max_block))
>>               max_block = block;
>> -        }
>>       }
>>         return max_block;
>> @@ -840,45 +1336,34 @@ alloc_from_freetree(struct gpu_buddy *mm,
>>               unsigned long flags)
>>   {
>>       struct gpu_buddy_block *block = NULL;
>> -    struct rb_root *root;
>> -    enum gpu_buddy_free_tree tree;
>>       unsigned int tmp;
>>       int err;
>>   -    tree = (flags & GPU_BUDDY_CLEAR_ALLOCATION) ?
>> -        GPU_BUDDY_CLEAR_TREE : GPU_BUDDY_DIRTY_TREE;
>> -
>>       if (flags & GPU_BUDDY_TOPDOWN_ALLOCATION) {
>> -        block = get_maxblock(mm, order, tree);
>> +        block = get_maxblock(mm, order);
>>           if (block)
>> -            /* Store the obtained block order */
>>               tmp = gpu_buddy_block_order(block);
>>       } else {
>> -        for (tmp = order; tmp <= mm->max_order; ++tmp) {
>> -            /* Get RB tree root for this order and tree */
>> -            root = &mm->free_trees[tree][tmp];
>> -            block = rbtree_last_free_block(root);
>> -            if (block)
>> -                break;
>> +        if (flags & GPU_BUDDY_CLEAR_ALLOCATION) {
>> +            for (tmp = order; tmp <= mm->max_order; ++tmp) {
>> +                block = 
>> rbtree_last_clear_free_block(&mm->free_tree[tmp],
>> +                                     GPU_BLOCK_MIXED);
>> +                if (block)
>> +                    break;
>> +            }
>>           }
>> -    }
>> -
>> -    if (!block) {
>> -        /* Try allocating from the other tree */
>> -        tree = (tree == GPU_BUDDY_CLEAR_TREE) ?
>> -            GPU_BUDDY_DIRTY_TREE : GPU_BUDDY_CLEAR_TREE;
>> -
>> -        for (tmp = order; tmp <= mm->max_order; ++tmp) {
>> -            root = &mm->free_trees[tree][tmp];
>> -            block = rbtree_last_free_block(root);
>> -            if (block)
>> -                break;
>> +        if (!block) {
>> +            for (tmp = order; tmp <= mm->max_order; ++tmp) {
>> +                block = rbtree_last_free_block(&mm->free_tree[tmp]);
>> +                if (block)
>> +                    break;
>> +            }
>>           }
>> -
>> -        if (!block)
>> -            return ERR_PTR(-ENOSPC);
>>       }
>>   +    if (!block)
>> +        return ERR_PTR(-ENOSPC);
>> +
>>       BUG_ON(!gpu_buddy_block_is_free(block));
>>         while (tmp != order) {
>> @@ -886,7 +1371,26 @@ alloc_from_freetree(struct gpu_buddy *mm,
>>           if (unlikely(err))
>>               goto err_undo;
>>   -        block = block->right;
>> +        if ((flags & GPU_BUDDY_CLEAR_ALLOCATION) &&
>> +            !(flags & GPU_BUDDY_TOPDOWN_ALLOCATION)) {
>> +            bool right_clear, left_clear;
>> +
>> +            right_clear = gpu_buddy_block_is_clear(block->right);
>> +            left_clear = gpu_buddy_block_is_clear(block->left);
>> +
>> +            if (right_clear)
>> +                block = block->right;
>> +            else if (left_clear)
>> +                block = block->left;
>> +            else if (block->right->has_clear)
>> +                block = block->right;
>> +            else if (block->left->has_clear)
>> +                block = block->left;
>> +            else
>> +                block = block->right;
>> +        } else {
>> +            block = block->right;
>> +        }
>>           tmp--;
>>       }
>>       return block;
>> @@ -913,12 +1417,10 @@ static bool 
>> gpu_buddy_subtree_can_satisfy(struct rb_node *node,
>>     static struct gpu_buddy_block *
>>   gpu_buddy_find_block_aligned(struct gpu_buddy *mm,
>> -                 enum gpu_buddy_free_tree tree,
>>                    unsigned int order,
>> -                 unsigned int alignment,
>> -                 unsigned long flags)
>> +                 unsigned int alignment)
>>   {
>> -    struct rb_root *root = &mm->free_trees[tree][order];
>> +    struct rb_root *root = &mm->free_tree[order];
>>       struct rb_node *rb = root->rb_node;
>>         while (rb) {
>> @@ -951,12 +1453,10 @@ gpu_buddy_find_block_aligned(struct gpu_buddy 
>> *mm,
>>   static struct gpu_buddy_block *
>>   gpu_buddy_offset_aligned_allocation(struct gpu_buddy *mm,
>>                       u64 size,
>> -                    u64 min_block_size,
>> -                    unsigned long flags)
>> +                    u64 min_block_size)
>>   {
>>       struct gpu_buddy_block *block = NULL;
>>       unsigned int order, tmp, alignment;
>> -    enum gpu_buddy_free_tree tree;
>>       unsigned long pages;
>>       int err;
>>   @@ -964,19 +1464,15 @@ gpu_buddy_offset_aligned_allocation(struct 
>> gpu_buddy *mm,
>>       pages = size >> ilog2(mm->chunk_size);
>>       order = fls(pages) - 1;
>>   -    tree = (flags & GPU_BUDDY_CLEAR_ALLOCATION) ?
>> -        GPU_BUDDY_CLEAR_TREE : GPU_BUDDY_DIRTY_TREE;
>> -
>> +    /*
>> +     * Offset-aligned allocation is a strict address-placement 
>> policy: the
>> +     * block is chosen purely by its offset alignment, regardless of 
>> its
>> +     * clear/dirty state. Clear state is re-derived from the dirty 
>> tracker
>> +     * once the allocation completes, and the driver is responsible for
>> +     * issuing the clear pass if a clear region is required.
>> +     */
>>       for (tmp = order; tmp <= mm->max_order; ++tmp) {
>> -        block = gpu_buddy_find_block_aligned(mm, tree, tmp,
>> -                             alignment, flags);
>> -        if (!block) {
>> -            tree = (tree == GPU_BUDDY_CLEAR_TREE) ?
>> -                GPU_BUDDY_DIRTY_TREE : GPU_BUDDY_CLEAR_TREE;
>> -            block = gpu_buddy_find_block_aligned(mm, tree, tmp,
>> -                                 alignment, flags);
>> -        }
>> -
>> +        block = gpu_buddy_find_block_aligned(mm, tmp, alignment);
>>           if (block)
>>               break;
>>       }
>> @@ -1015,6 +1511,7 @@ gpu_buddy_offset_aligned_allocation(struct 
>> gpu_buddy *mm,
>>   static int __alloc_range(struct gpu_buddy *mm,
>>                struct list_head *dfs,
>>                u64 start, u64 size,
>> +             unsigned long flags,
>>                struct list_head *blocks,
>>                u64 *total_allocated_on_err)
>>   {
>> @@ -1051,16 +1548,33 @@ static int __alloc_range(struct gpu_buddy *mm,
>>             if (contains(start, end, block_start, block_end)) {
>>               if (gpu_buddy_block_is_free(block)) {
>> +                bool block_clear = false;
>> +                u64 block_offset;
>> +                u64 block_size;
>> +
>> +                block_size = gpu_buddy_block_size(mm, block);
>> +                block_offset = gpu_buddy_block_offset(block);
>> +
>> +                if (flags & GPU_BUDDY_CLEAR_ALLOCATION)
>> +                    block_clear = gpu_buddy_block_is_clear(block);
>> +
>> +                if (!gpu_buddy_block_is_clear(block))
>> + gpu_dirty_tracker_remove_range(&mm->dirty,
>> +                                       block_offset,
>> +                                       block_size);
>> +
>>                   mark_allocated(mm, block);
>> -                total_allocated += gpu_buddy_block_size(mm, block);
>> -                mm->avail -= gpu_buddy_block_size(mm, block);
>> -                if (gpu_buddy_block_is_clear(block))
>> -                    mm->clear_avail -= gpu_buddy_block_size(mm, block);
>> +                total_allocated += block_size;
>> +                mm->avail -= block_size;
>> +
>> +                block->header &= ~GPU_BUDDY_HEADER_CLEAR;
>> +                if (block_clear)
>> +                    block->header |= GPU_BUDDY_HEADER_CLEAR;
>> +
>> +                gpu_buddy_sync_clear_avail(mm);
>> +
>>                   list_add_tail(&block->link, &allocated);
>>                   continue;
>> -            } else if (!mm->clear_avail) {
>> -                err = -ENOSPC;
>> -                goto err_free;
>>               }
>>           }
>>   @@ -1105,6 +1619,7 @@ static int __alloc_range(struct gpu_buddy *mm,
>>   static int __gpu_buddy_alloc_range(struct gpu_buddy *mm,
>>                      u64 start,
>>                      u64 size,
>> +                   unsigned long flags,
>>                      u64 *total_allocated_on_err,
>>                      struct list_head *blocks)
>>   {
>> @@ -1114,20 +1629,23 @@ static int __gpu_buddy_alloc_range(struct 
>> gpu_buddy *mm,
>>       for (i = 0; i < mm->n_roots; ++i)
>>           list_add_tail(&mm->roots[i]->tmp_link, &dfs);
>>   -    return __alloc_range(mm, &dfs, start, size,
>> +    return __alloc_range(mm, &dfs, start, size, flags,
>>                    blocks, total_allocated_on_err);
>>   }
>>     static int __alloc_contig_try_harder(struct gpu_buddy *mm,
>>                        u64 size,
>>                        u64 min_block_size,
>> +                     unsigned long flags,
>>                        struct list_head *blocks)
>>   {
>>       u64 rhs_offset, lhs_offset, lhs_size, filled;
>>       struct gpu_buddy_block *block;
>> -    unsigned int tree, order;
>>       LIST_HEAD(blocks_lhs);
>> +    struct rb_root *root;
>> +    struct rb_node *iter;
>>       unsigned long pages;
>> +    unsigned int order;
>>       u64 modify_size;
>>       int err;
>>   @@ -1137,45 +1655,40 @@ static int __alloc_contig_try_harder(struct 
>> gpu_buddy *mm,
>>       if (order == 0)
>>           return -ENOSPC;
>>   -    for_each_free_tree(tree) {
>> -        struct rb_root *root;
>> -        struct rb_node *iter;
>> -
>> -        root = &mm->free_trees[tree][order];
>> -        if (rbtree_is_empty(root))
>> -            continue;
>> +    root = &mm->free_tree[order];
>> +    if (RB_EMPTY_ROOT(root))
>> +        return -ENOSPC;
>>   -        iter = rb_last(root);
>> -        while (iter) {
>> -            block = rbtree_get_free_block(iter);
>> -
>> -            /* Allocate blocks traversing RHS */
>> -            rhs_offset = gpu_buddy_block_offset(block);
>> -            err =  __gpu_buddy_alloc_range(mm, rhs_offset, size,
>> -                               &filled, blocks);
>> -            if (!err || err != -ENOSPC)
>> -                return err;
>> -
>> -            lhs_size = max((size - filled), min_block_size);
>> -            if (!IS_ALIGNED(lhs_size, min_block_size))
>> -                lhs_size = round_up(lhs_size, min_block_size);
>> -
>> -            /* Allocate blocks traversing LHS */
>> -            lhs_offset = gpu_buddy_block_offset(block) - lhs_size;
>> -            err =  __gpu_buddy_alloc_range(mm, lhs_offset, lhs_size,
>> -                               NULL, &blocks_lhs);
>> -            if (!err) {
>> -                list_splice(&blocks_lhs, blocks);
>> -                return 0;
>> -            } else if (err != -ENOSPC) {
>> -                gpu_buddy_free_list_internal(mm, blocks);
>> -                return err;
>> -            }
>> -            /* Free blocks for the next iteration */
>> +    iter = rb_last(root);
>> +    while (iter) {
>> +        block = rbtree_get_free_block(iter);
>> +
>> +        /* Allocate blocks traversing RHS */
>> +        rhs_offset = gpu_buddy_block_offset(block);
>> +        err =  __gpu_buddy_alloc_range(mm, rhs_offset, size,
>> +                           flags, &filled, blocks);
>> +        if (!err || err != -ENOSPC)
>> +            return err;
>> +
>> +        lhs_size = max((size - filled), min_block_size);
>> +        if (!IS_ALIGNED(lhs_size, min_block_size))
>> +            lhs_size = round_up(lhs_size, min_block_size);
>> +
>> +        /* Allocate blocks traversing LHS */
>> +        lhs_offset = gpu_buddy_block_offset(block) - lhs_size;
>> +        err =  __gpu_buddy_alloc_range(mm, lhs_offset, lhs_size,
>> +                           flags, NULL, &blocks_lhs);
>> +        if (!err) {
>> +            list_splice(&blocks_lhs, blocks);
>> +            return 0;
>> +        } else if (err != -ENOSPC) {
>>               gpu_buddy_free_list_internal(mm, blocks);
>> -
>> -            iter = rb_prev(iter);
>> +            return err;
>>           }
>> +        /* Free blocks for the next iteration */
>> +        gpu_buddy_free_list_internal(mm, blocks);
>> +
>> +        iter = rb_prev(iter);
>>       }
>>         return -ENOSPC;
>> @@ -1209,6 +1722,7 @@ int gpu_buddy_block_trim(struct gpu_buddy *mm,
>>       struct gpu_buddy_block *block;
>>       u64 block_start, block_end;
>>       LIST_HEAD(dfs);
>> +    bool was_clear;
>>       u64 new_start;
>>       int err;
>>   @@ -1251,22 +1765,38 @@ int gpu_buddy_block_trim(struct gpu_buddy *mm,
>>       }
>>         list_del(&block->link);
>> -    mark_free(mm, block);
>> +
>> +    was_clear = gpu_buddy_block_is_clear(block);
>> +    block->header &= ~GPU_BUDDY_HEADER_CLEAR;
>> +
>> +    if (!was_clear)
>> +        gpu_dirty_tracker_mark_dirty(&mm->dirty,
>> +                         gpu_buddy_block_offset(block),
>> +                         gpu_buddy_block_size(mm, block));
>> +
>> +    __mark_free(mm, block, was_clear ? GPU_BLOCK_CLEAR : 
>> GPU_BLOCK_DIRTY);
>>       mm->avail += gpu_buddy_block_size(mm, block);
>> -    if (gpu_buddy_block_is_clear(block))
>> -        mm->clear_avail += gpu_buddy_block_size(mm, block);
>> +    gpu_buddy_sync_clear_avail(mm);
>>         /* Prevent recursively freeing this node */
>>       parent = block->parent;
>>       block->parent = NULL;
>>         list_add(&block->tmp_link, &dfs);
>> -    err =  __alloc_range(mm, &dfs, new_start, new_size, blocks, NULL);
>> +    err =  __alloc_range(mm, &dfs, new_start, new_size,
>> +                 was_clear ? GPU_BUDDY_CLEAR_ALLOCATION : 0,
>> +                 blocks, NULL);
>>       if (err) {
>>           mark_allocated(mm, block);
>>           mm->avail -= gpu_buddy_block_size(mm, block);
>> -        if (gpu_buddy_block_is_clear(block))
>> -            mm->clear_avail -= gpu_buddy_block_size(mm, block);
>> +        if (!was_clear) {
>> +            gpu_dirty_tracker_remove_range(&mm->dirty,
>> +                               gpu_buddy_block_offset(block),
>> +                               gpu_buddy_block_size(mm, block));
>> +        }
>> +        if (was_clear)
>> +            block->header |= GPU_BUDDY_HEADER_CLEAR;
>> +        gpu_buddy_sync_clear_avail(mm);
>>           list_add(&block->link, blocks);
>>       }
>>   @@ -1275,6 +1805,22 @@ int gpu_buddy_block_trim(struct gpu_buddy *mm,
>>   }
>>   EXPORT_SYMBOL(gpu_buddy_block_trim);
>>   +static bool dirty_steer_window(struct gpu_buddy *mm, u64 req_size,
>> +                   u64 *start, u64 *end, unsigned long *flags)
>> +{
>> +    u64 aligned_start;
>> +    struct gpu_dirty_extent *ext =
>> +        gpu_dirty_tracker_find(&mm->dirty, req_size, &aligned_start);
>> +
>> +    if (!ext)
>> +        return false;
>> +
>> +    *start  = aligned_start;
>> +    *end    = ext->end;
>> +    *flags |= GPU_BUDDY_RANGE_ALLOCATION;
>> +    return true;
>> +}
>> +
>>   static struct gpu_buddy_block *
>>   __gpu_buddy_alloc_blocks(struct gpu_buddy *mm,
>>                u64 start, u64 end,
>> @@ -1282,18 +1828,36 @@ __gpu_buddy_alloc_blocks(struct gpu_buddy *mm,
>>                unsigned int order,
>>                unsigned long flags)
>>   {
>> -    if (flags & GPU_BUDDY_RANGE_ALLOCATION)
>> +    struct gpu_buddy_block *block;
>> +    bool steered = false;
>> +
>> +    /* Allocate from dirty tracker */
>> +    if (!(flags & GPU_BUDDY_RANGE_ALLOCATION) &&
>> +        !(flags & GPU_BUDDY_CLEAR_ALLOCATION) &&
>> +        size >= min_block_size &&
>> +        mm->clear_avail && mm->dirty.total_dirty) {
>> +        u64 block_size = mm->chunk_size << order;
>> +
>> +        steered = dirty_steer_window(mm, block_size,
>> +                         &start, &end, &flags);
>> +    }
>> +
>> +    if (flags & GPU_BUDDY_RANGE_ALLOCATION) {
>>           /* Allocate traversing within the range */
>> -        return  __gpu_buddy_alloc_range_bias(mm, start, end,
>> -                             order, flags);
>> -    else if (size < min_block_size)
>> +        block = __alloc_range_bias(mm, start, end, order, flags);
>> +        if (!IS_ERR(block) || !steered)
>> +            return block;
>> +
>> +        flags &= ~GPU_BUDDY_RANGE_ALLOCATION;
>> +    }
>> +
>> +    if (size < min_block_size)
>>           /* Allocate from an offset-aligned region without size 
>> rounding */
>>           return gpu_buddy_offset_aligned_allocation(mm, size,
>> -                               min_block_size,
>> -                               flags);
>> -    else
>> -        /* Allocate from freetree */
>> -        return alloc_from_freetree(mm, order, flags);
>> +                               min_block_size);
>> +
>> +    /* Allocate from freetree */
>> +    return alloc_from_freetree(mm, order, flags);
>>   }
>>     /**
>> @@ -1354,7 +1918,7 @@ int gpu_buddy_alloc_blocks(struct gpu_buddy *mm,
>>           if (!IS_ALIGNED(start | end, min_block_size))
>>               return -EINVAL;
>>   -        return __gpu_buddy_alloc_range(mm, start, size, NULL, 
>> blocks);
>> +        return __gpu_buddy_alloc_range(mm, start, size, flags, NULL, 
>> blocks);
>>       }
>>         original_size = size;
>> @@ -1380,12 +1944,15 @@ int gpu_buddy_alloc_blocks(struct gpu_buddy *mm,
>>           if ((flags & GPU_BUDDY_CONTIGUOUS_ALLOCATION) &&
>>               !(flags & GPU_BUDDY_RANGE_ALLOCATION))
>>               return __alloc_contig_try_harder(mm, original_size,
>> -                             original_min_size, blocks);
>> +                             original_min_size,
>> +                             flags, blocks);
>>             return -EINVAL;
>>       }
>>         do {
>> +        bool block_clear = false;
>> +
>>           order = min(order, (unsigned int)fls(pages) - 1);
>>           BUG_ON(order > mm->max_order);
>>           /*
>> @@ -1395,8 +1962,6 @@ int gpu_buddy_alloc_blocks(struct gpu_buddy *mm,
>>           BUG_ON(size >= min_block_size && order < min_order);
>>             do {
>> -            unsigned int fallback_order;
>> -
>>               block = __gpu_buddy_alloc_blocks(mm, start,
>>                                end,
>>                                size,
>> @@ -1406,48 +1971,48 @@ int gpu_buddy_alloc_blocks(struct gpu_buddy *mm,
>>               if (!IS_ERR(block))
>>                   break;
>>   -            if (size < min_block_size) {
>> -                fallback_order = order;
>> -            } else if (order == min_order) {
>> -                fallback_order = min_order;
>> -            } else {
>> +            if (size >= min_block_size && order > min_order) {
>>                   order--;
>>                   continue;
>>               }
>>   -            /* Try allocation through force merge method */
>> -            if (mm->clear_avail &&
>> -                !__force_merge(mm, start, end, fallback_order)) {
>> -                block = __gpu_buddy_alloc_blocks(mm, start,
>> -                                 end,
>> -                                 size,
>> -                                 min_block_size,
>> -                                 fallback_order,
>> -                                 flags);
>> -                if (!IS_ERR(block)) {
>> -                    order = fallback_order;
>> -                    break;
>> -                }
>> -            }
>> -
>>               /*
>>                * Try contiguous block allocation through
>>                * try harder method.
>>                */
>>               if (flags & GPU_BUDDY_CONTIGUOUS_ALLOCATION &&
>> -                !(flags & GPU_BUDDY_RANGE_ALLOCATION))
>> -                return __alloc_contig_try_harder(mm,
>> -                                 original_size,
>> -                                 original_min_size,
>> -                                 blocks);
>> +                !(flags & GPU_BUDDY_RANGE_ALLOCATION)) {
>> +                err = __alloc_contig_try_harder(mm,
>> +                                original_size,
>> +                                original_min_size,
>> +                                flags,
>> +                                blocks);
>> +                if (!err)
>> +                    return 0;
>> +                if (err != -ENOSPC)
>> +                    return err;
>> +                goto err_free;
>> +            }
>>               err = -ENOSPC;
>>               goto err_free;
>>           } while (1);
>>   +        if (flags & GPU_BUDDY_CLEAR_ALLOCATION)
>> +            block_clear = gpu_buddy_block_is_clear(block);
>> +
>> +        if (!gpu_buddy_block_is_clear(block))
>> +            gpu_dirty_tracker_remove_range(&mm->dirty,
>> +                               gpu_buddy_block_offset(block),
>> +                               gpu_buddy_block_size(mm, block));
>> +
>>           mark_allocated(mm, block);
>>           mm->avail -= gpu_buddy_block_size(mm, block);
>> -        if (gpu_buddy_block_is_clear(block))
>> -            mm->clear_avail -= gpu_buddy_block_size(mm, block);
>> +
>> +        block->header &= ~GPU_BUDDY_HEADER_CLEAR;
>> +        if (block_clear)
>> +            block->header |= GPU_BUDDY_HEADER_CLEAR;
>> +
>> +        gpu_buddy_sync_clear_avail(mm);
>>           kmemleak_update_trace(block);
>>           list_add_tail(&block->link, &allocated);
>>   @@ -1542,6 +2107,7 @@ EXPORT_SYMBOL(gpu_buddy_print);
>>     static void gpu_buddy_module_exit(void)
>>   {
>> +    kmem_cache_destroy(slab_extents);
>>       kmem_cache_destroy(slab_blocks);
>>   }
>>   @@ -1551,7 +2117,15 @@ static int __init gpu_buddy_module_init(void)
>>       if (!slab_blocks)
>>           return -ENOMEM;
>>   +    slab_extents = KMEM_CACHE(gpu_dirty_extent, 0);
>> +    if (!slab_extents)
>> +        goto err_destroy_blocks;
>> +
>>       return 0;
>> +
>> +err_destroy_blocks:
>> +    kmem_cache_destroy(slab_blocks);
>> +    return -ENOMEM;
>>   }
>>     module_init(gpu_buddy_module_init);
>> diff --git a/drivers/gpu/tests/gpu_buddy_test.c 
>> b/drivers/gpu/tests/gpu_buddy_test.c
>> index 89698563c61b..f8e56da5058e 100644
>> --- a/drivers/gpu/tests/gpu_buddy_test.c
>> +++ b/drivers/gpu/tests/gpu_buddy_test.c
>> @@ -38,7 +38,7 @@ static void 
>> gpu_test_buddy_subtree_offset_alignment_stress(struct kunit *test)
>>       };
>>       struct list_head allocated[ARRAY_SIZE(alignments)];
>>       unsigned int i, max_subtree_align = 0;
>> -    int ret, tree, order;
>> +    int ret, order;
>>       struct gpu_buddy mm;
>>         KUNIT_ASSERT_FALSE_MSG(test, gpu_buddy_init(&mm, mm_size, 
>> SZ_4K),
>> @@ -78,15 +78,11 @@ static void 
>> gpu_test_buddy_subtree_offset_alignment_stress(struct kunit *test)
>>           }
>>             for (order = mm.max_order; order >= 0 && !root; order--) {
>> -            for (tree = 0; tree < 2; tree++) {
>> -                node = mm.free_trees[tree][order].rb_node;
>> -                if (node) {
>> -                    root = container_of(node,
>> -                                struct gpu_buddy_block,
>> -                                rb);
>> -                    break;
>> -                }
>> -            }
>> +            node = mm.free_tree[order].rb_node;
>> +            if (node)
>> +                root = container_of(node,
>> +                            struct gpu_buddy_block,
>> +                            rb);
>>           }
>>             KUNIT_ASSERT_NOT_NULL(test, root);
>> @@ -97,15 +93,13 @@ static void 
>> gpu_test_buddy_subtree_offset_alignment_stress(struct kunit *test)
>>           gpu_buddy_free_list(&mm, &allocated[i], 0);
>>             for (order = 0; order <= mm.max_order; order++) {
>> -            for (tree = 0; tree < 2; tree++) {
>> -                node = mm.free_trees[tree][order].rb_node;
>> -                if (!node)
>> -                    continue;
>> -
>> -                block = container_of(node, struct gpu_buddy_block, rb);
>> -                max_subtree_align = max(max_subtree_align,
>> -                            block->subtree_max_alignment);
>> -            }
>> +            node = mm.free_tree[order].rb_node;
>> +            if (!node)
>> +                continue;
>> +
>> +            block = container_of(node, struct gpu_buddy_block, rb);
>> +            max_subtree_align = max(max_subtree_align,
>> +                        block->subtree_max_alignment);
>>           }
>>             KUNIT_EXPECT_GE(test, max_subtree_align, 
>> ilog2(alignments[i]));
>> diff --git a/include/linux/gpu_buddy.h b/include/linux/gpu_buddy.h
>> index e037714563d8..899b84298cd8 100644
>> --- a/include/linux/gpu_buddy.h
>> +++ b/include/linux/gpu_buddy.h
>> @@ -8,6 +8,7 @@
>>     #include <linux/bitops.h>
>>   #include <linux/list.h>
>> +#include <linux/mempool.h>
>>   #include <linux/slab.h>
>>   #include <linux/sched.h>
>>   #include <linux/rbtree.h>
>> @@ -43,8 +44,8 @@
>>   /**
>>    * GPU_BUDDY_CLEAR_ALLOCATION - Prefer pre-cleared (zeroed) memory
>>    *
>> - * Attempt to allocate from the clear tree first. If insufficient clear
>> - * memory is available, falls back to dirty memory. Useful when the
>> + * Attempt to allocate outside dirty-tracked ranges first. If 
>> insufficient
>> + * clear memory is available, falls back to dirty memory. Useful 
>> when the
>>    * caller needs zeroed memory and wants to avoid GPU clear operations.
>>    */
>>   #define GPU_BUDDY_CLEAR_ALLOCATION        BIT(3)
>> @@ -53,8 +54,8 @@
>>    * GPU_BUDDY_CLEARED - Mark returned blocks as cleared
>>    *
>>    * Used with gpu_buddy_free_list() to indicate that the memory being
>> - * freed has been cleared (zeroed). The blocks will be placed in the
>> - * clear tree for future GPU_BUDDY_CLEAR_ALLOCATION requests.
>> + * freed has been cleared (zeroed). The blocks will be removed from the
>> + * dirty tracker for future GPU_BUDDY_CLEAR_ALLOCATION requests.
>>    */
>>   #define GPU_BUDDY_CLEARED            BIT(4)
>>   @@ -67,15 +68,6 @@
>>    */
>>   #define GPU_BUDDY_TRIM_DISABLE            BIT(5)
>>   -enum gpu_buddy_free_tree {
>> -    GPU_BUDDY_CLEAR_TREE = 0,
>> -    GPU_BUDDY_DIRTY_TREE,
>> -    GPU_BUDDY_MAX_FREE_TREES,
>> -};
>> -
>> -#define for_each_free_tree(tree) \
>> -    for ((tree) = 0; (tree) < GPU_BUDDY_MAX_FREE_TREES; (tree)++)
>> -
>>   /**
>>    * struct gpu_buddy_block - Block within a buddy allocator
>>    *
>> @@ -88,6 +80,17 @@ enum gpu_buddy_free_tree {
>>    * @private: Private data owned by the allocator user (e.g., 
>> driver-specific data)
>>    * @link: List node for user ownership while block is allocated
>>    */
>> +/*
>> + * Clear/dirty state of a free block. Ordered so a numerically 
>> larger value
>> + * is "more clear" (DIRTY < MIXED < CLEAR) which lets 
>> subtree_block_state be
>> + * maintained as a simple max-augment over the per-order free tree.
>> + */
>> +enum gpu_block_state {
>> +    GPU_BLOCK_DIRTY = 0,
>> +    GPU_BLOCK_MIXED = 1,
>> +    GPU_BLOCK_CLEAR = 2,
>> +};
>> +
>>   struct gpu_buddy_block {
>>   /* private: */
>>       /*
>> @@ -103,6 +106,13 @@ struct gpu_buddy_block {
>>   #define   GPU_BUDDY_ALLOCATED       (1 << 10)
>>   #define   GPU_BUDDY_FREE       (2 << 10)
>>   #define   GPU_BUDDY_SPLIT       (3 << 10)
>> +/*
>> + * GPU_BUDDY_HEADER_CLEAR has two roles:
>> + *  - FREE state:      set when the block's full range is cleared 
>> (dirty
>> + *                     tracker confirmed no overlap).
>> + *  - ALLOCATED state: set when the block was served from cleared 
>> memory,
>> + *                     informing the caller that no GPU clear pass 
>> is needed.
>> + */
>>   #define GPU_BUDDY_HEADER_CLEAR  GENMASK_ULL(9, 9)
>>   /* Free to be used, if needed in the future */
>>   #define GPU_BUDDY_HEADER_UNUSED GENMASK_ULL(8, 6)
>> @@ -128,13 +138,51 @@ struct gpu_buddy_block {
>>           struct list_head link;
>>       };
>>   /* private: */
>> -    struct list_head tmp_link;
>> +    enum gpu_block_state subtree_block_state;
>>       unsigned int subtree_max_alignment;
>> +    struct list_head tmp_link;
>> +    bool has_clear;
>>   };
>>     /* Order-zero must be at least SZ_4K */
>>   #define GPU_BUDDY_MAX_ORDER (63 - 12)
>>   +/**
>> + * struct gpu_dirty_extent - a contiguous dirty address range
>> + *
>> + * Tracks a single contiguous address range whose memory content is 
>> known
>> + * to be dirty.  Extents are non-overlapping and stored in an augmented
>> + * red-black tree sorted by @start.  The augmented value 
>> @subtree_max_size
>> + * allows O(log N) search for an extent of at least a given size.
>> + */
>> +struct gpu_dirty_extent {
>> +/* private: */
>> +    struct rb_node    rb;
>> +    u64        start;
>> +    u64        end;
>> +    u64        subtree_max_size;
>> +};
>> +
>> +/**
>> + * struct gpu_dirty_tracker - tracks dirty address intervals
>> + *
>> + * Maintains a set of non-overlapping dirty extents as an augmented
>> + * red-black tree.
>> + *
>> + * @total_dirty: Total bytes of dirty memory currently tracked.
>> + * @extent_pool: Mempool backing extent node allocations. sashiko 
>> reported
>> + *         that a __GFP_NOFAIL allocation on the free path could
>> + *         deadlock during memory reclaim, so a per-tracker mempool is
>> + *         used to guarantee extent nodes without __GFP_NOFAIL.
>> + */
>> +struct gpu_dirty_tracker {
>> +/* private: */
>> +    struct rb_root    root;
>> +    mempool_t    *extent_pool;
>> +/* public: */
>> +    u64        total_dirty;
>> +};
>> +
>>   /**
>>    * struct gpu_buddy - GPU binary buddy allocator
>>    *
>> @@ -152,20 +200,25 @@ struct gpu_buddy_block {
>>    * @chunk_size: Minimum allocation granularity in bytes. Must be at 
>> least SZ_4K.
>>    * @size: Total size of the address space managed by this allocator 
>> in bytes.
>>    * @avail: Total free space currently available for allocation in 
>> bytes.
>> - * @clear_avail: Free space available in the clear tree (zeroed 
>> memory) in bytes.
>> - *               This is a subset of @avail.
>> + * @clear_avail: Free space that is clear (zeroed) in bytes. A 
>> subset of @avail.
>> + *               Maintained as @avail - dirty.total_dirty, since the 
>> tracker
>> + *               records the dirty extents. Zero at init, as a fresh 
>> pool is
>> + *               fully dirty.
>>    * @lock_dep_map: Annotates gpu_buddy API with a driver provided lock.
>>    */
>>   struct gpu_buddy {
>>   /* private: */
>> +    /* Tracker of dirty address ranges (decoupled from free_tree). */
>> +    struct gpu_dirty_tracker dirty;
>>       /*
>> -     * Array of red-black trees for free block management.
>> -     * Indexed as free_trees[clear/dirty][order] where:
>> -     * - Index 0 (GPU_BUDDY_CLEAR_TREE): blocks with zeroed content
>> -     * - Index 1 (GPU_BUDDY_DIRTY_TREE): blocks with unknown content
>> -     * Each tree holds free blocks of the corresponding order.
>> +     * One RB-tree per order containing all free blocks (clear and
>> +     * dirty alike).  The augment field subtree_block_state (a max over
>> +     * the subtree of each block's state) lets clear allocations
>> +     * find the right-most fully-clear or mixed block in O(log N).
>> +     * Dirty free blocks coexist here but are also indexed by the
>> +     * @dirty tracker for fast dirty allocation lookups.
>>        */
>> -    struct rb_root **free_trees;
>> +    struct rb_root *free_tree;
>>       /*
>>        * Array of root blocks representing the top-level blocks of the
>>        * binary tree(s). Multiple roots exist when the total size is not
>>
>> base-commit: 9c950822f0fa923ccd344d7a143872d25efe89a3
>