dspk_ops/core/
scratch.rs

1use cubecl::prelude::*;
2use cubecl::server::Handle;
3
4use super::buffer;
5
6/// A device buffer reused across calls, reallocated only when a call needs more than it holds.
7#[derive(Debug, Clone, Default)]
8pub struct Scratch {
9    handle: Option<Handle>,
10    capacity_bytes: usize,
11}
12
13impl Scratch {
14    pub fn new() -> Self {
15        Self::default()
16    }
17
18    /// A buffer of at least `len` elements of `E` (contents unspecified).
19    pub fn get<E: CubeElement>(&mut self, client: &Client, len: usize) -> Handle {
20        let need = buffer::bytes::<E>(len);
21        match &self.handle {
22            Some(handle) if self.capacity_bytes >= need => handle.clone(),
23            _ => {
24                let handle = client.empty(need);
25                self.handle = Some(handle.clone());
26                self.capacity_bytes = need;
27                handle
28            }
29        }
30    }
31}
dspk_ops/core/scratch.rs

dspk_ops/core/
scratch.rs
