ferrotensor_backend/
pool.rs

1use std::collections::HashMap;
2use std::sync::{Arc, Mutex};
3use wgpu::{Buffer, BufferDescriptor, BufferUsages, Device};
4
5/// Default cap on the bytes a pool keeps for reuse (512 MiB).
6pub const DEFAULT_MAX_POOLED_BYTES: u64 = 512 * 1024 * 1024;
7
8/// Free lists of GPU buffers keyed by exact byte size.
9///
10/// Returned buffers are kept for reuse until the pool holds
11/// `max_pooled_bytes`; beyond that they are released. Without the cap a
12/// workload with many distinct tensor sizes (e.g. training on sequences of
13/// varying length) would keep every buffer it ever freed.
14pub struct BufferPool {
15    buffers: Mutex<PoolState>,
16    usage: BufferUsages,
17}
18
19#[derive(Default)]
20struct PoolState {
21    free: HashMap<u64, Vec<Arc<Buffer>>>,
22    pooled_bytes: u64,
23    max_pooled_bytes: u64,
24}
25
26impl std::fmt::Debug for BufferPool {
27    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
28        let state = self.buffers.lock().unwrap();
29        let counts: HashMap<u64, usize> = state
30            .free
31            .iter()
32            .map(|(sz, list)| (*sz, list.len()))
33            .collect();
34        f.debug_struct("BufferPool")
35            .field("usage", &self.usage)
36            .field("buffer_counts", &counts)
37            .field("pooled_bytes", &state.pooled_bytes)
38            .finish()
39    }
40}
41
42impl BufferPool {
43    /// A pool whose buffers are created with `usage`.
44    pub fn new(usage: BufferUsages) -> Self {
45        Self {
46            buffers: Mutex::new(PoolState {
47                max_pooled_bytes: DEFAULT_MAX_POOLED_BYTES,
48                ..PoolState::default()
49            }),
50            usage,
51        }
52    }
53
54    /// Change how many bytes of freed buffers the pool keeps (0 disables
55    /// reuse). Buffers above the new cap are released immediately.
56    pub fn set_max_pooled_bytes(&self, max: u64) {
57        let mut state = self.buffers.lock().unwrap();
58        state.max_pooled_bytes = max;
59        if state.pooled_bytes > max {
60            state.free.clear();
61            state.pooled_bytes = 0;
62        }
63    }
64
65    /// Bytes currently held for reuse.
66    pub fn pooled_bytes(&self) -> u64 {
67        self.buffers.lock().unwrap().pooled_bytes
68    }
69
70    /// A free buffer of exactly `size` bytes, or a new one.
71    pub fn get(&self, device: &Device, size: u64) -> Arc<Buffer> {
72        let mut state = self.buffers.lock().unwrap();
73        if let Some(buffer) = state.free.get_mut(&size).and_then(Vec::pop) {
74            state.pooled_bytes -= size;
75            return buffer;
76        }
77
78        // Create new if none available
79        Arc::new(device.create_buffer(&BufferDescriptor {
80            label: Some("Pooled Buffer"),
81            size,
82            usage: self.usage,
83            mapped_at_creation: false,
84        }))
85    }
86
87    /// Returns `buffer` (of `size` bytes) to the pool for reuse.
88    pub fn return_buffer(&self, size: u64, buffer: Arc<Buffer>) {
89        let mut state = self.buffers.lock().unwrap();
90        if state.pooled_bytes + size > state.max_pooled_bytes {
91            // Dropping the last reference releases the GPU memory.
92            return;
93        }
94        state.pooled_bytes += size;
95        state.free.entry(size).or_default().push(buffer);
96    }
97
98    /// Drops every pooled buffer.
99    pub fn clear(&self) {
100        let mut state = self.buffers.lock().unwrap();
101        state.free.clear();
102        state.pooled_bytes = 0;
103    }
104}
ferrotensor_backend/pool.rs

ferrotensor_backend/
pool.rs
