ranim_render/
utils.rs

1use std::{fmt::Debug, marker::PhantomData, ops::Deref};
2
3use tracing::{info, warn};
4use wgpu::util::DeviceExt;
5
6/// Wgpu context
7#[derive(bevy_ecs::prelude::Resource, Clone)]
8pub struct WgpuContext {
9    /// The wgpu instance
10    pub instance: wgpu::Instance,
11    /// The wgpu adapter
12    pub adapter: wgpu::Adapter,
13    /// The wgpu device
14    pub device: wgpu::Device,
15    /// The wgpu queue
16    pub queue: wgpu::Queue,
17}
18
19impl WgpuContext {
20    /// Create a new wgpu context
21    pub async fn new() -> Self {
22        let instance = wgpu::Instance::default();
23        let adapter = instance
24            .request_adapter(&wgpu::RequestAdapterOptions {
25                power_preference: wgpu::PowerPreference::HighPerformance,
26                ..Default::default()
27            })
28            .await
29            .unwrap();
30        info!("wgpu adapter info: {:?}", adapter.get_info());
31        let required_limits = adapter.limits();
32
33        #[cfg(feature = "profiling")]
34        let (device, queue) = adapter
35            .request_device(&wgpu::DeviceDescriptor {
36                required_features: wgpu_profiler::GpuProfiler::ALL_WGPU_TIMER_FEATURES,
37                required_limits,
38                ..Default::default()
39            })
40            .await
41            .unwrap();
42        #[cfg(not(feature = "profiling"))]
43        let (device, queue) = adapter
44            .request_device(&wgpu::DeviceDescriptor {
45                required_limits,
46                ..Default::default()
47            })
48            .await
49            .unwrap();
50
51        Self {
52            instance,
53            adapter,
54            device,
55            queue,
56        }
57    }
58}
59
60#[allow(unused)]
61pub(crate) struct WgpuBuffer<T: bytemuck::Pod + bytemuck::Zeroable + Debug> {
62    label: Option<&'static str>,
63    buffer: wgpu::Buffer,
64    usage: wgpu::BufferUsages,
65    inner: T,
66}
67
68impl<T: bytemuck::Pod + bytemuck::Zeroable + Debug> AsRef<wgpu::Buffer> for WgpuBuffer<T> {
69    fn as_ref(&self) -> &wgpu::Buffer {
70        &self.buffer
71    }
72}
73
74#[allow(unused)]
75impl<T: bytemuck::Pod + bytemuck::Zeroable + Debug> WgpuBuffer<T> {
76    pub(crate) fn new_init(
77        ctx: &WgpuContext,
78        label: Option<&'static str>,
79        usage: wgpu::BufferUsages,
80        data: T,
81    ) -> Self {
82        assert!(
83            usage.contains(wgpu::BufferUsages::COPY_DST),
84            "Buffer {label:?} does not contains COPY_DST"
85        );
86        // trace!("[WgpuBuffer]: new_init, {} {:?}", data.len(), usage);
87        Self {
88            label,
89            buffer: ctx
90                .device
91                .create_buffer_init(&wgpu::util::BufferInitDescriptor {
92                    label,
93                    contents: bytemuck::bytes_of(&data),
94                    usage,
95                }),
96            usage,
97            inner: data,
98        }
99    }
100
101    pub(crate) fn get(&self) -> &T {
102        &self.inner
103    }
104
105    pub(crate) fn set(&mut self, ctx: &WgpuContext, data: T) {
106        {
107            let mut view = ctx
108                .queue
109                .write_buffer_with(
110                    &self.buffer,
111                    0,
112                    wgpu::BufferSize::new(std::mem::size_of_val(&data) as u64).unwrap(),
113                )
114                .unwrap();
115            view.copy_from_slice(bytemuck::bytes_of(&data));
116        }
117        // ctx.queue.submit([]);
118        self.inner = data;
119    }
120
121    #[allow(unused)]
122    pub(crate) fn read_buffer(&self, ctx: &WgpuContext) -> Vec<u8> {
123        let size = std::mem::size_of::<T>();
124        let staging_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
125            label: Some("Debug Staging Buffer"),
126            size: size as u64,
127            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
128            mapped_at_creation: false,
129        });
130
131        let mut encoder = ctx
132            .device
133            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
134                label: Some("Debug Read Encoder"),
135            });
136
137        encoder.copy_buffer_to_buffer(&self.buffer, 0, &staging_buffer, 0, size as u64);
138        ctx.queue.submit(Some(encoder.finish()));
139
140        let buffer_slice = staging_buffer.slice(..);
141        let (tx, rx) = async_channel::bounded(1);
142        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
143            pollster::block_on(tx.send(result)).unwrap()
144        });
145        ctx.device
146            .poll(wgpu::PollType::wait_indefinitely())
147            .unwrap();
148        pollster::block_on(rx.recv()).unwrap().unwrap();
149
150        buffer_slice.get_mapped_range().unwrap().to_vec()
151    }
152}
153
154pub(crate) struct WgpuVecBuffer<T: Default + bytemuck::Pod + bytemuck::Zeroable + Debug> {
155    label: Option<&'static str>,
156    pub(crate) buffer: wgpu::Buffer,
157    usage: wgpu::BufferUsages,
158    /// Keep match to the buffer size
159    len: usize,
160    _phantom: PhantomData<T>,
161    // inner: Vec<T>,
162}
163
164impl<T: Default + bytemuck::Pod + bytemuck::Zeroable + Debug> WgpuVecBuffer<T> {
165    pub(crate) fn new(
166        ctx: &WgpuContext,
167        label: Option<&'static str>,
168        usage: wgpu::BufferUsages,
169        len: usize,
170    ) -> Self {
171        assert!(
172            usage.contains(wgpu::BufferUsages::COPY_DST),
173            "Buffer {label:?} does not contains COPY_DST"
174        );
175        let size = (std::mem::size_of::<T>() * len) as u64;
176        Self {
177            label,
178            buffer: ctx.device.create_buffer(&wgpu::BufferDescriptor {
179                label,
180                size,
181                usage,
182                mapped_at_creation: false,
183            }),
184            usage,
185            len: 0,
186            _phantom: PhantomData,
187            // inner: vec![],
188        }
189    }
190
191    #[allow(unused)]
192    pub(crate) fn new_init(
193        ctx: &WgpuContext,
194        label: Option<&'static str>,
195        usage: wgpu::BufferUsages,
196        data: &[T],
197    ) -> Self {
198        let mut buffer = Self::new(ctx, label, usage, data.len());
199        buffer.set(ctx, data);
200        buffer
201    }
202
203    #[allow(unused)]
204    pub(crate) fn len(&self) -> usize {
205        self.len
206    }
207    // pub(crate) fn get(&self) -> &[T] {
208    //     self.inner.as_ref()
209    // }
210
211    #[allow(unused)]
212    pub(crate) fn resize(&mut self, ctx: &WgpuContext, len: usize) -> bool {
213        let size = (std::mem::size_of::<T>() * len) as u64;
214        let realloc = self.buffer.size() != size;
215        if realloc {
216            self.len = len;
217            // self.inner.resize(len, T::default());
218            self.buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
219                label: self.label,
220                size,
221                usage: self.usage,
222                mapped_at_creation: false,
223            })
224        }
225        realloc
226    }
227
228    pub(crate) fn set(&mut self, ctx: &WgpuContext, data: &[T]) -> bool {
229        // trace!("{} {}", self.inner.len(), data.len());
230        // self.inner.resize(data.len(), T::default());
231        // self.inner.copy_from_slice(data);
232        self.len = data.len();
233        let realloc = self.buffer.size() != std::mem::size_of_val(data) as u64;
234
235        if realloc {
236            // info!("realloc");
237            // NOTE: create_buffer_init sometimes causes freezing in wasm
238            let buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
239                label: self.label,
240                size: (std::mem::size_of_val(data)) as u64,
241                usage: self.usage,
242                mapped_at_creation: false,
243            });
244            ctx.queue
245                .write_buffer(&buffer, 0, bytemuck::cast_slice(data));
246            // info!("new");
247            self.buffer = buffer;
248        } else {
249            // info!("queue copy");
250            {
251                let mut view = ctx
252                    .queue
253                    .write_buffer_with(
254                        &self.buffer,
255                        0,
256                        wgpu::BufferSize::new((std::mem::size_of_val(data)) as u64).unwrap(),
257                    )
258                    .unwrap();
259                view.copy_from_slice(bytemuck::cast_slice(data));
260            }
261            // ctx.queue.submit([]);
262        }
263        // info!("done");
264        realloc
265    }
266
267    #[allow(unused)]
268    pub(crate) fn read_buffer(&self, ctx: &WgpuContext) -> Option<Vec<u8>> {
269        let size = std::mem::size_of::<T>() * self.len;
270        let staging_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
271            label: Some("Debug Staging Buffer"),
272            size: size as u64,
273            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
274            mapped_at_creation: false,
275        });
276
277        let mut encoder = ctx
278            .device
279            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
280                label: Some("Debug Read Encoder"),
281            });
282
283        encoder.copy_buffer_to_buffer(&self.buffer, 0, &staging_buffer, 0, size as u64);
284        ctx.queue.submit(Some(encoder.finish()));
285
286        let buffer_slice = staging_buffer.slice(..);
287        let (tx, rx) = async_channel::bounded(1);
288        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
289            tx.try_send(result).unwrap()
290        });
291        ctx.device
292            .poll(wgpu::PollType::wait_indefinitely())
293            .unwrap();
294        pollster::block_on(rx.recv()).unwrap().unwrap();
295
296        let x = buffer_slice.get_mapped_range().unwrap().to_vec();
297        Some(x)
298    }
299}
300
301pub struct WgpuTexture {
302    inner: wgpu::Texture,
303}
304
305impl WgpuTexture {
306    pub fn new(ctx: &WgpuContext, desc: &wgpu::TextureDescriptor) -> Self {
307        Self {
308            inner: ctx.device.create_texture(desc),
309        }
310    }
311}
312
313impl Deref for WgpuTexture {
314    type Target = wgpu::Texture;
315    fn deref(&self) -> &Self::Target {
316        &self.inner
317    }
318}
319
320/// A [`WgpuTexture`] with [`wgpu::TextureUsages::COPY_SRC`] usage and wrapped with a staging buffer and
321/// a cpu side bytes `Vec<T>` buffer to read back from the texture.
322pub struct ReadbackWgpuTexture {
323    inner: WgpuTexture,
324    aligned_bytes_per_row: usize,
325    staging_buffer: wgpu::Buffer,
326    bytes: Vec<u8>,
327    /// Pending async readback receiver. Present when `start_readback` has been called
328    /// but `finish_readback` has not yet completed.
329    pending_rx: Option<async_channel::Receiver<Result<(), wgpu::BufferAsyncError>>>,
330}
331
332impl Deref for ReadbackWgpuTexture {
333    type Target = WgpuTexture;
334    fn deref(&self) -> &Self::Target {
335        &self.inner
336    }
337}
338
339const ALIGNMENT: usize = 256;
340impl ReadbackWgpuTexture {
341    pub fn new(ctx: &WgpuContext, desc: &wgpu::TextureDescriptor) -> Self {
342        if !desc.usage.contains(wgpu::TextureUsages::COPY_SRC) {
343            warn!(
344                "ReadbackWgpuTexture should have COPY_SRC usage, but got {:?}, will auto add this usage",
345                desc.usage
346            );
347        }
348        let texture = WgpuTexture::new(
349            ctx,
350            &wgpu::TextureDescriptor {
351                usage: desc.usage | wgpu::TextureUsages::COPY_SRC,
352                ..*desc
353            },
354        );
355        let block_size = desc.format.block_copy_size(None).unwrap();
356        let bytes_per_row =
357            (texture.size().width * block_size).div_ceil(ALIGNMENT as u32) as usize * ALIGNMENT;
358
359        let staging_buffer_label = desc.label.map(|s| format!("{s} Staging Buffer"));
360        let staging_buffer = ctx.device.create_buffer(&wgpu::BufferDescriptor {
361            label: staging_buffer_label.as_deref(),
362            size: (bytes_per_row * texture.size().height as usize) as u64,
363            usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
364            mapped_at_creation: false,
365        });
366        let len = texture.size().width * texture.size().height * block_size;
367        let bytes = vec![0u8; len as usize];
368
369        Self {
370            inner: texture,
371            aligned_bytes_per_row: bytes_per_row,
372            staging_buffer,
373            bytes,
374            pending_rx: None,
375        }
376    }
377    pub fn texture_data(&self) -> &[u8] {
378        &self.bytes
379    }
380
381    /// Start an async readback: copy texture to staging buffer, submit, and begin mapping.
382    ///
383    /// This is non-blocking. Call [`finish_readback`](Self::finish_readback) later to
384    /// poll the device and copy the data into the CPU-side buffer.
385    pub fn start_readback(&mut self, ctx: &WgpuContext) {
386        let size = self.size();
387
388        let mut encoder = ctx
389            .device
390            .create_command_encoder(&wgpu::CommandEncoderDescriptor {
391                label: Some("Readback Copy Encoder"),
392            });
393        encoder.copy_texture_to_buffer(
394            wgpu::TexelCopyTextureInfo {
395                aspect: wgpu::TextureAspect::All,
396                texture: self,
397                mip_level: 0,
398                origin: wgpu::Origin3d::ZERO,
399            },
400            wgpu::TexelCopyBufferInfo {
401                buffer: &self.staging_buffer,
402                layout: wgpu::TexelCopyBufferLayout {
403                    offset: 0,
404                    bytes_per_row: Some(self.aligned_bytes_per_row as u32),
405                    rows_per_image: Some(size.height),
406                },
407            },
408            size,
409        );
410        ctx.queue.submit(Some(encoder.finish()));
411
412        let buffer_slice = self.staging_buffer.slice(..);
413        let (tx, rx) = async_channel::bounded(1);
414        buffer_slice.map_async(wgpu::MapMode::Read, move |result| {
415            let _ = tx.try_send(result);
416        });
417        self.pending_rx = Some(rx);
418    }
419
420    /// Finish a pending async readback: poll the device, copy data from the staging
421    /// buffer into the CPU-side buffer, and unmap.
422    ///
423    /// If no readback is pending, this is a no-op.
424    pub fn finish_readback(&mut self, ctx: &WgpuContext) {
425        let Some(rx) = self.pending_rx.take() else {
426            return;
427        };
428
429        ctx.device
430            .poll(wgpu::PollType::wait_indefinitely())
431            .unwrap();
432        pollster::block_on(rx.recv()).unwrap().unwrap();
433
434        self.copy_staging_to_bytes();
435    }
436
437    /// Try to finish a pending readback without blocking.
438    /// Returns `true` if completed (or nothing was pending), `false` if GPU isn't done yet.
439    pub fn try_finish_readback(&mut self, ctx: &WgpuContext) -> bool {
440        let Some(rx) = self.pending_rx.as_ref() else {
441            return true;
442        };
443
444        // Non-blocking poll to nudge the GPU
445        let _ = ctx.device.poll(wgpu::PollType::Poll);
446
447        // Check if the mapping callback has fired
448        match rx.try_recv() {
449            Ok(result) => {
450                result.unwrap();
451                self.pending_rx = None;
452                self.copy_staging_to_bytes();
453                true
454            }
455            Err(async_channel::TryRecvError::Empty) => false,
456            Err(async_channel::TryRecvError::Closed) => {
457                self.pending_rx = None;
458                true
459            }
460        }
461    }
462
463    fn copy_staging_to_bytes(&mut self) {
464        let size = self.size();
465        let buffer_slice = self.staging_buffer.slice(..);
466        let view = buffer_slice.get_mapped_range().unwrap();
467        let block_size = self.inner.format().block_copy_size(None).unwrap();
468        let bytes_in_row = (size.width * block_size) as usize;
469
470        for y in 0..size.height as usize {
471            let src_row_start = y * self.aligned_bytes_per_row;
472            let dst_row_start = y * bytes_in_row;
473
474            self.bytes[dst_row_start..dst_row_start + bytes_in_row]
475                .copy_from_slice(&view[src_row_start..src_row_start + bytes_in_row]);
476        }
477        drop(view);
478        self.staging_buffer.unmap();
479    }
480
481    /// Synchronous readback: start + finish in one call.
482    pub fn update_texture_data(&mut self, ctx: &WgpuContext) -> &[u8] {
483        self.start_readback(ctx);
484        self.finish_readback(ctx);
485        &self.bytes
486    }
487}
488
489#[cfg(test)]
490mod test {
491    #[test]
492    fn test() {
493        // let x = vec![0, 1, 2, 3];
494        // assert_eq!(
495        //     bytemuck::bytes_of(&[x.as_slice()]),
496        //     bytemuck::bytes_of(&x)
497        // )
498    }
499}
ranim_render/utils.rs

ranim_render/
utils.rs
