cubecl_wgpu/compute/
storage.rs

1use cubecl_core::server::IoError;
2use cubecl_environment::backtrace::BackTrace;
3use cubecl_environment::collections::HashMap;
4use cubecl_runtime::storage::{ComputeStorage, StorageHandle, StorageId, StorageUtilization};
5use std::num::NonZeroU64;
6use wgpu::BufferUsages;
7
8/// Minimum buffer size in bytes. The WebGPU spec requires buffer sizes > 0, and shaders
9/// declare typed arrays (e.g. `array<vec4<f32>>`) that impose a minimum binding size.
10/// 32 bytes covers the largest possible binding type (`vec4<f64>`).
11const MIN_BUFFER_SIZE: u64 = 32;
12
13/// Buffer storage for wgpu.
14pub struct WgpuStorage {
15    memory: HashMap<StorageId, WgpuMemory>,
16    device: wgpu::Device,
17    buffer_usages: BufferUsages,
18    mem_alignment: usize,
19    #[allow(unused, reason = "keep it simple")]
20    vk_storage: bool,
21}
22
23impl core::fmt::Debug for WgpuStorage {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.write_str(format!("WgpuStorage {{ device: {:?} }}", self.device).as_str())
26    }
27}
28
29/// The memory resource that can be allocated for wgpu.
30#[derive(new, Debug, Clone)]
31pub struct WgpuResource {
32    /// The wgpu buffer.
33    pub buffer: wgpu::Buffer,
34    /// The buffer device address, if supported
35    pub address: Option<NonZeroU64>,
36    /// The buffer offset.
37    pub offset: u64,
38    /// The size of the resource.
39    ///
40    /// # Notes
41    ///
42    /// The result considers the offset.
43    pub size: u64,
44}
45
46/// The memory that can be allocated for wgpu.
47#[derive(new, Debug)]
48pub struct WgpuMemory {
49    /// The wgpu buffer.
50    pub buffer: wgpu::Buffer,
51    /// The buffer device address, if supported
52    pub address: Option<NonZeroU64>,
53}
54
55impl WgpuResource {
56    /// Return the binding view of the buffer.
57    pub fn as_wgpu_bind_resource(&self) -> wgpu::BindingResource<'_> {
58        // wgpu enforces 4-byte alignment for buffer binding sizes per the WebGPU spec.
59        // - https://github.com/gfx-rs/wgpu/pull/8041
60        //
61        // This padding is safe because:
62        // 1. In checked mode, bounds checks prevent reading beyond the logical size.
63        // 2. In unchecked mode, OOB access is already undefined behavior.
64        //
65        // For zero-sized resources, pass None (use rest of buffer from offset).
66        // The allocator guarantees the buffer is at least MIN_BUFFER_SIZE bytes.
67        let size = NonZeroU64::new(self.size.next_multiple_of(4));
68
69        let binding = wgpu::BufferBinding {
70            buffer: &self.buffer,
71            offset: self.offset,
72            size,
73        };
74        wgpu::BindingResource::Buffer(binding)
75    }
76}
77
78/// Keeps actual wgpu buffer references in a hashmap with ids as key.
79impl WgpuStorage {
80    /// Create a new storage on the given [device](wgpu::Device).
81    pub fn new(
82        mem_alignment: usize,
83        device: wgpu::Device,
84        usages: BufferUsages,
85        vk_storage: bool,
86    ) -> Self {
87        Self {
88            memory: HashMap::new(),
89            device,
90            buffer_usages: usages,
91            mem_alignment,
92            vk_storage,
93        }
94    }
95}
96
97impl ComputeStorage for WgpuStorage {
98    type Resource = WgpuResource;
99
100    fn alignment(&self) -> usize {
101        self.mem_alignment
102    }
103
104    fn get(&mut self, handle: &StorageHandle) -> Result<Self::Resource, IoError> {
105        let memory = self
106            .memory
107            .get(&handle.id)
108            .ok_or_else(|| IoError::StorageHandleNotFound {
109                reason: format!("{} in the wgpu buffer storage", handle.id).into(),
110                backtrace: BackTrace::capture(),
111            })?;
112        Ok(WgpuResource::new(
113            memory.buffer.clone(),
114            memory.address,
115            handle.offset(),
116            handle.size(),
117        ))
118    }
119
120    #[cfg_attr(
121        feature = "tracing",
122        tracing::instrument(level = "trace", skip(self, size))
123    )]
124    fn alloc(&mut self, size: u64) -> Result<StorageHandle, IoError> {
125        let id = StorageId::new();
126
127        let alloc_size = size.max(MIN_BUFFER_SIZE);
128
129        let memory = self.create_buffer(&wgpu::BufferDescriptor {
130            label: None,
131            size: alloc_size,
132            usage: self.buffer_usages,
133            mapped_at_creation: false,
134        })?;
135
136        self.memory.insert(id, memory);
137        Ok(StorageHandle::new(
138            id,
139            StorageUtilization { offset: 0, size },
140        ))
141    }
142
143    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
144    fn dealloc(&mut self, id: StorageId) {
145        self.memory.remove(&id);
146    }
147
148    fn flush(&mut self) {
149        // We don't wait for dealloc
150    }
151}
152
153impl WgpuStorage {
154    #[cfg(feature = "spirv")]
155    fn create_buffer(&self, desc: &wgpu::BufferDescriptor<'_>) -> Result<WgpuMemory, IoError> {
156        if self.vk_storage {
157            // wgpu currently doesn't expose this, even though it's used internally for acceleration
158            // structures. While we could use the acceleration structure input flag, that would
159            // then require ray tracing to be supported. So we need to allocate manually, then import
160            // the native buffer into wgpu using `from_raw_managed`.
161            // This actually skips some of the buffer batching stuff we don't really want in `gpu_allocator`.
162            let (buffer, addr) = crate::backend::vulkan::create_storage_buffer(&self.device, desc)?;
163            Ok(WgpuMemory::new(buffer, NonZeroU64::new(addr)))
164        } else {
165            Ok(WgpuMemory::new(self.device.create_buffer(desc), None))
166        }
167    }
168
169    #[cfg(not(feature = "spirv"))]
170    fn create_buffer(&self, desc: &wgpu::BufferDescriptor<'_>) -> Result<WgpuMemory, IoError> {
171        Ok(WgpuMemory::new(self.device.create_buffer(desc), None))
172    }
173}
cubecl_wgpu/compute/storage.rs

cubecl_wgpu/compute/
storage.rs
