1use crate::context::{GpuContext, Kernel, PipelineKey};
4use crate::image::{GpuImage, GpuPixelFormat, GpuTile};
5use crate::kernels::{self, Params};
6use flashraw_core::{Error, FrameInfo, Orientation, PackOrder, PackedPayload, Rect, Result};
7use std::sync::Arc;
8
9#[derive(Debug, Clone, Copy)]
11pub enum GpuInput<'a> {
12 Packed {
14 data: &'a [u8],
15 bits: u32,
16 order: PackOrder,
17 row_stride: usize,
19 },
20 U16(&'a [u16]),
22 F32(&'a [f32]),
24}
25
26#[derive(Debug, Clone, Copy)]
28pub struct GpuFrame<'a> {
29 pub input: GpuInput<'a>,
30 pub raster_width: u32,
32 pub raster_height: u32,
33 pub cpp: u32,
34 pub window: Rect,
36 pub origin: (u32, u32),
38 pub linearization: Option<&'a [u16]>,
40}
41
42impl<'a> GpuFrame<'a> {
43 pub fn from_packed(p: &PackedPayload<'a>) -> Self {
44 GpuFrame {
45 input: GpuInput::Packed { data: p.data, bits: p.bits, order: p.order, row_stride: p.row_stride },
46 raster_width: p.width,
47 raster_height: p.height,
48 cpp: 1,
49 window: Rect::full(p.width, p.height),
50 origin: (0, 0),
51 linearization: None,
52 }
53 }
54 pub fn from_u16(data: &'a [u16], width: u32, height: u32, cpp: u32) -> Self {
55 Self::plain(GpuInput::U16(data), width, height, cpp)
56 }
57 pub fn from_f32(data: &'a [f32], width: u32, height: u32, cpp: u32) -> Self {
58 Self::plain(GpuInput::F32(data), width, height, cpp)
59 }
60 fn plain(input: GpuInput<'a>, width: u32, height: u32, cpp: u32) -> Self {
61 GpuFrame {
62 input,
63 raster_width: width,
64 raster_height: height,
65 cpp,
66 window: Rect::full(width, height),
67 origin: (0, 0),
68 linearization: None,
69 }
70 }
71 pub fn with_window(mut self, r: Rect) -> Self {
73 self.origin = (self.origin.0 + r.x - self.window.x, self.origin.1 + r.y - self.window.y);
74 self.window = r;
75 self
76 }
77 pub fn with_origin(mut self, origin: (u32, u32)) -> Self {
78 self.origin = origin;
79 self
80 }
81 pub fn with_linearization(mut self, lut: Option<&'a [u16]>) -> Self {
82 self.linearization = lut.filter(|l| !l.is_empty());
83 self
84 }
85 pub fn upload_bytes(&self) -> usize {
87 match self.input {
88 GpuInput::Packed { data, .. } => data.len(),
89 GpuInput::U16(d) => d.len() * 2,
90 GpuInput::F32(d) => d.len() * 4,
91 }
92 }
93}
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq)]
97pub enum DevelopOutput {
98 Rgba8,
100 LinearRgbF32,
102 Rgb16,
105 LinearRgb16,
108 LinearRgbF16,
111}
112
113#[derive(Debug, Clone, Copy)]
115pub struct DevelopSpec {
116 pub crop: Rect,
118 pub half_size: bool,
120 pub wb: [f32; 4],
122 pub cam_to_rgb: [[f32; 3]; 3],
123 pub orientation: Orientation,
124 pub output: DevelopOutput,
125 pub levels: Option<[f32; 2]>,
128 pub curve: Option<GpuCurve>,
130 pub edge: Option<EdgeSpec>,
133}
134
135#[derive(Debug, Clone, Copy)]
139pub struct EdgeSpec {
140 pub period: u32,
142 pub cx: u32,
144 pub cy: u32,
145 pub beta: f32,
147 pub halo: u32,
149 pub canon: &'static [u8],
151 pub g0: &'static [f32],
153 pub dir: &'static [f32],
155 pub ahd: bool,
158}
159
160const EDGE_SCRATCH_MAX: u64 = 512 << 20;
162const EDGE_PLANES: u64 = 8;
164const AHD_PLANES: u64 = 11;
166
167impl EdgeSpec {
168 fn planes(&self) -> u64 {
169 if self.ahd { AHD_PLANES } else { EDGE_PLANES }
170 }
171}
172
173struct EdgeRun {
175 params: wgpu::Buffer,
177 stride: u64,
178 tables: Arc<wgpu::Buffer>,
179 scratch: wgpu::Buffer,
180 bands: Vec<(u32, (u32, u32), u64)>,
182 passes: Vec<wgpu::ComputePipeline>,
184}
185
186#[derive(Debug, Clone, Copy)]
188pub struct GpuCurve {
189 pub lut8: &'static [u8],
191 pub lut16: &'static [f32],
193}
194
195#[derive(Debug, Clone, Copy)]
197pub enum GpuOp {
198 Unpack,
200 Normalize,
202 Develop(DevelopSpec),
203}
204
205#[derive(Debug, Clone, Copy, Default, PartialEq)]
207pub struct GpuTimings {
208 pub upload_ms: f64,
211 pub compute_ms: f64,
213 pub kernel_ms: Option<f64>,
215 pub bands: u32,
217}
218
219#[derive(Debug)]
221pub struct GpuRun {
222 pub image: GpuImage,
223 pub timings: GpuTimings,
224}
225
226pub(crate) fn buffer_init(device: &wgpu::Device, label: &str, data: &[u8], usage: wgpu::BufferUsages) -> wgpu::Buffer {
227 let size = (data.len() as u64).div_ceil(4).max(1) * 4;
228 let buf =
229 device.create_buffer(&wgpu::BufferDescriptor { label: Some(label), size, usage, mapped_at_creation: true });
230 {
231 let mut view = buf.get_mapped_range_mut(..).expect("buffer mapped at creation");
232 let mut padded = data.to_vec();
233 padded.resize(size as usize, 0);
234 view.copy_from_slice(&padded);
235 }
236 buf.unmap();
237 buf
238}
239
240fn write_padded(queue: &wgpu::Queue, buf: &wgpu::Buffer, data: &[u8]) {
243 let bulk = data.len() & !3;
244 if bulk > 0 {
245 queue.write_buffer(buf, 0, &data[..bulk]);
246 }
247 if bulk < data.len() {
248 let mut tail = [0u8; 4];
249 tail[..data.len() - bulk].copy_from_slice(&data[bulk..]);
250 queue.write_buffer(buf, bulk as u64, &tail);
251 }
252}
253
254fn grid_1d(n: u64) -> (u32, u32) {
256 let groups = n.div_ceil(256).max(1);
257 let x = groups.min(32768);
258 (x as u32, groups.div_ceil(x) as u32)
259}
260
261fn u32_checked(v: u64, what: &str) -> Result<u32> {
262 u32::try_from(v).map_err(|_| Error::LimitExceeded(format!("{what} too large for GPU addressing")))
263}
264
265fn round4(v: u64) -> u64 {
266 v.div_ceil(4).max(1) * 4
267}
268
269#[derive(Debug, Clone, Copy, PartialEq, Eq)]
271enum OutLayout {
272 Single,
274 SubRange,
276 Tiled,
278}
279
280#[derive(Debug, Clone, Copy, PartialEq, Eq)]
283pub(crate) struct Band {
284 pub y0: u32,
285 pub y1: u32,
286 pub in0: u32,
287 pub in1: u32,
288}
289
290#[derive(Debug, Clone, Copy)]
292pub(crate) struct BandGeometry {
293 pub rows: u32,
295 pub in_row_bytes: u64,
297 pub out_row_bytes: u64,
299 pub kind: BandKind,
301 pub frame_h: u32,
303 pub row_align: u32,
305 pub tile_rows: u32,
307 pub max_rows: u32,
308}
309
310#[derive(Debug, Clone, Copy, PartialEq, Eq)]
311pub(crate) enum BandKind {
312 Direct,
314 Bilinear { crop_y: u32 },
316 Half { crop_y: u32 },
318 Edge { crop_y: u32, halo: u32 },
320}
321
322impl BandGeometry {
323 fn input_rows(&self, y0: u32, y1: u32) -> (u32, u32) {
324 let h = self.frame_h;
325 match self.kind {
326 BandKind::Direct => (y0, y1.min(h)),
327 BandKind::Bilinear { crop_y } => ((crop_y + y0).saturating_sub(1), (crop_y + y1 + 1).min(h)),
328 BandKind::Half { crop_y } => ((crop_y + 2 * y0).min(h), (crop_y + 2 * y1).min(h)),
329 BandKind::Edge { crop_y, halo } => ((crop_y + y0).saturating_sub(halo), (crop_y + y1 + halo).min(h)),
330 }
331 }
332
333 pub fn plan(&self, budget: u64, in_slack: u64, out_whole: bool, out_slack: u64) -> Result<Vec<Band>> {
336 let rows_in = budget.saturating_sub(in_slack) / self.in_row_bytes.max(1);
337 let n_in = match self.kind {
338 BandKind::Direct => rows_in,
339 BandKind::Bilinear { .. } => rows_in.saturating_sub(2),
340 BandKind::Half { .. } => rows_in / 2,
341 BandKind::Edge { halo, .. } => rows_in.saturating_sub(2 * halo as u64),
342 };
343 let n_out = if out_whole { u64::MAX } else { budget.saturating_sub(out_slack) / self.out_row_bytes.max(1) };
344 let mut n = n_in.min(n_out).min(self.max_rows as u64).min(self.rows as u64) as u32;
345 if n < self.rows {
346 if n >= self.tile_rows {
347 n -= n % self.tile_rows;
348 }
349 n -= n % self.row_align;
350 }
351 if n == 0 {
352 return Err(Error::LimitExceeded(format!(
353 "one band of {} rows exceeds the GPU buffer limit ({:.1} MB)",
354 self.row_align,
355 budget as f64 / 1e6
356 )));
357 }
358 let mut bands = Vec::with_capacity(self.rows.div_ceil(n) as usize);
359 let mut y0 = 0;
360 while y0 < self.rows {
361 let y1 = (y0 + n).min(self.rows);
362 let (in0, in1) = self.input_rows(y0, y1);
363 bands.push(Band { y0, y1, in0, in1 });
364 y0 = y1;
365 }
366 Ok(bands)
367 }
368}
369
370pub(crate) fn band_rect(y0: u32, y1: u32, dw: u32, dh: u32, orientation: u32) -> Rect {
373 let n = y1 - y0;
374 match orientation {
375 3 | 4 => Rect::new(0, dh - y1, dw, n),
376 5 | 8 => Rect::new(y0, 0, n, dw),
377 6 | 7 => Rect::new(dh - y1, 0, n, dw),
378 _ => Rect::new(0, y0, dw, n),
379 }
380}
381
382struct InputSlice<'a> {
384 bytes: &'a [u8],
385 offset: u32,
386 len: u32,
387}
388
389struct Slot {
391 src: wgpu::Buffer,
392 staging: Option<wgpu::Buffer>,
393 last: Option<wgpu::SubmissionIndex>,
394}
395
396impl GpuContext {
397 pub async fn run(self: &Arc<Self>, frame: &GpuFrame<'_>, info: &FrameInfo, op: &GpuOp) -> Result<GpuRun> {
405 self.run_banded(frame, info, op, None).await
406 }
407
408 pub(crate) async fn run_banded(
411 self: &Arc<Self>,
412 frame: &GpuFrame<'_>,
413 info: &FrameInfo,
414 op: &GpuOp,
415 budget: Option<u64>,
416 ) -> Result<GpuRun> {
417 let (w, h, cpp) = (frame.window.width, frame.window.height, frame.cpp.max(1));
418 let win = frame.window;
419 if w == 0 || h == 0 || win.right() > frame.raster_width || win.bottom() > frame.raster_height {
420 return Err(Error::InvalidArgument(format!(
421 "window {win} outside {}x{} raster",
422 frame.raster_width, frame.raster_height
423 )));
424 }
425 let mut curve: Option<GpuCurve> = None;
426 let mut edge: Option<EdgeSpec> = None;
427 let mut p = Params {
428 srgb_n: 4096,
429 width: w,
430 height: h,
431 cpp,
432 src_x: win.x,
433 origin_x: frame.origin.0,
434 origin_y: frame.origin.1,
435 orientation: 1,
436 wb: [1.0; 4],
437 ..Default::default()
438 };
439 let raster_samples = frame.raster_width as u64 * frame.raster_height as u64 * cpp as u64;
441 let raster_row = frame.raster_width as u64 * cpp as u64;
442 let (input_mode, order, in_row_bytes): (u32, u32, u64) = match frame.input {
443 GpuInput::Packed { bits, order, row_stride, .. } => {
444 if !(1..=16).contains(&bits) {
445 return Err(Error::InvalidArgument(format!("unsupported packed bit depth {bits}")));
446 }
447 let row_bytes = (raster_row * bits as u64).div_ceil(8);
448 p.bits = bits;
449 p.src_stride = u32_checked(row_stride as u64, "row stride")?;
450 p.row_bytes = u32_checked(row_bytes, "row size")?;
451 (2, kernels::order_index(order), (row_stride as u64).max(row_bytes))
452 }
453 GpuInput::U16(d) => {
454 if (d.len() as u64) < raster_samples {
455 return Err(Error::InvalidArgument("u16 raster shorter than its dimensions".into()));
456 }
457 p.src_stride = u32_checked(raster_row, "row")?;
458 (0, 0, raster_row * 2)
459 }
460 GpuInput::F32(d) => {
461 if (d.len() as u64) < raster_samples {
462 return Err(Error::InvalidArgument("f32 raster shorter than its dimensions".into()));
463 }
464 p.src_stride = u32_checked(raster_row, "row")?;
465 (1, 0, raster_row * 4)
466 }
467 };
468 let use_lut = frame.linearization.is_some() && input_mode != 1;
469 let mut calib = kernels::calib_table(info, cpp as usize);
470 p.pat_w = calib.pat_w;
471 p.pat_h = calib.pat_h;
472 p.off_scale = calib.off_scale;
473 p.off_rowd = calib.off_rowd;
474 p.off_cold = calib.off_cold;
475 p.n_rowd = calib.n_rowd;
476 p.n_cold = calib.n_cold;
477 let (cfa_w, cfa_h, sites) = kernels::cfa_sites(info);
478 p.cfa_w = cfa_w;
479 p.cfa_h = cfa_h;
480
481 let row_samples = w as u64 * cpp as u64;
483 let mut geo = BandGeometry {
484 rows: h,
485 in_row_bytes,
486 out_row_bytes: 0,
487 kind: BandKind::Direct,
488 frame_h: h,
489 row_align: 1,
490 tile_rows: 1,
491 max_rows: u32::MAX,
492 };
493 let (key, out_w, out_h, channels, format) = match op {
494 GpuOp::Unpack => {
495 if input_mode == 1 {
496 return Err(Error::InvalidArgument("unpack needs integer input".into()));
497 }
498 geo.out_row_bytes = row_samples * 2;
499 geo.row_align = if row_samples % 2 == 1 { 2 } else { 1 };
501 let key = PipelineKey { kernel: Kernel::Unpack, input_mode, order, use_lut, layout: 0, out_mode: 0 };
502 (key, w, h, cpp, GpuPixelFormat::U16)
503 }
504 GpuOp::Normalize => {
505 geo.out_row_bytes = row_samples * 4;
506 let key = PipelineKey { kernel: Kernel::Normalize, input_mode, order, use_lut, layout: 0, out_mode: 0 };
507 (key, w, h, cpp, GpuPixelFormat::F32)
508 }
509 GpuOp::Develop(d) => {
510 let crop = d.crop.clamp_to(w, h);
511 let (dw, dh) = if d.half_size { (crop.width / 2, crop.height / 2) } else { (crop.width, crop.height) };
512 if dw == 0 || dh == 0 {
513 return Err(Error::InvalidArgument("develop region too small".into()));
514 }
515 p.crop_x = crop.x;
516 p.crop_y = crop.y;
517 p.dev_w = dw;
518 p.dev_h = dh;
519 p.wb = d.wb;
520 curve = d.curve;
521 let m = d.cam_to_rgb;
522 p.m0 = [m[0][0], m[0][1], m[0][2], 0.0];
523 p.m1 = [m[1][0], m[1][1], m[1][2], 0.0];
524 p.m2 = [m[2][0], m[2][1], m[2][2], 0.0];
525 if let Some([scale, offset]) = d.levels {
526 p.levels = [scale, offset, 1.0, 0.0];
527 }
528 edge = d.edge.filter(|_| !d.half_size && cpp == 1 && info.cfa.is_some());
529 let layout = if cpp == 1 && info.cfa.is_some_and(|c| c.is_bayer()) {
530 3
531 } else if cpp == 1 && info.cfa.is_some() {
532 1
533 } else if cpp >= 3 {
534 2
535 } else {
536 0
537 };
538 let (out_mode, ow, oh, ch, fmt) = match d.output {
539 DevelopOutput::Rgba8 => {
540 p.orientation = d.orientation.exif();
541 let (ow, oh) = if d.orientation.swaps_dimensions() { (dh, dw) } else { (dw, dh) };
542 (0, ow, oh, 4, GpuPixelFormat::Rgba8)
543 }
544 DevelopOutput::LinearRgbF32 => (1, dw, dh, 3, GpuPixelFormat::F32),
545 DevelopOutput::Rgb16 => {
546 p.orientation = d.orientation.exif();
547 let (ow, oh) = if d.orientation.swaps_dimensions() { (dh, dw) } else { (dw, dh) };
548 (2, ow, oh, 4, GpuPixelFormat::U16)
549 }
550 DevelopOutput::LinearRgb16 => (3, dw, dh, 4, GpuPixelFormat::U16),
551 DevelopOutput::LinearRgbF16 => (4, dw, dh, 4, GpuPixelFormat::U16),
552 };
553 p.out_w = ow;
554 p.out_h = oh;
555 geo.rows = dh;
556 geo.out_row_bytes = dw as u64 * ch as u64 * fmt.bytes_per_sample() as u64;
557 geo.kind = match (&edge, d.half_size) {
558 (Some(e), _) => BandKind::Edge { crop_y: crop.y, halo: e.halo },
559 (None, true) => BandKind::Half { crop_y: crop.y },
560 (None, false) => BandKind::Bilinear { crop_y: crop.y },
561 };
562 geo.tile_rows = 16;
563 geo.max_rows = 16 * self.device.limits().max_compute_workgroups_per_dimension;
564 if let Some(e) = &edge {
565 let table = e.period as usize * e.period as usize;
567 if !matches!(e.period, 2 | 6)
568 || e.canon.len() != table
569 || e.g0.len() != table * 49
570 || e.dir.len() != table * 3 * 4 * 49
571 || (e.ahd && (e.period != 2 || e.beta != 0.0))
572 {
573 return Err(Error::InvalidArgument("inconsistent edge-aware filter tables".into()));
574 }
575 let row_bytes = (dw as u64 + 2 * e.halo as u64) * e.planes() * 4;
576 let avail = EDGE_SCRATCH_MAX.min(self.max_binding_bytes());
577 let rows = (avail / row_bytes).saturating_sub(2 * e.halo as u64);
578 if rows == 0 {
579 return Err(Error::LimitExceeded(
580 "frame too wide for the edge-aware demosaic on this GPU".into(),
581 ));
582 }
583 geo.max_rows = geo.max_rows.min(rows.min(u32::MAX as u64) as u32);
584 }
585 let key = PipelineKey {
586 kernel: if edge.is_some() {
587 Kernel::EdgeDevelop
588 } else if d.half_size {
589 Kernel::DevelopHalf
590 } else {
591 Kernel::Develop
592 },
593 input_mode,
594 order,
595 use_lut,
596 layout,
597 out_mode,
598 };
599 (key, ow, oh, ch, fmt)
600 }
601 };
602 if key.out_mode == 2 {
603 p.off_srgb16 = calib.tab.len() as u32;
604 match curve {
605 Some(c) if c.lut16.len() == kernels::srgb16_table().len() => calib.tab.extend_from_slice(c.lut16),
606 _ => calib.tab.extend_from_slice(kernels::srgb16_table()),
607 }
608 }
609 let bps = format.bytes_per_sample() as u64;
610 let out_bytes = out_w as u64 * out_h as u64 * channels as u64 * bps;
611 let columns = p.orientation >= 5;
613
614 let limits = self.device.limits();
616 let max_buffer = self.max_buffer_bytes();
617 let binding = budget.unwrap_or_else(|| self.max_binding_bytes()).min(self.max_binding_bytes()) & !3;
618 let align = limits.min_storage_buffer_offset_alignment.max(4) as u64;
619 let layout = if round4(out_bytes) <= binding {
620 OutLayout::Single
621 } else if out_bytes <= max_buffer && !columns {
622 OutLayout::SubRange
623 } else {
624 OutLayout::Tiled
625 };
626 let bands = geo.plan(binding, 8, layout == OutLayout::Single, align + 4)?;
627
628 let mut lut: Vec<u32> = Vec::new();
630 if let Some(l) = frame.linearization.filter(|_| use_lut) {
631 lut.extend(l.iter().map(|&v| v as u32));
632 p.lut_len = l.len() as u32;
633 }
634 p.off_srgb = lut.len() as u32;
635 match curve {
636 Some(c) if key.out_mode == 0 && c.lut8.len() > 1 => {
637 p.srgb_n = (c.lut8.len() - 1) as u32;
638 lut.extend(c.lut8.iter().map(|&v| v as u32));
639 }
640 _ => lut.extend_from_slice(&self.srgb),
641 }
642 p.off_cfa = lut.len() as u32;
643 lut.extend_from_slice(&sites);
644
645 let slice_of = |b: &Band| -> InputSlice<'_> {
647 let (r0, r1) = ((win.y + b.in0) as u64, (win.y + b.in1) as u64);
648 match frame.input {
649 GpuInput::Packed { data, row_stride, .. } => {
650 let len = data.len() as u64;
651 let start = (r0 * row_stride as u64).min(len);
652 let end = if r1 > r0 { ((r1 - 1) * row_stride as u64 + in_row_bytes).min(len) } else { start };
653 let al = start & !3;
654 InputSlice {
655 bytes: &data[al as usize..end as usize],
656 offset: (start - al) as u32,
657 len: (end - start) as u32,
658 }
659 }
660 GpuInput::U16(d) => {
661 let s = &d[(r0 * raster_row) as usize..(r1 * raster_row) as usize];
662 InputSlice { bytes: bytemuck::cast_slice(s), offset: 0, len: 0 }
663 }
664 GpuInput::F32(d) => {
665 let s = &d[(r0 * raster_row) as usize..(r1 * raster_row) as usize];
666 InputSlice { bytes: bytemuck::cast_slice(s), offset: 0, len: 0 }
667 }
668 }
669 };
670 let slices: Vec<InputSlice<'_>> = bands.iter().map(slice_of).collect();
671 let in_max = slices.iter().map(|s| round4(s.bytes.len() as u64)).max().unwrap_or(4);
672 let develop = matches!(op, GpuOp::Develop(_));
673 let mut params = Vec::with_capacity(bands.len());
676 let mut out_ranges = Vec::with_capacity(bands.len());
677 for (b, s) in bands.iter().zip(&slices) {
678 let mut q = p;
679 q.src_y = b.in0;
680 q.src_offset = s.offset;
681 q.src_len = s.len;
682 q.band_y = b.y0;
683 q.band_end = b.y1;
684 q.total = u32_checked((b.y1 - b.y0) as u64 * row_samples, "band")?;
685 let rect = if develop {
686 band_rect(b.y0, b.y1, p.dev_w, p.dev_h, p.orientation)
687 } else {
688 Rect::new(0, b.y0, w, b.y1 - b.y0)
689 };
690 let row_out = out_w as u64 * channels as u64 * bps;
691 let (offset, size) = match layout {
692 OutLayout::Single => {
693 q.dst_stride = out_w;
694 q.dst_off = if develop { 0 } else { u32_checked(rect.y as u64 * row_out / 4, "output")? };
695 (0, round4(out_bytes))
696 }
697 OutLayout::SubRange => {
698 let start = rect.y as u64 * row_out;
699 let a = start / align * align;
700 let end = round4(rect.bottom() as u64 * row_out);
701 q.dst_y = rect.y;
702 q.dst_stride = out_w;
703 q.dst_off = ((start - a) / 4) as u32;
704 (a, end - a)
705 }
706 OutLayout::Tiled => {
707 q.dst_x = rect.x;
708 q.dst_y = rect.y;
709 q.dst_stride = rect.width;
710 (0, round4(rect.width as u64 * rect.height as u64 * channels as u64 * bps))
711 }
712 };
713 params.push(q);
714 out_ranges.push((rect, offset, size));
715 }
716
717 let dev = &self.device;
718 let validation = dev.push_error_scope(wgpu::ErrorFilter::Validation);
719 let oom = dev.push_error_scope(wgpu::ErrorFilter::OutOfMemory);
720
721 let ustride = (std::mem::size_of::<Params>() as u64)
723 .div_ceil(limits.min_uniform_buffer_offset_alignment as u64)
724 * limits.min_uniform_buffer_offset_alignment as u64;
725 let mut ubytes = vec![0u8; ustride as usize * params.len()];
726 for (i, q) in params.iter().enumerate() {
727 let o = i * ustride as usize;
728 ubytes[o..o + std::mem::size_of::<Params>()].copy_from_slice(bytemuck::bytes_of(q));
729 }
730 let ubuf = buffer_init(dev, "flashraw.params", &ubytes, wgpu::BufferUsages::UNIFORM);
731 let tab = buffer_init(dev, "flashraw.calib", bytemuck::cast_slice(&calib.tab), wgpu::BufferUsages::STORAGE);
732 let lutb = buffer_init(dev, "flashraw.lut", bytemuck::cast_slice(&lut), wgpu::BufferUsages::STORAGE);
733 let mut dst_usage = wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC;
734 if self.uma {
735 dst_usage |= wgpu::BufferUsages::MAP_READ;
736 }
737 let outputs: Vec<wgpu::Buffer> = match layout {
738 OutLayout::Single | OutLayout::SubRange => {
739 vec![self.acquire("flashraw.output", round4(out_bytes), dst_usage)]
740 }
741 OutLayout::Tiled => {
742 out_ranges.iter().map(|&(_, _, size)| self.acquire("flashraw.output", size, dst_usage)).collect()
743 }
744 };
745 let staged = !self.uma && in_max >= 1 << 20;
748 let nslots = bands.len().min(2);
749 let mut slots: Vec<Slot> = (0..nslots)
750 .map(|_| {
751 let src_usage = if self.uma {
752 wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::MAP_WRITE
753 } else {
754 wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST
755 };
756 Slot {
757 src: self.acquire("flashraw.src", in_max, src_usage),
758 staging: staged.then(|| {
759 self.acquire(
760 "flashraw.upload",
761 in_max,
762 wgpu::BufferUsages::MAP_WRITE | wgpu::BufferUsages::COPY_SRC,
763 )
764 }),
765 last: None,
766 }
767 })
768 .collect();
769 let ts = (self.info.timestamps && bands.len() * 2 <= wgpu::QUERY_SET_MAX_QUERIES as usize).then(|| {
770 let n = bands.len() as u32 * 2;
771 let qs = dev.create_query_set(&wgpu::QuerySetDescriptor {
772 label: Some("flashraw.timestamps"),
773 ty: wgpu::QueryType::Timestamp,
774 count: n,
775 });
776 let resolve = dev.create_buffer(&wgpu::BufferDescriptor {
777 label: Some("flashraw.ts_resolve"),
778 size: n as u64 * 8,
779 usage: wgpu::BufferUsages::QUERY_RESOLVE | wgpu::BufferUsages::COPY_SRC,
780 mapped_at_creation: false,
781 });
782 let read = dev.create_buffer(&wgpu::BufferDescriptor {
783 label: Some("flashraw.ts_read"),
784 size: n as u64 * 8,
785 usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
786 mapped_at_creation: false,
787 });
788 (qs, resolve, read)
789 });
790 let pipeline = self.pipeline(key);
791 let edge_run = match &edge {
792 Some(e) => {
793 let mut eps = Vec::with_capacity(bands.len());
794 let mut info_b = Vec::with_capacity(bands.len());
795 let table = e.period * e.period;
796 for b in &bands {
797 let (hl, fw, fh) = (e.halo, p.width, p.height);
800 let (ix0, ix1) = (p.crop_x.max(hl), (p.crop_x + p.dev_w).min(fw.saturating_sub(hl)));
801 let (iy0, iy1) = ((p.crop_y + b.y0).max(hl), (p.crop_y + b.y1).min(fh.saturating_sub(hl)));
802 let (ex0, ey0, ew, eh) = if ix1 > ix0 && iy1 > iy0 {
803 (ix0 - hl, iy0 - hl, ix1 - ix0 + 2 * hl, iy1 - iy0 + 2 * hl)
804 } else {
805 (0, 0, 0, 0)
806 };
807 let n = u32_checked(ew as u64 * eh as u64, "edge-aware band")?;
808 eps.push(kernels::EdgeParams {
809 period: e.period,
810 cx: e.cx,
811 cy: e.cy,
812 off_g0: table,
813 off_dir: table + table * 49,
814 ex0,
815 ey0,
816 ew,
817 eh,
818 n,
819 halo: hl,
820 smooth_on: (e.beta > 0.0) as u32,
821 beta: e.beta,
822 pad: [0.0; 3],
823 wb: match op {
824 GpuOp::Develop(d) => d.wb,
825 _ => [1.0; 4],
826 },
827 });
828 info_b.push((n, (ew.div_ceil(16), eh.div_ceil(16)), round4(n as u64 * e.planes() * 4).max(16)));
829 }
830 let stride = (std::mem::size_of::<kernels::EdgeParams>() as u64)
831 .div_ceil(limits.min_uniform_buffer_offset_alignment as u64)
832 * limits.min_uniform_buffer_offset_alignment as u64;
833 let mut bytes = vec![0u8; stride as usize * eps.len()];
834 for (i, q) in eps.iter().enumerate() {
835 let o = i * stride as usize;
836 bytes[o..o + std::mem::size_of::<kernels::EdgeParams>()].copy_from_slice(bytemuck::bytes_of(q));
837 }
838 let tables = self.edge_table(e.period, || {
839 let mut v: Vec<f32> = e.canon.iter().map(|&c| c as f32).collect();
840 v.extend_from_slice(e.g0);
841 for pc in 0..table as usize * 3 {
844 let f = &e.dir[pc * 196..(pc + 1) * 196];
845 v.extend((0..49 * 4).map(|i| f[(i % 4) * 49 + i / 4]));
846 }
847 v
848 });
849 let scratch_bytes = info_b.iter().map(|b| b.2).max().unwrap_or(16);
850 let mk = |kernel| self.pipeline(PipelineKey { kernel, ..key });
851 let passes = if e.ahd {
852 [Kernel::EdgeMosaic, Kernel::AhdGreen, Kernel::AhdLab, Kernel::AhdHom, Kernel::AhdBlend]
853 .map(mk)
854 .to_vec()
855 } else {
856 let d1 = if e.period == 2 { Kernel::EdgeBayerD1 } else { Kernel::EdgeXtransD1 };
857 [Kernel::EdgeMosaic, Kernel::EdgeG0, Kernel::EdgeRgb, d1, Kernel::EdgeBayerV2].map(mk).to_vec()
858 };
859 Some(EdgeRun {
860 params: buffer_init(dev, "flashraw.edge_params", &bytes, wgpu::BufferUsages::UNIFORM),
861 stride,
862 tables,
863 scratch: self.acquire("flashraw.edge_scratch", scratch_bytes, wgpu::BufferUsages::STORAGE),
864 bands: info_b,
865 passes,
866 })
867 }
868 None => None,
869 };
870
871 let t0 = flashraw_core::profile::now_ms();
873 let mut upload_ms = 0.0;
874 for (i, (b, s)) in bands.iter().zip(&slices).enumerate() {
875 let slot = &mut slots[i % nslots];
876 let in_size = round4(s.bytes.len() as u64);
877 upload_ms += if self.uma {
880 self.write_mapped(&slot.src, s.bytes, slot.last.take()).await?
881 } else if let Some(st) = &slot.staging {
882 self.write_mapped(st, s.bytes, slot.last.take()).await?
883 } else {
884 let tu = flashraw_core::profile::now_ms();
885 write_padded(&self.queue, &slot.src, s.bytes);
886 flashraw_core::profile::now_ms() - tu
887 };
888 let (_, out_offset, out_size) = out_ranges[i];
889 let out_buf = if layout == OutLayout::Tiled { &outputs[i] } else { &outputs[0] };
890 let bg = dev.create_bind_group(&wgpu::BindGroupDescriptor {
891 label: Some("flashraw.kernel"),
892 layout: &self.layout,
893 entries: &[
894 wgpu::BindGroupEntry {
895 binding: 0,
896 resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
897 buffer: &ubuf,
898 offset: i as u64 * ustride,
899 size: wgpu::BufferSize::new(std::mem::size_of::<Params>() as u64),
900 }),
901 },
902 wgpu::BindGroupEntry {
903 binding: 1,
904 resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
905 buffer: &slot.src,
906 offset: 0,
907 size: wgpu::BufferSize::new(in_size),
908 }),
909 },
910 wgpu::BindGroupEntry { binding: 2, resource: tab.as_entire_binding() },
911 wgpu::BindGroupEntry { binding: 3, resource: lutb.as_entire_binding() },
912 wgpu::BindGroupEntry {
913 binding: 4,
914 resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
915 buffer: out_buf,
916 offset: out_offset,
917 size: wgpu::BufferSize::new(out_size),
918 }),
919 },
920 ],
921 });
922 let groups = match op {
923 GpuOp::Unpack => grid_1d((params[i].total as u64).div_ceil(2)),
924 GpuOp::Normalize => grid_1d(params[i].total as u64),
925 GpuOp::Develop(_) => (p.dev_w.div_ceil(16), (b.y1 - b.y0).div_ceil(16)),
926 };
927 let edge_bg = edge_run.as_ref().map(|er| {
928 dev.create_bind_group(&wgpu::BindGroupDescriptor {
929 label: Some("flashraw.edge"),
930 layout: &self.edge_layout,
931 entries: &[
932 wgpu::BindGroupEntry {
933 binding: 0,
934 resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
935 buffer: &er.params,
936 offset: i as u64 * er.stride,
937 size: wgpu::BufferSize::new(std::mem::size_of::<kernels::EdgeParams>() as u64),
938 }),
939 },
940 wgpu::BindGroupEntry { binding: 1, resource: er.tables.as_entire_binding() },
941 wgpu::BindGroupEntry { binding: 3, resource: er.tables.as_entire_binding() },
942 wgpu::BindGroupEntry {
943 binding: 2,
944 resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
945 buffer: &er.scratch,
946 offset: 0,
947 size: wgpu::BufferSize::new(er.bands[i].2),
948 }),
949 },
950 ],
951 })
952 });
953 let mut enc = dev.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: Some("flashraw.band") });
954 if let Some(st) = &slot.staging {
955 enc.copy_buffer_to_buffer(st, 0, &slot.src, 0, in_size);
956 }
957 {
958 let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor {
959 label: Some("flashraw.kernel"),
960 timestamp_writes: ts.as_ref().map(|(qs, _, _)| wgpu::ComputePassTimestampWrites {
961 query_set: qs,
962 beginning_of_pass_write_index: Some(2 * i as u32),
963 end_of_pass_write_index: Some(2 * i as u32 + 1),
964 }),
965 });
966 pass.set_bind_group(0, &bg, &[]);
967 if let (Some(er), Some(ebg)) = (&edge_run, &edge_bg) {
968 pass.set_bind_group(1, ebg, &[]);
969 let (n, (gx, gy), _) = er.bands[i];
970 if n > 0 {
971 for pl in &er.passes {
972 pass.set_pipeline(pl);
973 pass.dispatch_workgroups(gx, gy, 1);
974 }
975 }
976 }
977 pass.set_pipeline(&pipeline);
978 pass.dispatch_workgroups(groups.0, groups.1, 1);
979 }
980 if let Some((qs, resolve, read)) = ts.as_ref().filter(|_| i + 1 == bands.len()) {
981 let n = bands.len() as u32 * 2;
982 enc.resolve_query_set(qs, 0..n, resolve, 0);
983 enc.copy_buffer_to_buffer(resolve, 0, read, 0, n as u64 * 8);
984 }
985 slot.last = Some(self.queue.submit([enc.finish()]));
986 }
987 self.wait_idle().await?;
988 let compute_ms = (flashraw_core::profile::now_ms() - t0 - upload_ms).max(0.0);
989 let kernel_ms = match &ts {
990 Some((_, _, read)) => {
991 let n = bands.len() * 2;
992 let ticks = self
993 .read_mapped(read, n as u64 * 8, |b| {
994 b.chunks_exact(16)
995 .map(|c| {
996 let t0 = bytemuck::pod_read_unaligned::<u64>(&c[0..8]);
997 let t1 = bytemuck::pod_read_unaligned::<u64>(&c[8..16]);
998 t1.saturating_sub(t0)
999 })
1000 .sum::<u64>()
1001 })
1002 .await?;
1003 let ns = ticks as f64 * self.queue.get_timestamp_period() as f64;
1004 (ns > 0.0).then_some(ns / 1e6)
1006 }
1007 None => None,
1008 };
1009 if let Some(er) = edge_run {
1010 self.release(er.scratch);
1011 }
1012 for s in slots {
1013 self.release(s.src);
1014 if let Some(st) = s.staging {
1015 self.release(st);
1016 }
1017 }
1018 let e2 = oom.pop().await;
1019 let e1 = validation.pop().await;
1020 if let Some(e) = e2.or(e1) {
1021 return Err(Error::GpuUnavailable(format!("GPU error: {e}")));
1022 }
1023 if let Some(e) = self.take_error() {
1024 return Err(Error::GpuUnavailable(format!("GPU error: {e}")));
1025 }
1026 let tiled = layout == OutLayout::Tiled;
1027 let parts: Vec<GpuTile> = if tiled {
1028 outputs.into_iter().zip(&out_ranges).map(|(buffer, &(rect, _, _))| GpuTile { rect, buffer }).collect()
1029 } else {
1030 outputs.into_iter().map(|buffer| GpuTile { rect: Rect::full(out_w, out_h), buffer }).collect()
1031 };
1032 Ok(GpuRun {
1033 image: GpuImage { ctx: self.clone(), parts, tiled, width: out_w, height: out_h, channels, format },
1034 timings: GpuTimings { upload_ms, compute_ms, kernel_ms, bands: bands.len() as u32 },
1035 })
1036 }
1037
1038 #[cfg(not(target_arch = "wasm32"))]
1040 pub fn run_blocking(self: &Arc<Self>, frame: &GpuFrame<'_>, info: &FrameInfo, op: &GpuOp) -> Result<GpuRun> {
1041 pollster::block_on(self.run(frame, info, op))
1042 }
1043}
1044
1045#[cfg(test)]
1046mod tests {
1047 use super::*;
1048
1049 fn geo(rows: u32, kind: BandKind, frame_h: u32) -> BandGeometry {
1050 BandGeometry {
1051 rows,
1052 in_row_bytes: 100,
1053 out_row_bytes: 400,
1054 kind,
1055 frame_h,
1056 row_align: 1,
1057 tile_rows: 16,
1058 max_rows: u32::MAX,
1059 }
1060 }
1061
1062 fn check_cover(bands: &[Band], rows: u32) {
1063 assert_eq!(bands.first().unwrap().y0, 0);
1064 assert_eq!(bands.last().unwrap().y1, rows);
1065 assert!(bands.windows(2).all(|w| w[0].y1 == w[1].y0));
1066 }
1067
1068 #[test]
1069 fn plan_untiled_when_it_fits() {
1070 let g = geo(100, BandKind::Bilinear { crop_y: 0 }, 100);
1071 let b = g.plan(1 << 20, 8, true, 0).unwrap();
1072 assert_eq!(b, vec![Band { y0: 0, y1: 100, in0: 0, in1: 100 }]);
1073 }
1074
1075 #[test]
1076 fn plan_bilinear_halo_and_tile_rounding() {
1077 let g = geo(100, BandKind::Bilinear { crop_y: 3 }, 110);
1079 let b = g.plan(4008, 8, true, 0).unwrap();
1080 check_cover(&b, 100);
1081 assert_eq!(b[0], Band { y0: 0, y1: 32, in0: 2, in1: 36 });
1082 assert_eq!(b[1], Band { y0: 32, y1: 64, in0: 34, in1: 68 });
1083 assert_eq!(b.last().unwrap().in1, 104);
1084 let b = g.plan(4000 + 8, 8, false, 8).unwrap();
1086 assert!(b.iter().all(|b| b.y1 - b.y0 <= 10));
1087 check_cover(&b, 100);
1088 }
1089
1090 #[test]
1091 fn plan_half_and_alignment() {
1092 let g = geo(50, BandKind::Half { crop_y: 1 }, 101);
1093 let b = g.plan(2008, 8, true, 0).unwrap();
1094 check_cover(&b, 50);
1095 assert_eq!(b[0], Band { y0: 0, y1: 10, in0: 1, in1: 21 });
1096 let mut g = geo(9, BandKind::Direct, 9);
1097 g.row_align = 2;
1098 g.tile_rows = 1;
1099 let b = g.plan(308, 8, true, 0).unwrap();
1100 assert!(b.iter().all(|b| b.y0 % 2 == 0));
1101 check_cover(&b, 9);
1102 assert!(g.plan(108, 8, true, 0).is_err());
1103 }
1104
1105 #[test]
1106 fn band_rects_cover_every_orientation() {
1107 let (dw, dh) = (7u32, 10u32);
1108 for o in 1..=8u32 {
1109 let (ow, oh) = if o >= 5 { (dh, dw) } else { (dw, dh) };
1110 let mut seen = vec![0u8; (ow * oh) as usize];
1111 for (y0, y1) in [(0, 3), (3, 4), (4, 10)] {
1112 let r = band_rect(y0, y1, dw, dh, o);
1113 assert!(r.right() <= ow && r.bottom() <= oh, "o {o} {r}");
1114 for dy in y0..y1 {
1116 for dx in 0..dw {
1117 let (ox, oy) = match o {
1118 2 => (dw - 1 - dx, dy),
1119 3 => (dw - 1 - dx, dh - 1 - dy),
1120 4 => (dx, dh - 1 - dy),
1121 5 => (dy, dx),
1122 6 => (dh - 1 - dy, dx),
1123 7 => (dh - 1 - dy, dw - 1 - dx),
1124 8 => (dy, dw - 1 - dx),
1125 _ => (dx, dy),
1126 };
1127 assert!(r.contains_rect(&Rect::new(ox, oy, 1, 1)), "o {o} ({ox},{oy}) outside {r}");
1128 seen[(oy * ow + ox) as usize] += 1;
1129 }
1130 }
1131 }
1132 assert!(seen.iter().all(|&s| s == 1), "o {o}");
1133 }
1134 }
1135}