flashraw_gpu/
kernels.rs

1//! Kernel sources, the uniform parameter block, and host-side table builders.
2//!
3//! The tables reproduce the CPU reference (`flashraw::process`) exactly: the same
4//! calibration pattern size, the same per-position black/scale values, the same
5//! 3x3 same-colour neighbour sets and the same sRGB quantisation table.
6
7use bytemuck::{Pod, Zeroable};
8use flashraw_core::{FrameInfo, PackOrder};
9
10/// `common.wgsl` + kernel file (WGSL has no `#include`).
11pub(crate) const UNPACK_WGSL: &str =
12    concat!(include_str!("../shaders/common.wgsl"), "\n", include_str!("../shaders/unpack.wgsl"));
13pub(crate) const NORMALIZE_WGSL: &str =
14    concat!(include_str!("../shaders/common.wgsl"), "\n", include_str!("../shaders/normalize.wgsl"));
15pub(crate) const DEVELOP_WGSL: &str =
16    concat!(include_str!("../shaders/common.wgsl"), "\n", include_str!("../shaders/develop.wgsl"));
17pub(crate) const TO_TEXTURE_WGSL: &str = include_str!("../shaders/to_texture.wgsl");
18
19/// WGSL sources of every kernel module, for embedding in other runtimes (e.g. a
20/// hand-written WebGPU pipeline in JavaScript). Each entry is `(name, source)`.
21pub fn wgsl_sources() -> [(&'static str, &'static str); 4] {
22    [("unpack", UNPACK_WGSL), ("normalize", NORMALIZE_WGSL), ("develop", DEVELOP_WGSL), ("to_texture", TO_TEXTURE_WGSL)]
23}
24
25/// Mirrors `struct Params` in `shaders/common.wgsl` (std140-compatible: 40 u32 then 5 vec4).
26#[repr(C)]
27#[derive(Debug, Clone, Copy, Pod, Zeroable, Default)]
28pub(crate) struct Params {
29    pub width: u32,
30    pub height: u32,
31    pub cpp: u32,
32    pub bits: u32,
33    pub src_stride: u32,
34    pub src_offset: u32,
35    pub src_len: u32,
36    pub row_bytes: u32,
37    pub src_x: u32,
38    pub src_y: u32,
39    pub origin_x: u32,
40    pub origin_y: u32,
41    pub pat_w: u32,
42    pub pat_h: u32,
43    pub cfa_w: u32,
44    pub cfa_h: u32,
45    pub n_rowd: u32,
46    pub n_cold: u32,
47    pub off_scale: u32,
48    pub off_rowd: u32,
49    pub off_cold: u32,
50    pub lut_len: u32,
51    pub off_srgb: u32,
52    pub off_cfa: u32,
53    pub crop_x: u32,
54    pub crop_y: u32,
55    pub dev_w: u32,
56    pub dev_h: u32,
57    pub out_w: u32,
58    pub out_h: u32,
59    pub orientation: u32,
60    pub total: u32,
61    pub band_y: u32,
62    pub band_end: u32,
63    pub dst_x: u32,
64    pub dst_y: u32,
65    pub dst_stride: u32,
66    pub dst_off: u32,
67    pub off_srgb16: u32,
68    /// Segments of the 8-bit output table at `off_srgb` (4096 for sRGB).
69    pub srgb_n: u32,
70    pub wb: [f32; 4],
71    pub m0: [f32; 4],
72    pub m1: [f32; 4],
73    pub m2: [f32; 4],
74    pub levels: [f32; 4],
75}
76
77pub(crate) fn order_index(o: PackOrder) -> u32 {
78    match o {
79        PackOrder::MsbBytes => 0,
80        PackOrder::LsbBytes => 1,
81        PackOrder::Msb16 => 2,
82        PackOrder::Msb32 => 3,
83    }
84}
85
86fn gcd(a: usize, b: usize) -> usize {
87    if b == 0 { a } else { gcd(b, a % b) }
88}
89fn lcm(a: usize, b: usize) -> usize {
90    a / gcd(a, b) * b
91}
92
93/// Calibration table (`tab` binding) and the matching `Params` fields.
94pub(crate) struct CalibTable {
95    pub tab: Vec<f32>,
96    pub pat_w: u32,
97    pub pat_h: u32,
98    pub off_scale: u32,
99    pub off_rowd: u32,
100    pub off_cold: u32,
101    pub n_rowd: u32,
102    pub n_cold: u32,
103}
104
105/// Per-position black and `1 / (white - black)` for each component, laid out as
106/// `black[c][pat]`, `scale[c][pat]`, then row and column deltas. Same pattern
107/// size and values as the CPU `Calib::new`.
108pub(crate) fn calib_table(info: &FrameInfo, cpp: usize) -> CalibTable {
109    let (cw, ch) = info.cfa.map(|p| (p.width(), p.height())).unwrap_or((1, 1));
110    let b = &info.black_level;
111    let (rows, cols, samples) = (b.rows.max(1) as usize, b.cols.max(1) as usize, b.samples.max(1) as usize);
112    let pw = lcm(cw, cols).min(64);
113    let ph = lcm(ch, rows).min(64);
114    let cpp = cpp.max(1);
115    let mut black = Vec::with_capacity(pw * ph * cpp);
116    let mut scale = Vec::with_capacity(pw * ph * cpp);
117    for c in 0..cpp {
118        for y in 0..ph {
119            for x in 0..pw {
120                let bl = b.values.get(((y % rows) * cols + x % cols) * samples + c % samples).copied().unwrap_or(0.0);
121                let color = info.cfa.map(|p| p.color_at(x, y) as usize).unwrap_or(c);
122                let wl = &info.white_level;
123                let white = wl.at(if wl.values.len() > 1 { color.min(wl.values.len() - 1) } else { 0 }) as f32;
124                black.push(bl);
125                scale.push(if white > bl { 1.0 / (white - bl) } else { 1.0 });
126            }
127        }
128    }
129    let off_scale = black.len();
130    let mut tab = black;
131    tab.extend_from_slice(&scale);
132    let off_rowd = tab.len();
133    tab.extend_from_slice(&b.row_deltas);
134    let off_cold = tab.len();
135    tab.extend_from_slice(&b.col_deltas);
136    // Reciprocals 1/n (n = 0..=9) at off_cold + n_cold, computed like the CPU's `1.0 / len as f32`.
137    tab.extend((0..=9).map(|n| if n == 0 { 0.0 } else { 1.0 / n as f32 }));
138    CalibTable {
139        tab,
140        pat_w: pw as u32,
141        pat_h: ph as u32,
142        off_scale: off_scale as u32,
143        off_rowd: off_rowd as u32,
144        off_cold: off_cold as u32,
145        n_rowd: b.row_deltas.len() as u32,
146        n_cold: b.col_deltas.len() as u32,
147    }
148}
149
150/// CFA site descriptors (see `cfa_site` in common.wgsl): colour (clamped to 3) and,
151/// for R/G/B, the mask of 3x3 neighbours of that colour. 1x1 green site without CFA.
152pub(crate) fn cfa_sites(info: &FrameInfo) -> (u32, u32, Vec<u32>) {
153    let Some(p) = info.cfa else { return (1, 1, vec![1]) };
154    let (pw, ph) = (p.width(), p.height());
155    let mut out = Vec::with_capacity(pw * ph);
156    for py in 0..ph {
157        for px in 0..pw {
158            let own = (p.color_at(px, py) as u32).min(3);
159            let mut d = own;
160            for dy in -1isize..=1 {
161                for dx in -1isize..=1 {
162                    let c = (p.color_at(
163                        (px as isize + dx + pw as isize * 4) as usize,
164                        (py as isize + dy + ph as isize * 4) as usize,
165                    ) as u32)
166                        .min(3);
167                    // A site's own colour is measured (centre only), not interpolated.
168                    if c < 3 && (c != own || (dx == 0 && dy == 0)) {
169                        let bit = ((dy + 1) * 3 + (dx + 1)) as u32;
170                        d |= 1 << (2 + 9 * c + bit);
171                    }
172                }
173            }
174            out.push(d);
175        }
176    }
177    (pw as u32, ph as u32, out)
178}
179
180/// Linear → sRGB 8-bit table with 4097 entries over [0, 1] (identical to the CPU LUT).
181pub(crate) fn srgb_table() -> Vec<u32> {
182    (0..=4096)
183        .map(|i| {
184            let x = i as f32 / 4096.0;
185            let s = if x <= 0.003_130_8 { 12.92 * x } else { 1.055 * x.powf(1.0 / 2.4) - 0.055 };
186            (s * 255.0 + 0.5).clamp(0.0, 255.0) as u8 as u32
187        })
188        .collect()
189}
190
191/// Linear -> sRGB curve scaled to 16 bits at 16385 evenly spaced points of `[0, 1]`
192/// (exact in f64, stored as f32; same table as the CPU `flashraw::process::srgb16_lut`).
193pub fn srgb16_table() -> &'static [f32] {
194    static T: std::sync::OnceLock<Vec<f32>> = std::sync::OnceLock::new();
195    T.get_or_init(|| {
196        (0..=16384)
197            .map(|i| {
198                let x = i as f64 / 16384.0;
199                let s = if x <= 0.003_130_8 { 12.92 * x } else { 1.055 * x.powf(1.0 / 2.4) - 0.055 };
200                (s * 65535.0) as f32
201            })
202            .collect()
203    })
204}
205
206/// Scalar models of the GPU kernels' addressing, written to mirror the WGSL line by
207/// line. They let the bit-level semantics be tested on machines without a GPU
208/// (against `flashraw_codecs::packed::unpack_row_scalar`), and serve as executable
209/// documentation of the shader behaviour.
210pub mod reference {
211    use flashraw_core::PackOrder;
212
213    /// Description of a packed raster as the `unpack` kernel sees it.
214    #[derive(Debug, Clone, Copy)]
215    pub struct PackedRaster {
216        pub bits: u32,
217        pub order: PackOrder,
218        /// Bytes from one row start to the next.
219        pub row_stride: usize,
220        /// Bytes of sample data in one row (bytes past this read as 0).
221        pub row_bytes: usize,
222    }
223
224    /// Sample `i` of row `y` exactly as `packed_sample` in common.wgsl computes it.
225    pub fn packed_sample(data: &[u8], r: &PackedRaster, i: usize, y: usize) -> u16 {
226        let row = y * r.row_stride;
227        let bitpos = i * r.bits as usize;
228        let (k, sh) = (bitpos >> 3, (bitpos & 7) as u32);
229        let byte = |k: usize| -> u32 {
230            let p = match r.order {
231                PackOrder::Msb16 => k ^ 1,
232                PackOrder::Msb32 => k ^ 3,
233                _ => k,
234            };
235            if p >= r.row_bytes {
236                return 0;
237            }
238            data.get(row + p).copied().unwrap_or(0) as u32
239        };
240        let (b0, b1, b2) = (byte(k), byte(k + 1), byte(k + 2));
241        let mask = (1u32 << r.bits) - 1;
242        let v = if r.order == PackOrder::LsbBytes {
243            ((b0 | (b1 << 8) | (b2 << 16)) >> sh) & mask
244        } else {
245            (((b0 << 16) | (b1 << 8) | b2) >> (24 - sh - r.bits)) & mask
246        };
247        v as u16
248    }
249
250    /// Whole-raster unpack (what the `unpack` kernel produces for a full frame).
251    pub fn unpack(data: &[u8], r: &PackedRaster, width: usize, height: usize) -> Vec<u16> {
252        let mut out = Vec::with_capacity(width * height);
253        for y in 0..height {
254            for x in 0..width {
255                out.push(packed_sample(data, r, x, y));
256            }
257        }
258        out
259    }
260}
261
262#[cfg(test)]
263mod tests {
264    use super::*;
265
266    #[test]
267    fn params_layout_matches_wgsl() {
268        // 40 scalars + 5 vec4<f32>: 240 bytes, 16-byte multiple (uniform buffer rules).
269        assert_eq!(std::mem::size_of::<Params>(), 240);
270        assert_eq!(std::mem::offset_of!(Params, wb), 160);
271        // Every WGSL field name appears in the host struct in the same order.
272        let src = include_str!("../shaders/common.wgsl");
273        let start = src.find("struct Params {").unwrap();
274        let body = &src[start..start + src[start..].find('}').unwrap()];
275        let names: Vec<&str> = body
276            .lines()
277            .skip(1)
278            .filter_map(|l| l.trim().split(':').next())
279            .filter(|n| !n.is_empty() && !n.starts_with("//"))
280            .collect();
281        let expected = [
282            "width",
283            "height",
284            "cpp",
285            "bits",
286            "src_stride",
287            "src_offset",
288            "src_len",
289            "row_bytes",
290            "src_x",
291            "src_y",
292            "origin_x",
293            "origin_y",
294            "pat_w",
295            "pat_h",
296            "cfa_w",
297            "cfa_h",
298            "n_rowd",
299            "n_cold",
300            "off_scale",
301            "off_rowd",
302            "off_cold",
303            "lut_len",
304            "off_srgb",
305            "off_cfa",
306            "crop_x",
307            "crop_y",
308            "dev_w",
309            "dev_h",
310            "out_w",
311            "out_h",
312            "orientation",
313            "total",
314            "band_y",
315            "band_end",
316            "dst_x",
317            "dst_y",
318            "dst_stride",
319            "dst_off",
320            "off_srgb16",
321            "srgb_n",
322            "wb",
323            "m0",
324            "m1",
325            "m2",
326            "levels",
327        ];
328        assert_eq!(names, expected);
329    }
330
331    #[test]
332    fn wgsl_modules_validate() {
333        for (name, src) in wgsl_sources() {
334            let module =
335                naga::front::wgsl::parse_str(src).unwrap_or_else(|e| panic!("{name}: {}", e.emit_to_string(src)));
336            naga::valid::Validator::new(naga::valid::ValidationFlags::all(), naga::valid::Capabilities::empty())
337                .validate(&module)
338                .unwrap_or_else(|e| panic!("{name}: {e:?}"));
339        }
340    }
341
342    #[test]
343    fn bayer_neighbour_masks() {
344        let mut info = FrameInfo::new(4, 4, 12);
345        info.cfa = Some(flashraw_core::CfaPattern::RGGB);
346        let (w, h, sites) = cfa_sites(&info);
347        assert_eq!((w, h), (2, 2));
348        // R site: own colour R only at the centre (bit 4), G at the cross, B at the corners.
349        let r = sites[0];
350        assert_eq!(r & 3, 0);
351        assert_eq!((r >> 2) & 0x1ff, 1 << 4);
352        assert_eq!((r >> 11) & 0x1ff, (1 << 1) | (1 << 3) | (1 << 5) | (1 << 7));
353        assert_eq!((r >> 20) & 0x1ff, (1 << 0) | (1 << 2) | (1 << 6) | (1 << 8));
354        // G site (top right): G only at the centre, not the diagonals.
355        assert_eq!((sites[1] >> 11) & 0x1ff, 1 << 4);
356    }
357
358    #[test]
359    fn srgb_table_endpoints() {
360        let t = srgb_table();
361        assert_eq!(t.len(), 4097);
362        assert_eq!(t[0], 0);
363        assert_eq!(t[4096], 255);
364        assert!(t.windows(2).all(|w| w[0] <= w[1]));
365    }
366}
flashraw_gpu/kernels.rs

flashraw_gpu/
kernels.rs
