ranim_render/
utils.rs

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

ranim_render/
utils.rs
