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

flashraw_gpu/
kernels.rs
