ruda_driver_wgpu/memory/
storage.rs

1use ruda_kernel::dsl::server::IoError;
2use ruda::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, wgpu::Buffer>,
15    device: wgpu::Device,
16    buffer_usages: BufferUsages,
17    mem_alignment: usize,
18    relocation_queue: Option<wgpu::Queue>,
19}
20
21impl core::fmt::Debug for WgpuStorage {
22    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        f.write_str(format!("WgpuStorage {{ device: {:?} }}", self.device).as_str())
24    }
25}
26
27/// The memory resource that can be allocated for wgpu.
28#[derive(new, Debug)]
29pub struct WgpuResource {
30    #[new(default)]
31    pin: Option<ruda::runtime::memory_management::MemoryResourcePin>,
32    /// The wgpu buffer.
33    pub buffer: wgpu::Buffer,
34    /// The buffer offset.
35    pub offset: u64,
36    /// The size of the resource.
37    ///
38    /// # Notes
39    ///
40    /// The result considers the offset.
41    pub size: u64,
42}
43
44impl WgpuResource {
45    pub(crate) fn address_pin(&self) -> Option<ruda::runtime::memory_management::MemoryResourcePin> {
46        self.pin.clone()
47    }
48
49    /// Return the binding view of the buffer.
50    pub fn as_wgpu_bind_resource(&self) -> wgpu::BindingResource<'_> {
51        // wgpu enforces 4-byte alignment for buffer binding sizes per the WebGPU spec.
52        // - https://github.com/gfx-rs/wgpu/pull/8041
53        //
54        // This padding is safe because:
55        // 1. In checked mode, bounds checks prevent reading beyond the logical size.
56        // 2. In unchecked mode, OOB access is already undefined behavior.
57        //
58        // For zero-sized resources, pass None (use rest of buffer from offset).
59        // The allocator guarantees the buffer is at least MIN_BUFFER_SIZE bytes.
60        let size = NonZeroU64::new(self.size.next_multiple_of(4));
61
62        let binding = wgpu::BufferBinding {
63            buffer: &self.buffer,
64            offset: self.offset,
65            size,
66        };
67        wgpu::BindingResource::Buffer(binding)
68    }
69}
70
71/// Keeps actual wgpu buffer references in a hashmap with ids as key.
72impl WgpuStorage {
73    /// Create a new storage on the given [device](wgpu::Device).
74    pub fn new(mem_alignment: usize, device: wgpu::Device, usages: BufferUsages) -> Self {
75        Self {
76            memory: HashMap::new(),
77            device,
78            buffer_usages: usages,
79            mem_alignment,
80            relocation_queue: None,
81        }
82    }
83
84    pub(crate) fn with_relocation_queue(mut self, queue: wgpu::Queue) -> Self {
85        self.relocation_queue = Some(queue);
86        self
87    }
88
89    fn wait_relocation(&self) -> Result<(), IoError> {
90        #[cfg(not(target_family = "wasm"))]
91        {
92            self.device.poll(wgpu::PollType::Wait { submission_index: None, timeout: None })
93                .map(|_| ()).map_err(|error| IoError::Unknown {
94                    description: format!("WGPU relocation wait: {error}"),
95                    backtrace: ruda_core::backtrace::BackTrace::capture(),
96                })
97        }
98        #[cfg(target_family = "wasm")]
99        {
100            Err(IoError::UnsupportedIoOperation { backtrace: ruda_core::backtrace::BackTrace::capture() })
101        }
102    }
103}
104
105impl ComputeStorage for WgpuStorage {
106    type Resource = WgpuResource;
107
108    fn alignment(&self) -> usize {
109        self.mem_alignment
110    }
111
112    fn get_pinned(&mut self, handle: &StorageHandle, binding: ruda::runtime::memory_management::ManagedMemoryBinding) -> Self::Resource {
113        let mut resource = self.get(handle);
114        resource.pin = Some(binding.pin());
115        resource
116    }
117
118    fn supports_relocation(&self) -> bool {
119        self.relocation_queue.is_some() && cfg!(not(target_family = "wasm"))
120    }
121
122    fn relocation_barrier(&mut self) -> Result<(), IoError> {
123        let queue = self.relocation_queue.as_ref().expect("relocation queue");
124        // Submit pending write_buffer work too, even when no compute encoder
125        // had commands. The stream submits its encoder before reserve.
126        queue.submit([]);
127        self.wait_relocation()
128    }
129
130    fn relocation_copy(&mut self, source: &StorageHandle, target: &StorageHandle) -> Result<(), IoError> {
131        let src = self.memory.get(&source.id).expect("relocation source");
132        let dst = self.memory.get(&target.id).expect("relocation target");
133        let size = source.size().next_multiple_of(wgpu::COPY_BUFFER_ALIGNMENT);
134        if source.offset() % wgpu::COPY_BUFFER_ALIGNMENT != 0 || target.offset() % wgpu::COPY_BUFFER_ALIGNMENT != 0
135            || source.offset().checked_add(size).is_none_or(|end| end > src.size())
136            || target.offset().checked_add(size).is_none_or(|end| end > dst.size()) {
137            return Err(IoError::UnsupportedIoOperation { backtrace: ruda_core::backtrace::BackTrace::capture() });
138        }
139        let validation = self.device.push_error_scope(wgpu::ErrorFilter::Validation);
140        let allocation = self.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory);
141        let internal = self.device.push_error_scope(wgpu::ErrorFilter::Internal);
142        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("RUDA adaptive memory copies") });
143        encoder.copy_buffer_to_buffer(src, source.offset(), dst, target.offset(), size);
144        self.relocation_queue.as_ref().expect("relocation queue").submit([encoder.finish()]);
145        let errors = [internal.pop(), allocation.pop(), validation.pop()];
146        #[cfg(not(target_family = "wasm"))]
147        {
148            let mut first = None;
149            for error in errors {
150                if let Some(error) = ruda_core::future::block_on(error) {
151                    if first.is_none() { first = Some(error); }
152                }
153            }
154            if let Some(error) = first {
155                return Err(IoError::Unknown { description: format!("WGPU relocation copy: {error}"),
156                    backtrace: ruda_core::backtrace::BackTrace::capture() });
157            }
158        }
159        #[cfg(target_family = "wasm")]
160        let _ = errors;
161        Ok(())
162    }
163
164    fn relocation_complete(&mut self) -> Result<(), IoError> { self.wait_relocation() }
165
166    fn get(&mut self, handle: &StorageHandle) -> Self::Resource {
167        let buffer = self.memory.get(&handle.id).unwrap();
168        WgpuResource::new(buffer.clone(), handle.offset(), handle.size())
169    }
170
171    #[cfg_attr(
172        feature = "tracing",
173        tracing::instrument(level = "trace", skip(self, size))
174    )]
175    fn alloc(&mut self, size: u64) -> Result<StorageHandle, IoError> {
176        let id = StorageId::new();
177
178        let alloc_size = size.max(MIN_BUFFER_SIZE);
179
180        let scopes = self.relocation_queue.as_ref().map(|_| [
181            self.device.push_error_scope(wgpu::ErrorFilter::Validation),
182            self.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory),
183            self.device.push_error_scope(wgpu::ErrorFilter::Internal),
184        ]);
185        let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
186            label: None,
187            size: alloc_size,
188            usage: self.buffer_usages,
189            mapped_at_creation: false,
190        });
191
192        if let Some(scopes) = scopes {
193            let errors = scopes.into_iter().rev().map(|scope| scope.pop()).collect::<Vec<_>>();
194            #[cfg(not(target_family = "wasm"))]
195            {
196                let mut first = None;
197                for error in errors {
198                    if let Some(error) = ruda_core::future::block_on(error) {
199                        if first.is_none() { first = Some(error); }
200                    }
201                }
202                if let Some(error) = first {
203                    return Err(IoError::Unknown { description: format!("WGPU adaptive allocation: {error}"),
204                        backtrace: ruda_core::backtrace::BackTrace::capture() });
205                }
206            }
207            #[cfg(target_family = "wasm")]
208            let _ = errors;
209        }
210
211        self.memory.insert(id, buffer);
212        Ok(StorageHandle::new(
213            id,
214            StorageUtilization { offset: 0, size },
215        ))
216    }
217
218    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
219    fn dealloc(&mut self, id: StorageId) {
220        self.memory.remove(&id);
221    }
222
223    fn flush(&mut self) {
224        // We don't wait for dealloc
225    }
226}
ruda_driver_wgpu/memory/storage.rs

ruda_driver_wgpu/memory/
storage.rs
