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 encode_relocation_copy(&self, encoder: &mut wgpu::CommandEncoder, source: &StorageHandle, target: &StorageHandle) -> Result<(), IoError> {
90        let src = self.memory.get(&source.id).expect("relocation source");
91        let dst = self.memory.get(&target.id).expect("relocation target");
92        let size = source.size().next_multiple_of(wgpu::COPY_BUFFER_ALIGNMENT);
93        if source.offset() % wgpu::COPY_BUFFER_ALIGNMENT != 0 || target.offset() % wgpu::COPY_BUFFER_ALIGNMENT != 0
94            || source.offset().checked_add(size).is_none_or(|end| end > src.size())
95            || target.offset().checked_add(size).is_none_or(|end| end > dst.size()) {
96            return Err(IoError::UnsupportedIoOperation { backtrace: ruda_core::backtrace::BackTrace::capture() });
97        }
98        encoder.copy_buffer_to_buffer(src, source.offset(), dst, target.offset(), size);
99        Ok(())
100    }
101
102    fn wait_relocation(&self) -> Result<(), IoError> {
103        #[cfg(not(target_family = "wasm"))]
104        {
105            self.device.poll(wgpu::PollType::Wait { submission_index: None, timeout: None })
106                .map(|_| ()).map_err(|error| IoError::Unknown {
107                    description: format!("WGPU relocation wait: {error}"),
108                    backtrace: ruda_core::backtrace::BackTrace::capture(),
109                })
110        }
111        #[cfg(target_family = "wasm")]
112        {
113            Err(IoError::UnsupportedIoOperation { backtrace: ruda_core::backtrace::BackTrace::capture() })
114        }
115    }
116}
117
118impl ComputeStorage for WgpuStorage {
119    type Resource = WgpuResource;
120
121    fn alignment(&self) -> usize {
122        self.mem_alignment
123    }
124
125    fn get_pinned(&mut self, handle: &StorageHandle, binding: ruda::runtime::memory_management::ManagedMemoryBinding) -> Self::Resource {
126        let mut resource = self.get(handle);
127        resource.pin = Some(binding.pin());
128        resource
129    }
130
131    fn supports_relocation(&self) -> bool {
132        self.relocation_queue.is_some() && cfg!(not(target_family = "wasm"))
133    }
134
135    fn relocation_barrier(&mut self) -> Result<(), IoError> {
136        let queue = self.relocation_queue.as_ref().expect("relocation queue");
137        // Submit pending write_buffer work too, even when no compute encoder
138        // had commands. The stream submits its encoder before reserve.
139        queue.submit([]);
140        self.wait_relocation()
141    }
142
143    fn relocation_copy(&mut self, source: &StorageHandle, target: &StorageHandle) -> Result<(), IoError> {
144        self.relocation_copy_batch(core::iter::once((source, target)))
145    }
146
147    fn relocation_copy_batch<'a>(
148        &mut self,
149        copies: impl IntoIterator<Item = (&'a StorageHandle, &'a StorageHandle)>,
150    ) -> Result<(), IoError> {
151        let mut copies = copies.into_iter().peekable();
152        if copies.peek().is_none() { return Ok(()); }
153        let validation = self.device.push_error_scope(wgpu::ErrorFilter::Validation);
154        let allocation = self.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory);
155        let internal = self.device.push_error_scope(wgpu::ErrorFilter::Internal);
156        let mut encoder = self.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("RUDA adaptive memory copies") });
157        let copied = copies.try_for_each(|(source, target)| self.encode_relocation_copy(&mut encoder, source, target));
158        self.relocation_queue.as_ref().expect("relocation queue").submit([encoder.finish()]);
159        let errors = [internal.pop(), allocation.pop(), validation.pop()];
160        #[cfg(not(target_family = "wasm"))]
161        {
162            let mut first = None;
163            for error in errors {
164                if let Some(error) = ruda_core::future::block_on(error) {
165                    if first.is_none() { first = Some(error); }
166                }
167            }
168            copied?;
169            if let Some(error) = first {
170                return Err(IoError::Unknown { description: format!("WGPU relocation copy: {error}"),
171                    backtrace: ruda_core::backtrace::BackTrace::capture() });
172            }
173        }
174        #[cfg(target_family = "wasm")]
175        { let _ = errors; copied?; }
176        Ok(())
177    }
178
179    fn relocation_complete(&mut self) -> Result<(), IoError> { self.wait_relocation() }
180
181    fn get(&mut self, handle: &StorageHandle) -> Self::Resource {
182        let buffer = self.memory.get(&handle.id).unwrap();
183        WgpuResource::new(buffer.clone(), handle.offset(), handle.size())
184    }
185
186    #[cfg_attr(
187        feature = "tracing",
188        tracing::instrument(level = "trace", skip(self, size))
189    )]
190    fn alloc(&mut self, size: u64) -> Result<StorageHandle, IoError> {
191        let id = StorageId::new();
192
193        let alloc_size = size.max(MIN_BUFFER_SIZE);
194
195        let scopes = self.relocation_queue.as_ref().map(|_| [
196            self.device.push_error_scope(wgpu::ErrorFilter::Validation),
197            self.device.push_error_scope(wgpu::ErrorFilter::OutOfMemory),
198            self.device.push_error_scope(wgpu::ErrorFilter::Internal),
199        ]);
200        let buffer = self.device.create_buffer(&wgpu::BufferDescriptor {
201            label: None,
202            size: alloc_size,
203            usage: self.buffer_usages,
204            mapped_at_creation: false,
205        });
206
207        if let Some(scopes) = scopes {
208            let errors = scopes.into_iter().rev().map(|scope| scope.pop()).collect::<Vec<_>>();
209            #[cfg(not(target_family = "wasm"))]
210            {
211                let mut first = None;
212                for error in errors {
213                    if let Some(error) = ruda_core::future::block_on(error) {
214                        if first.is_none() { first = Some(error); }
215                    }
216                }
217                if let Some(error) = first {
218                    return Err(IoError::Unknown { description: format!("WGPU adaptive allocation: {error}"),
219                        backtrace: ruda_core::backtrace::BackTrace::capture() });
220                }
221            }
222            #[cfg(target_family = "wasm")]
223            let _ = errors;
224        }
225
226        self.memory.insert(id, buffer);
227        Ok(StorageHandle::new(
228            id,
229            StorageUtilization { offset: 0, size },
230        ))
231    }
232
233    #[cfg_attr(feature = "tracing", tracing::instrument(level = "trace", skip(self)))]
234    fn dealloc(&mut self, id: StorageId) {
235        self.memory.remove(&id);
236    }
237
238    fn flush(&mut self) {
239        // We don't wait for dealloc
240    }
241}
ruda_driver_wgpu/memory/storage.rs

ruda_driver_wgpu/memory/
storage.rs
