sim_lib_compute_wgpu/
arena.rs

1//! Bounded resident allocation arena for wgpu tensor planning.
2
3use std::collections::VecDeque;
4
5/// Opaque resident allocation id inside a planned wgpu arena.
6#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
7pub struct WgpuAllocationId(u64);
8
9/// Allocation record returned by the arena.
10#[derive(Clone, Debug, PartialEq, Eq)]
11pub struct WgpuArenaAllocation {
12    /// Allocation id.
13    pub id: WgpuAllocationId,
14    /// Allocation size in bytes.
15    pub bytes: u64,
16}
17
18#[derive(Clone, Debug)]
19struct LiveAllocation {
20    allocation: WgpuArenaAllocation,
21    active: bool,
22}
23
24/// Snapshot of resident arena pressure.
25#[derive(Clone, Debug, Default, PartialEq, Eq)]
26pub struct WgpuArenaSnapshot {
27    /// Active resident allocation count.
28    pub live_allocations: usize,
29    /// Active resident bytes.
30    pub resident_bytes: u64,
31    /// Evictions performed to stay inside the bound.
32    pub evictions: usize,
33}
34
35/// Bounded resident arena with oldest-first eviction.
36#[derive(Clone, Debug)]
37pub struct WgpuResidentArena {
38    max_bytes: u64,
39    next_id: u64,
40    resident_bytes: u64,
41    evictions: usize,
42    allocations: VecDeque<LiveAllocation>,
43}
44
45impl WgpuResidentArena {
46    /// Creates an arena with a byte ceiling.
47    pub fn new(max_bytes: u64) -> Self {
48        Self {
49            max_bytes,
50            next_id: 0,
51            resident_bytes: 0,
52            evictions: 0,
53            allocations: VecDeque::new(),
54        }
55    }
56
57    /// Allocates resident bytes, evicting older allocations when necessary.
58    pub fn allocate(&mut self, bytes: u64) -> Result<WgpuArenaAllocation, String> {
59        if bytes > self.max_bytes {
60            return Err("wgpu resident allocation exceeds arena".to_owned());
61        }
62        while self.resident_bytes.saturating_add(bytes) > self.max_bytes {
63            let Some(mut allocation) = self.allocations.pop_front() else {
64                break;
65            };
66            if allocation.active {
67                allocation.active = false;
68                self.resident_bytes = self
69                    .resident_bytes
70                    .saturating_sub(allocation.allocation.bytes);
71                self.evictions += 1;
72            }
73            self.allocations.push_back(allocation);
74        }
75        self.next_id += 1;
76        let allocation = WgpuArenaAllocation {
77            id: WgpuAllocationId(self.next_id),
78            bytes,
79        };
80        self.resident_bytes += bytes;
81        self.allocations.push_back(LiveAllocation {
82            allocation: allocation.clone(),
83            active: true,
84        });
85        Ok(allocation)
86    }
87
88    /// Returns whether an allocation is still resident.
89    pub fn contains(&self, id: WgpuAllocationId) -> bool {
90        self.allocations
91            .iter()
92            .any(|allocation| allocation.active && allocation.allocation.id == id)
93    }
94
95    /// Returns current arena pressure.
96    pub fn snapshot(&self) -> WgpuArenaSnapshot {
97        WgpuArenaSnapshot {
98            live_allocations: self
99                .allocations
100                .iter()
101                .filter(|allocation| allocation.active)
102                .count(),
103            resident_bytes: self.resident_bytes,
104            evictions: self.evictions,
105        }
106    }
107}
sim_lib_compute_wgpu/arena.rs

sim_lib_compute_wgpu/
arena.rs
