cubecl_wgpu/compute/
storage.rs

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

cubecl_wgpu/compute/
storage.rs
