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

ranim_render/
utils.rs
