flashraw_gpu/
image.rs

1//! GPU-resident images.
2
3use crate::context::GpuContext;
4use flashraw_core::{Error, Rect, Result};
5use std::sync::Arc;
6
7/// Element format of a [`GpuImage`] buffer.
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
9pub enum GpuPixelFormat {
10    /// u16 samples, two per 32-bit word (byte-identical to a host `Vec<u16>`).
11    U16,
12    /// f32 samples.
13    F32,
14    /// sRGB RGBA8, one pixel per 32-bit word (`R | G << 8 | B << 16 | A << 24`).
15    Rgba8,
16}
17
18impl GpuPixelFormat {
19    pub fn bytes_per_sample(self) -> usize {
20        match self {
21            GpuPixelFormat::U16 => 2,
22            GpuPixelFormat::F32 => 4,
23            GpuPixelFormat::Rgba8 => 1,
24        }
25    }
26}
27
28/// One band of a [`GpuImage`]: the pixels of `rect` (image coordinates), stored
29/// row-major with `rect.width * channels` samples per row from byte 0 of `buffer`.
30#[derive(Debug)]
31pub struct GpuTile {
32    pub rect: Rect,
33    pub buffer: wgpu::Buffer,
34}
35
36/// An image living in GPU storage buffers (`STORAGE | COPY_SRC`, plus `MAP_READ` on
37/// unified-memory adapters), on the device of the [`GpuContext`] that produced it.
38///
39/// Most images are a single buffer: use [`GpuImage::buffer`] to bind it in your own
40/// wgpu/WebGPU pipelines, [`GpuImage::to_texture`] for an `Rgba8Unorm` texture, or
41/// download it. Images larger than the adapter's storage-binding limit are **tiled**
42/// ([`GpuImage::is_tiled`]): [`GpuImage::tiles`] lists one buffer per band with its
43/// rectangle in the image (full-width row bands, or full-height column bands for
44/// RGBA8 outputs rotated by 90/270 degrees). `download` and `to_texture` reassemble
45/// tiled images.
46///
47/// Buffers come from FlashRaw's buffer pool and are **returned to the pool when the
48/// `GpuImage` is dropped** (large allocations are expensive to recreate). Keep the
49/// `GpuImage` alive while the GPU may still use them, or take ownership with
50/// [`GpuImage::into_buffer`] / [`GpuImage::into_tiles`]. A buffer may be larger than
51/// the data it holds.
52pub struct GpuImage {
53    pub(crate) ctx: Arc<GpuContext>,
54    /// Untiled: one part covering the image. Empty only after `into_tiles`.
55    pub(crate) parts: Vec<GpuTile>,
56    pub(crate) tiled: bool,
57    pub(crate) width: u32,
58    pub(crate) height: u32,
59    pub(crate) channels: u32,
60    pub(crate) format: GpuPixelFormat,
61}
62
63impl std::fmt::Debug for GpuImage {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct("GpuImage")
66            .field("width", &self.width)
67            .field("height", &self.height)
68            .field("channels", &self.channels)
69            .field("format", &self.format)
70            .field("tiles", &if self.tiled { self.parts.len() } else { 0 })
71            .field("adapter", &self.ctx.info.name)
72            .finish()
73    }
74}
75
76/// Identity semantics: two `GpuImage`s are equal when they are the same buffers.
77impl PartialEq for GpuImage {
78    fn eq(&self, other: &Self) -> bool {
79        self.parts.len() == other.parts.len() && self.parts.iter().zip(&other.parts).all(|(a, b)| a.buffer == b.buffer)
80    }
81}
82
83impl Drop for GpuImage {
84    fn drop(&mut self) {
85        for t in self.parts.drain(..) {
86            self.ctx.release(t.buffer);
87        }
88    }
89}
90
91/// Copy `src` into `dst` using several threads for large buffers (native): the
92/// destination is usually fresh memory, and page faults dominate a serial copy.
93fn par_copy(dst: &mut [u8], src: &[u8]) {
94    #[cfg(not(target_arch = "wasm32"))]
95    {
96        const CHUNK: usize = 8 << 20;
97        let threads = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1).min(8);
98        if src.len() >= 2 * CHUNK && threads > 1 {
99            let per = src.len().div_ceil(threads).div_ceil(4096) * 4096;
100            std::thread::scope(|s| {
101                for (d, c) in dst.chunks_mut(per).zip(src.chunks(per)) {
102                    s.spawn(move || d.copy_from_slice(c));
103                }
104            });
105            return;
106        }
107    }
108    dst.copy_from_slice(src);
109}
110
111/// [`par_copy`] of 4-sample 16-bit pixels keeping the first 3 samples of each
112/// (`src`: 8 bytes per pixel, `dst`: 6).
113fn par_copy_rgb16(dst: &mut [u8], src: &[u8]) {
114    let run = |d: &mut [u8], s: &[u8]| {
115        for (d, s) in d.chunks_exact_mut(6).zip(s.chunks_exact(8)) {
116            d.copy_from_slice(&s[..6]);
117        }
118    };
119    #[cfg(not(target_arch = "wasm32"))]
120    {
121        const CHUNK: usize = 8 << 20;
122        let threads = std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1).min(8);
123        if dst.len() >= 2 * CHUNK && threads > 1 {
124            let px = (dst.len() / 6).div_ceil(threads).div_ceil(4096) * 4096;
125            std::thread::scope(|sc| {
126                for (d, s) in dst.chunks_mut(px * 6).zip(src.chunks(px * 8)) {
127                    sc.spawn(move || run(d, s));
128                }
129            });
130            return;
131        }
132    }
133    run(dst, src);
134}
135
136/// Host copy of a downloaded [`GpuImage`].
137#[derive(Debug, Clone, PartialEq)]
138pub enum HostPixels {
139    U8(Vec<u8>),
140    U16(Vec<u16>),
141    F32(Vec<f32>),
142}
143
144impl GpuImage {
145    pub fn width(&self) -> u32 {
146        self.width
147    }
148    pub fn height(&self) -> u32 {
149        self.height
150    }
151    /// Samples per pixel (1 for CFA/mono, 3 for RGB, 4 for RGBA8).
152    pub fn channels(&self) -> u32 {
153        self.channels
154    }
155    pub fn format(&self) -> GpuPixelFormat {
156        self.format
157    }
158    /// Number of samples (`width * height * channels`; bytes for RGBA8).
159    pub fn len(&self) -> usize {
160        self.width as usize * self.height as usize * self.channels as usize
161    }
162    pub fn is_empty(&self) -> bool {
163        self.len() == 0
164    }
165    /// Bytes of image data (buffers may be larger).
166    pub fn byte_len(&self) -> usize {
167        self.len() * self.format.bytes_per_sample()
168    }
169    /// The storage buffer holding the whole image (row-major, from byte 0), or `None`
170    /// for a tiled image (see [`GpuImage::tiles`]). See the type docs for lifetimes.
171    pub fn buffer(&self) -> Option<&wgpu::Buffer> {
172        if self.tiled { None } else { self.parts.first().map(|t| &t.buffer) }
173    }
174    /// The image is split into several buffers (it is larger than one storage binding).
175    pub fn is_tiled(&self) -> bool {
176        self.tiled
177    }
178    /// The buffers and the rectangle each one holds (one tile covering the image when
179    /// untiled). Tiles are disjoint, cover the image, and are ordered by position.
180    pub fn tiles(&self) -> &[GpuTile] {
181        &self.parts
182    }
183    /// Take ownership of the buffer of an untiled image (it will not be recycled).
184    /// A tiled image is returned unchanged as the error.
185    pub fn into_buffer(mut self) -> std::result::Result<wgpu::Buffer, GpuImage> {
186        if self.tiled || self.parts.is_empty() {
187            return Err(self);
188        }
189        Ok(self.parts.remove(0).buffer)
190    }
191    /// Take ownership of all buffers (they will not be recycled).
192    pub fn into_tiles(mut self) -> Vec<GpuTile> {
193        std::mem::take(&mut self.parts)
194    }
195    pub fn device(&self) -> &wgpu::Device {
196        &self.ctx.device
197    }
198    pub fn queue(&self) -> &wgpu::Queue {
199        &self.ctx.queue
200    }
201    pub fn context(&self) -> &Arc<GpuContext> {
202        &self.ctx
203    }
204
205    /// Download to host memory (tiled images are reassembled). Returns the pixels and
206    /// the wall time in ms.
207    pub async fn download_async(&self) -> Result<(HostPixels, f64)> {
208        self.download_impl(false).await
209    }
210
211    /// Download a 4-channel 16-bit image (the 16-bit develop outputs) as 3-channel RGB:
212    /// the fourth sample of every pixel is dropped while copying out of the device
213    /// buffer, so no 4-channel host copy is made. Other images download unchanged.
214    pub async fn download_rgb_async(&self) -> Result<(HostPixels, f64)> {
215        self.download_impl(self.format == GpuPixelFormat::U16 && self.channels == 4).await
216    }
217
218    /// Blocking [`GpuImage::download_rgb_async`] (native only).
219    #[cfg(not(target_arch = "wasm32"))]
220    pub fn download_rgb(&self) -> Result<(HostPixels, f64)> {
221        pollster::block_on(self.download_rgb_async())
222    }
223
224    async fn download_impl(&self, drop_alpha: bool) -> Result<(HostPixels, f64)> {
225        let t0 = flashraw_core::profile::now_ms();
226        let n = if drop_alpha { self.len() / 4 * 3 } else { self.len() };
227        let mut px = match self.format {
228            GpuPixelFormat::Rgba8 => HostPixels::U8(vec![0u8; n]),
229            GpuPixelFormat::U16 => HostPixels::U16(vec![0u16; n]),
230            GpuPixelFormat::F32 => HostPixels::F32(vec![0f32; n]),
231        };
232        let out: &mut [u8] = match &mut px {
233            HostPixels::U8(v) => v,
234            HostPixels::U16(v) => bytemuck::cast_slice_mut(v),
235            HostPixels::F32(v) => bytemuck::cast_slice_mut(v),
236        };
237        let ctx = &self.ctx;
238        let bpp = self.channels as usize * self.format.bytes_per_sample();
239        // Bytes per pixel on the host (6 when the fourth 16-bit sample is dropped).
240        let hbpp = if drop_alpha { 6 } else { bpp };
241        let row = self.width as usize * hbpp;
242        // Untiled: the image starts at byte 0 of the single buffer, row-major.
243        for tile in &self.parts {
244            let r = tile.rect;
245            let tile_row = r.width as usize * bpp;
246            let bytes = tile_row * r.height as usize;
247            if bytes == 0 {
248                continue;
249            }
250            let size = (bytes as u64).div_ceil(4) * 4;
251            let host_tile_row = r.width as usize * hbpp;
252            let copy_out = |src: &[u8]| {
253                if host_tile_row == row {
254                    let o = r.y as usize * row;
255                    let len = host_tile_row * r.height as usize;
256                    if drop_alpha {
257                        par_copy_rgb16(&mut out[o..o + len], &src[..bytes]);
258                    } else {
259                        par_copy(&mut out[o..o + len], &src[..bytes]);
260                    }
261                } else {
262                    for (y, s) in src[..bytes].chunks_exact(tile_row).enumerate() {
263                        let o = (r.y as usize + y) * row + r.x as usize * hbpp;
264                        let d = &mut out[o..o + host_tile_row];
265                        if drop_alpha {
266                            for (d, s) in d.chunks_exact_mut(6).zip(s.chunks_exact(8)) {
267                                d.copy_from_slice(&s[..6]);
268                            }
269                        } else {
270                            d.copy_from_slice(s);
271                        }
272                    }
273                }
274            };
275            let buf = &tile.buffer;
276            if buf.usage().contains(wgpu::BufferUsages::MAP_READ) {
277                // Unified memory: read the storage buffer in place.
278                ctx.read_mapped(buf, size, copy_out).await?;
279            } else {
280                let staging =
281                    ctx.acquire("flashraw.readback", size, wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST);
282                let mut enc =
283                    ctx.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("readback") });
284                enc.copy_buffer_to_buffer(buf, 0, &staging, 0, size);
285                ctx.queue.submit([enc.finish()]);
286                let r = ctx.read_mapped(&staging, size, copy_out).await;
287                ctx.release(staging);
288                r?;
289            }
290        }
291        if let Some(e) = ctx.take_error() {
292            return Err(Error::GpuUnavailable(format!("GPU error during download: {e}")));
293        }
294        Ok((px, flashraw_core::profile::now_ms() - t0))
295    }
296
297    /// Blocking download (native only).
298    #[cfg(not(target_arch = "wasm32"))]
299    pub fn download(&self) -> Result<(HostPixels, f64)> {
300        pollster::block_on(self.download_async())
301    }
302
303    /// Copy an RGBA8 image (tiled or not) into a new `Rgba8Unorm` texture (usage
304    /// `STORAGE_BINDING | TEXTURE_BINDING | COPY_SRC`) for sampling/rendering.
305    /// The copy is queued; it completes before any later submission on the queue.
306    pub fn to_texture(&self) -> Result<wgpu::Texture> {
307        if self.format != GpuPixelFormat::Rgba8 {
308            return Err(Error::InvalidArgument("`to_texture` requires an RGBA8 `GpuImage`".into()));
309        }
310        let ctx = &self.ctx;
311        let max = ctx.device.limits().max_texture_dimension_2d;
312        if self.width > max || self.height > max {
313            return Err(Error::LimitExceeded(format!(
314                "{}x{} exceeds the adapter's {max} px texture limit",
315                self.width, self.height
316            )));
317        }
318        let tex = ctx.device.create_texture(&wgpu::TextureDescriptor {
319            label: Some("flashraw.rgba8"),
320            size: wgpu::Extent3d { width: self.width, height: self.height, depth_or_array_layers: 1 },
321            mip_level_count: 1,
322            sample_count: 1,
323            dimension: wgpu::TextureDimension::D2,
324            format: wgpu::TextureFormat::Rgba8Unorm,
325            usage: wgpu::TextureUsages::STORAGE_BINDING
326                | wgpu::TextureUsages::TEXTURE_BINDING
327                | wgpu::TextureUsages::COPY_SRC,
328            view_formats: &[],
329        });
330        let view = tex.create_view(&wgpu::TextureViewDescriptor::default());
331        let mut enc = ctx.device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("to_texture") });
332        for tile in &self.parts {
333            let r = tile.rect;
334            if r.is_empty() {
335                continue;
336            }
337            let params: [u32; 4] = [r.width, r.height, r.x, r.y];
338            let ubuf = crate::run::buffer_init(
339                &ctx.device,
340                "flashraw.tex_params",
341                bytemuck::cast_slice(&params),
342                wgpu::BufferUsages::UNIFORM,
343            );
344            let bytes = r.width as u64 * r.height as u64 * 4;
345            let bg = ctx.device.create_bind_group(&wgpu::BindGroupDescriptor {
346                label: Some("flashraw.to_texture"),
347                layout: &ctx.tex_layout,
348                entries: &[
349                    wgpu::BindGroupEntry { binding: 0, resource: ubuf.as_entire_binding() },
350                    wgpu::BindGroupEntry {
351                        binding: 1,
352                        resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
353                            buffer: &tile.buffer,
354                            offset: 0,
355                            size: wgpu::BufferSize::new(bytes),
356                        }),
357                    },
358                    wgpu::BindGroupEntry { binding: 2, resource: wgpu::BindingResource::TextureView(&view) },
359                ],
360            });
361            let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
362            pass.set_pipeline(ctx.tex_pipeline());
363            pass.set_bind_group(0, &bg, &[]);
364            pass.dispatch_workgroups(r.width.div_ceil(16), r.height.div_ceil(16), 1);
365        }
366        ctx.queue.submit([enc.finish()]);
367        if let Some(e) = ctx.take_error() {
368            return Err(Error::GpuUnavailable(format!("GPU error in `to_texture`: {e}")));
369        }
370        Ok(tex)
371    }
372}
flashraw_gpu/image.rs

flashraw_gpu/
image.rs
