diff --git a/crates/engine/src/background/offload.rs b/crates/engine/src/background/offload.rs index 23cff4da6..f9038accb 100644 --- a/crates/engine/src/background/offload.rs +++ b/crates/engine/src/background/offload.rs @@ -558,6 +558,7 @@ async fn fill(s: &mut Session, plan: crate::commands::content_fill::FillPlan, po true }; let accel = job_accel(); + let fill = accel.as_deref().and_then(|a| a.inpaint()); let files = crate::commands::content_fill::run_plan( &plan, &s.project, @@ -566,6 +567,7 @@ async fn fill(s: &mut Session, plan: crate::commands::content_fill::FillPlan, po &s.layer_cache, s.services.as_ref(), accel.as_deref(), + fill, &mut progress, &mut |_| {}, ) diff --git a/crates/engine/src/cmd/content_fill.rs b/crates/engine/src/cmd/content_fill.rs index d96be073d..967651a37 100644 --- a/crates/engine/src/cmd/content_fill.rs +++ b/crates/engine/src/cmd/content_fill.rs @@ -13,7 +13,7 @@ use std::sync::Arc; use effectcraft_color::Label; use effectcraft_project::{Footage, FootageKind, ItemId, ItemKind, LayerId, LayerSource, build}; use effectcraft_raster::Image; -use effectcraft_raster::inpaint::{FillInput, FillMethod, FillOpts, dilate, fill_sequence}; +use effectcraft_raster::inpaint::{FillInput, FillMethod, FillOpts, dilate}; use effectcraft_render::{Accelerator, EvalCtx, ExprHost, FootageSource, LayerCache, Renderer}; use effectcraft_time::Tick; use serde::{Deserialize, Serialize}; @@ -185,8 +185,21 @@ fn memory_shortfall(frames: usize, w: u32, h: u32, available: u64) -> Option &'static str { + if on_device && method != FillMethod::EdgeBlend { "Filling (GPU)" } else { "Filling" } +} + /// Run a fill: the layer's frames (rendered in passes with a job worker's GPU), the fill, the /// PNG sequence written into `plan.dir` (`Services::store_file`). Returns the files. +/// +/// `accel` renders the frames (a job worker's GPU, as for every analysis: `None` on the +/// desktop, whose analysis threads render on the CPU by design); `fill` is the fill's own +/// backend ([`effectcraft_raster::inpaint::FillBackend`]), which is what makes Content-Aware +/// Fill itself run on the GPU. They are separate on purpose: a desktop fill can run on the GPU +/// while its input frames are composited on the CPU. #[allow(clippy::too_many_arguments)] pub(crate) async fn run_plan( plan: &FillPlan, @@ -196,6 +209,7 @@ pub(crate) async fn run_plan( cache: &LayerCache, services: &dyn crate::Services, accel: Option<&dyn Accelerator>, + fill: Option<&dyn effectcraft_raster::inpaint::FillBackend>, progress: &mut dyn FnMut(u64, u64) -> bool, message: &mut dyn FnMut(&str), ) -> std::result::Result, String> { @@ -233,10 +247,13 @@ pub(crate) async fn run_plan( references.push((k, img)); } } - message("Filling"); let opts = FillOpts { method: FillMethod::from_name(&plan.method).unwrap_or_default(), lighting: plan.lighting, ..Default::default() }; + message(filling_status(fill.is_some(), opts.method)); let input = FillInput { frames, holes, references }; - let out = fill_sequence(&input, &opts, &mut |d, t| progress(n + (d as u64 * n / t.max(1) as u64), n * 3)).ok_or(crate::render_queue::CANCELLED)?; + let cpu = effectcraft_raster::inpaint::CpuFill; + let backend: &dyn effectcraft_raster::inpaint::FillBackend = fill.unwrap_or(&cpu); + let out = effectcraft_raster::inpaint::fill_sequence_with(&input, &opts, backend, &mut |d, t| progress(n + (d as u64 * n / t.max(1) as u64), n * 3)) + .ok_or(crate::render_queue::CANCELLED)?; message("Writing the fill layer"); let dir = &plan.dir; let mut paths = vec![]; @@ -324,7 +341,12 @@ fn generate(s: &mut Session, p: &Value) -> Result { let services = s.services.clone(); let importer = s.importer.clone(); let name = layer.name.clone(); + // The fill's own GPU backend (Mercury GPU Acceleration; none under Software Only or + // without a compositor). The frames themselves still render on the CPU here, as for every + // desktop analysis (`offload::job_accel`): only the fill's own work moves to the device. + let accel = if s.project.settings.gpu_acceleration { s.accel.clone() } else { None }; s.spawn_task("contentFill", format!("Content-Aware Fill: {name}"), wait, move |ctl| { + let fill = accel.as_deref().and_then(|a| a.inpaint()); let paths = effectcraft_render::passes::block_on(run_plan( &plan, &lf.project, @@ -333,6 +355,7 @@ fn generate(s: &mut Session, p: &Value) -> Result { &lf.cache, services.as_ref(), None, + fill, &mut |d, t| ctl.progress(d, t), &mut |m| ctl.message(m), ))?; @@ -445,7 +468,18 @@ pub fn specs() -> Vec { #[cfg(test)] mod tests { - use super::memory_shortfall; + use super::{filling_status, memory_shortfall}; + + /// Edge Blend fills with the membrane, which the device backend leaves to the CPU reference: + /// it must not claim a GPU fill just because a backend is available. + #[test] + fn edge_blend_is_not_reported_as_a_device_fill() { + use effectcraft_raster::inpaint::FillMethod; + assert_eq!(filling_status(true, FillMethod::Object), "Filling (GPU)"); + assert_eq!(filling_status(true, FillMethod::Surface), "Filling (GPU)"); + assert_eq!(filling_status(true, FillMethod::EdgeBlend), "Filling"); + assert_eq!(filling_status(false, FillMethod::Object), "Filling"); + } /// #1789: 10 s of 1080p30 (300 frames) needs at least 20 GB, which 4 GB free can't hold. #[test] diff --git a/crates/gpu/src/context.rs b/crates/gpu/src/context.rs index 543c79283..713f6f83d 100644 --- a/crates/gpu/src/context.rs +++ b/crates/gpu/src/context.rs @@ -803,6 +803,7 @@ impl GpuContext { include_str!("shaders/fx_time.wgsl"), include_str!("shaders/fx_pixel2.wgsl"), include_str!("shaders/fx_gen2.wgsl"), + include_str!("shaders/inpaint.wgsl"), include_str!("shaders/sky.wgsl"), ] .concat(); @@ -911,6 +912,7 @@ impl GpuContext { .chain(crate::fx_time::KERNELS) .chain(crate::fx_pixel2::KERNELS) .chain(crate::fx_gen2::KERNELS) + .chain(crate::inpaint::KERNELS) { match build(e, &layout) { Ok(p) => { @@ -1893,23 +1895,31 @@ impl<'g> Enc<'g> { data.resize(layout.len / 16, [0.0; 4]); let mut out = Image { width: img.width, height: img.height, data }; let dst: &mut [u8] = bytemuck::cast_slice_mut(&mut out.data); - self.read_texture_into(&img.texture, img.width, img.height, 16, dst)?; + self.read_texture_into(&img.texture, [0, 0, img.width, img.height], 16, dst)?; Some(out) } /// Tightly packed texel bytes of a texture (`bpp` bytes per pixel). pub fn read_texture(&mut self, texture: &wgpu::Texture, w: u32, h: u32, bpp: u32) -> Option> { + self.read_rect(texture, 0, 0, w, h, bpp) + } + + /// [`Enc::read_texture`] of a sub-rectangle at (`x`, `y`): a fill only changes its hole, so + /// only the hole's bounds need to cross back (`bytes_per_row` is still row-aligned inside). + pub fn read_rect(&mut self, texture: &wgpu::Texture, x: u32, y: u32, w: u32, h: u32, bpp: u32) -> Option> { let layout = crate::readback::Layout::new(w, h, bpp, self.g.device.limits().max_buffer_size).map_err(|e| log::error!("gpu texture readback: {e}")).ok()?; let mut out = Vec::new(); out.try_reserve_exact(layout.len).map_err(|e| log::error!("gpu readback allocation: {e}")).ok()?; out.resize(layout.len, 0); - self.read_texture_into(texture, w, h, bpp, &mut out)?; + self.read_texture_into(texture, [x, y, w, h], bpp, &mut out)?; Some(out) } - /// [`Enc::read_texture`] into `out` (`w × bpp × h` bytes), through a reused staging buffer. - fn read_texture_into(&mut self, texture: &wgpu::Texture, w: u32, h: u32, bpp: u32, out: &mut [u8]) -> Option<()> { + /// [`Enc::read_rect`] into `out` (`w × bpp × h` bytes), through a reused staging buffer. + /// `rect` is `[x, y, w, h]`. + fn read_texture_into(&mut self, texture: &wgpu::Texture, rect: [u32; 4], bpp: u32, out: &mut [u8]) -> Option<()> { + let [x, y, w, h] = rect; if !self.g.can_wait() { return None; } @@ -1932,7 +1942,7 @@ impl<'g> Enc<'g> { }) }); self.encoder().copy_texture_to_buffer( - wgpu::TexelCopyTextureInfo { texture, mip_level: 0, origin: wgpu::Origin3d::ZERO, aspect: wgpu::TextureAspect::All }, + wgpu::TexelCopyTextureInfo { texture, mip_level: 0, origin: wgpu::Origin3d { x, y, z: 0 }, aspect: wgpu::TextureAspect::All }, wgpu::TexelCopyBufferInfo { buffer: &buf, layout: wgpu::TexelCopyBufferLayout { offset: 0, bytes_per_row: Some(padded), rows_per_image: Some(h) } }, wgpu::Extent3d { width: w, height: h, depth_or_array_layers: 1 }, ); diff --git a/crates/gpu/src/inpaint.rs b/crates/gpu/src/inpaint.rs new file mode 100644 index 000000000..25cbc31b9 --- /dev/null +++ b/crates/gpu/src/inpaint.rs @@ -0,0 +1,729 @@ +//! Content-Aware Fill's per-frame synthesis and flow on the GPU: PatchMatch's multi-scale EM +//! loop, propagation, random search and vote of `effectcraft_raster::inpaint`, and the block +//! matching of `effectcraft_raster::flow::block_flow`, as compute kernels following the CPU's +//! steps (`shaders/inpaint.wgsl`). +//! +//! The fill's input frames are composited on the CPU by design (analyses do not use the job +//! worker's GPU: `effectcraft_engine::background::offload::job_accel`), so this module is about the fill +//! itself — the part that runs for minutes on a large mask, and whose parallel search is what a +//! device wins on. The membrane (harmonic) solve, the flow's hole completion, the orchestration +//! ([`effectcraft_raster::inpaint::fill_sequence_with`]), the flow-guided propagation and the +//! per-frame flow control stay on the CPU: the membrane is bandwidth-bound (see +//! [`FillBackend::patch`]), and the propagation's coupling is sequential, needing every frame of +//! the range resident at once. +//! +//! Parity with the CPU reference is a quality contract, not bit-exactness: the CPU's +//! scan-order propagation becomes a two-pass checkerboard, and its in-place relaxation becomes +//! Jacobi reads. Both solve the same problem to the same quality (`tests/inpaint.rs`). + +use effectcraft_raster::Image; +use effectcraft_raster::flow::{Flow, block_flow}; +use effectcraft_raster::inpaint::{FillBackend, complete_flow, dilate}; + +use crate::Gpu; +use crate::context::{Enc, GpuContext, GpuImage, KernelImages, Params}; + +/// The kernels this module adds to the shared pipeline registry (`context.rs`). +pub(crate) const KERNELS: &[&str] = &[ + "imp_down", + "imp_up", + "imp_jacobi", + "imp_pm_down", + "imp_pm_up", + "imp_nnf_up", + "imp_pm_init", + "imp_pm_prop", + "imp_pm_rand", + "imp_nnf_store", + "imp_pm_vote", + "imp_hist", + "imp_fl_half", + "imp_fl_search", + "imp_fl_subpix", + "imp_fl_median", +]; + +/// Distance histogram bins (`imp_hist`); bin 0's `.g` holds the largest finite distance the +/// dispatch saw, for the HDR re-bin. +const BINS: usize = 1024; + +/// The distance range `imp_hist`'s bins cover for SDR input. An image whose channels are all ≤ 1 +/// has `Σ_c (a − b)² ≤ 4`, so no patch distance exceeds it and the bins need no rescaling. +const SDR_DISTANCE: f32 = 4.0; + +/// `u`/`f` blocks of a kernel's parameter buffer (see `inpaint.wgsl`'s header). +fn params(u0: [u32; 4], u1: [u32; 4], u2: [u32; 4], f0: [f32; 4]) -> Params { + Params { u: [u0, u1, u2, [0; 4]], f: [f0, [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4], [0.0; 4]] } +} + +/// `p` with this level's mask offset in `u[3].x` (`inpaint.wgsl`'s `data`). +fn with_off(p: Params, off: u32) -> Params { + Params { u: [p.u[0], p.u[1], p.u[2], [off, 0, 0, 0]], f: p.f } +} + +/// A relaxation sweep count of `membrane_core`. The CPU uses 200 for the small levels and 40 for +/// the large ones; this uses 64 and 40. The small levels are where the difference lies: Jacobi +/// needs O(n²) sweeps to converge on an n-wide grid, so at the levels the coarse-to-fine schedule +/// actually reaches (n ≤ 16, where n² ≤ 256) 64 sweeps is already converged, and finer levels +/// start from a good guess. Matching the CPU's 200 measured ≈ 6% slower on the whole fill with no +/// change in the result to two decimals (67.0 dB either way), so the schedule is deliberately +/// not the CPU's — Jacobi and Gauss–Seidel do not converge alike, and the parity contract is +/// quality, not bit-exactness. +fn sweep_count(w: usize, h: usize) -> usize { + if w * h <= 4096 { 64 } else { 40 } +} + +/// PatchMatch's patch radius at a level (`pm_level`). +fn radius(patch: usize, w: u32, h: u32) -> u32 { + (patch / 2).min(w.min(h).saturating_sub(1) as usize / 2).max(1) as u32 +} + +/// The bounding box of `mask`'s set pixels, `[x0, y0, x1, y1)`, or `None` when it is empty. +fn bbox(mask: &[bool], w: usize, h: usize) -> Option<[u32; 4]> { + let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize); + for y in 0..h { + for x in 0..w { + if mask.get(y * w + x).copied().unwrap_or(false) { + x0 = x0.min(x); + y0 = y0.min(y); + x1 = x1.max(x + 1); + y1 = y1.max(y + 1); + } + } + } + if x0 == usize::MAX { None } else { Some([x0 as u32, y0 as u32, x1 as u32, y1 as u32]) } +} + +/// [`bbox`] grown by `n` pixels, clipped to the image. +fn grow(b: [u32; 4], n: u32, w: u32, h: u32) -> [u32; 4] { + [b[0].saturating_sub(n), b[1].saturating_sub(n), b[2].saturating_add(n).min(w), b[3].saturating_add(n).min(h)] +} + +/// The bounding box of the pixels a membrane level still has to solve (`known` false), or +/// `None` when every pixel is known. +fn unknown_bbox(known: &[f32], w: usize, h: usize) -> Option<[u32; 4]> { + let (mut x0, mut y0, mut x1, mut y1) = (usize::MAX, usize::MAX, 0usize, 0usize); + for y in 0..h { + for x in 0..w { + if known.get(y * w + x).copied().unwrap_or(1.0) <= 0.5 { + x0 = x0.min(x); + y0 = y0.min(y); + x1 = x1.max(x + 1); + y1 = y1.max(y + 1); + } + } + } + if x0 == usize::MAX { None } else { Some([x0 as u32, y0 as u32, x1 as u32, y1 as u32]) } +} + +/// One dispatch's images, its per-pixel `data` table and the region it covers. Grouped into a +/// struct because a kernel's bindings plus its bounds exceed clippy's argument limit, and the +/// workspace denies `too_many_arguments`. +struct Pass<'a> { + src: &'a GpuImage, + aux: Option<&'a GpuImage>, + out: &'a GpuImage, + data: Option<&'a wgpu::Buffer>, + bb: [u32; 4], +} + +/// Dispatch a kernel over `pass.bb`. The box's origin is written into `P.u[1].xy` (the kernels +/// place their other `u[1]` fields in `.zw`): deriving it from the box the dispatch is sized from +/// keeps a dispatch's covered region and the origin its shader reads from drifting apart, which +/// is the subtle way an edge of the region goes unwritten. +fn over(e: &mut Enc, entry: &str, p: &Params, pass: Pass<'_>) { + let bb = pass.bb; + let p = Params { u: [p.u[0], [bb[0], bb[1], p.u[1][2], p.u[1][3]], p.u[2], p.u[3]], f: p.f }; + let groups = ((bb[2] - bb[0]).div_ceil(16), (bb[3] - bb[1]).div_ceil(16)); + e.dispatch(entry, &p, KernelImages { src: pass.src, aux: pass.aux, out: pass.out }, pass.data, groups); +} + +// ------------------------------------------------------------------ membrane + +/// Relax `fine`'s unknown pixels — the harmonic interpolation of their border — with the CPU's +/// coarse-to-fine schedule (`membrane_core`). `spare` is a free texture of the same size; the +/// result is returned (it may be `fine`, `spare` or one of the coarse levels). +fn membrane_tex(e: &mut Enc, fine: &GpuImage, spare: &GpuImage, hole: &[bool], w: u32, h: u32) -> Option { + let (uw, uh) = (w as usize, h as usize); + let n = uw * uh; + if hole.len() < n || !hole.iter().take(n).any(|b| *b) || !hole.iter().take(n).any(|b| !*b) { + return Some(fine.clone()); + } + // The recursion's levels: halve while the level is wider than 8. + let mut dims = vec![(uw, uh)]; + while let Some(&(cw, ch)) = dims.last() { + if !(cw > 8 && ch > 8) || dims.len() >= 32 { + break; + } + dims.push((cw.div_ceil(2), ch.div_ceil(2))); + } + // The known mask of every level: a coarse pixel is known where any of its 2×2 is. + let mut knowns: Vec> = Vec::with_capacity(dims.len()); + knowns.push(hole.iter().take(n).map(|b| if *b { 0.0 } else { 1.0 }).collect()); + for k in 1..dims.len() { + let (pw, ph) = dims[k - 1]; + let (cw, ch) = dims[k]; + let prev = &knowns[k - 1]; + let mut cur = vec![0.0f32; cw * ch]; + for y in 0..ch { + for x in 0..cw { + for (dx, dy) in [(0usize, 0usize), (1, 0), (0, 1), (1, 1)] { + let (fx, fy) = (2 * x + dx, 2 * y + dy); + if fx < pw && fy < ph && prev[fy * pw + fx] > 0.5 { + cur[y * cw + x] = 1.0; + break; + } + } + } + } + knowns.push(cur); + } + // One buffer for every level's mask (a buffer per level costs a device allocation each). + let mut packed: Vec = Vec::with_capacity(knowns.iter().map(|k| k.len()).sum()); + let mut offs: Vec = Vec::with_capacity(knowns.len()); + for k in &knowns { + offs.push(packed.len() as u32); + packed.extend_from_slice(k); + } + let datas = e.data(&packed); + let moff = |k: usize| offs.get(k).copied().unwrap_or(0); + // The bounds of each level's unknown pixels: the relaxation only touches them (known pixels + // never move), which keeps a fill's cost proportional to the hole, not the frame. + let mut uns: Vec> = Vec::with_capacity(dims.len()); + for (k, &(cw, ch)) in dims.iter().enumerate() { + uns.push(unknown_bbox(knowns.get(k)?, cw, ch)); + } + // Two textures per coarse level, for the ping-pong sweeps. + let mut lv: Vec = Vec::with_capacity(2 * dims.len()); + for k in 1..dims.len() { + let (cw, ch) = dims[k]; + lv.push(e.scratch(cw as u32, ch as u32)); + lv.push(e.scratch(cw as u32, ch as u32)); + } + // Fine to coarse. + for k in 1..dims.len() { + let (pw, ph) = dims[k - 1]; + let (cw, ch) = dims[k]; + let p = with_off(params([pw as u32, ph as u32, cw as u32, ch as u32], [0; 4], [0; 4], [0.0; 4]), moff(k - 1)); + let src = if k == 1 { fine } else { lv.get(2 * (k - 2))? }; + e.pixels("imp_down", &p, src, None, lv.get(2 * (k - 1))?, Some(&datas)); + } + // Relax, coarse to fine. `first[k]` says whether level k's image is the first of its pair + // (level 0's pair is `fine`/`spare`). + let top = dims.len() - 1; + let mut first = vec![true; dims.len()]; + if top >= 1 { + let (a, b) = (lv.get(2 * (top - 1))?, lv.get(2 * (top - 1) + 1)?); + // Both textures must hold the level before the sweeps: they now write only the hole. + e.copy_into(a, b, 0, 0); + first[top] = sweeps(e, &datas, moff(top), a, b, Region { w: dims[top].0 as u32, h: dims[top].1 as u32, bb: uns.get(top).copied().flatten() }); + } + for k in (1..top).rev() { + let (cw, ch) = dims[k]; + let (pw, ph) = dims[k + 1]; + let coarse = lv.get(2 * k + usize::from(!first[k + 1]))?; + let (fa, fb) = (lv.get(2 * (k - 1))?, lv.get(2 * (k - 1) + 1)?); + e.copy_into(fa, fb, 0, 0); + if let Some(bb) = uns.get(k).copied().flatten() { + let p = with_off(params([cw as u32, ch as u32, pw as u32, ph as u32], [0; 4], [0; 4], [0.0; 4]), moff(k)); + over(e, "imp_up", &p, Pass { src: coarse, aux: None, out: fb, data: Some(&datas), bb }); + } + // `imp_up` wrote into `fb`, the pair's second. + first[k] = !sweeps(e, &datas, moff(k), fb, fa, Region { w: cw as u32, h: ch as u32, bb: uns.get(k).copied().flatten() }); + } + if top >= 1 { + let coarse = lv.get(usize::from(!first[1]))?; + e.copy_into(fine, spare, 0, 0); + let bb = uns.first().copied().flatten()?; + let p = with_off(params([w, h, dims[1].0 as u32, dims[1].1 as u32], [0; 4], [0; 4], [0.0; 4]), moff(0)); + over(e, "imp_up", &p, Pass { src: coarse, aux: None, out: spare, data: Some(&datas), bb }); + return Some(if sweeps(e, &datas, moff(0), spare, fine, Region { w, h, bb: Some(bb) }) { spare.clone() } else { fine.clone() }); + } + let bb = uns.first().copied().flatten()?; + e.copy_into(fine, spare, 0, 0); + Some(if sweeps(e, &datas, moff(0), fine, spare, Region { w, h, bb: Some(bb) }) { fine.clone() } else { spare.clone() }) +} + +/// A texture's dimensions and the region a dispatch covers (`None`: nothing to do). +struct Region { + w: u32, + h: u32, + bb: Option<[u32; 4]>, +} + +/// Jacobi relaxation sweeps of a texture pair, over the unknown pixels' bounds `region.bb` (the +/// state starts in `a`); returns whether the result is in `a`. +fn sweeps(e: &mut Enc, data: &wgpu::Buffer, off: u32, a: &GpuImage, b: &GpuImage, region: Region) -> bool { + let Some(bb) = region.bb else { return true }; + let p = with_off(params([region.w, region.h, 0, 0], [0; 4], [0; 4], [0.0; 4]), off); + let mut in_a = true; + for _ in 0..sweep_count(region.w as usize, region.h as usize) { + let (src, out) = if in_a { (a, b) } else { (b, a) }; + over(e, "imp_jacobi", &p, Pass { src, aux: None, out, data: Some(data), bb }); + in_a = !in_a; + } + in_a +} + +/// `img` with the pixels `mask` marks inside `bb` replaced by `rect` — the tightly packed RGBA f32 +/// rows of the rectangle that was read back. Pixels outside the mask stay the input's, bit for bit. +fn patch_back(img: &Image, mask: &[bool], bb: [u32; 4], rect: &[u8]) -> Image { + let (w, rw) = (img.width as usize, (bb[2] - bb[0]) as usize); + let mut out = img.clone(); + for y in bb[1]..bb[3] { + for x in bb[0]..bb[2] { + let i = y as usize * w + x as usize; + if !mask.get(i).copied().unwrap_or(false) { + continue; + } + let off = ((y - bb[1]) as usize * rw + (x - bb[0]) as usize) * 16; + let Some(b) = rect.get(off..off + 16) else { continue }; + let mut px = [0.0f32; 4]; + for (c, v) in px.iter_mut().enumerate() { + *v = f32::from_le_bytes([b[c * 4], b[c * 4 + 1], b[c * 4 + 2], b[c * 4 + 3]]); + } + out.data[i] = px; + } + } + out +} + +impl FillBackend for Gpu { + /// PatchMatch synthesis — the fill's dominant cost, and the stage a device wins on: two to + /// three times faster than the CPU's 18-rayon-thread loop at 1920×1080 with a 300×200 hole + /// on an Intel Arc iGPU, and far more on a discrete one. + /// + /// The membrane ([`FillBackend::membrane`]) deliberately stays on the CPU: a harmonic solve + /// over a hole is only a few tens of milliseconds, and on the same measurement it is *slower* + /// on the device (≈ 90 ms against ≈ 55 ms). Most of that is the fill's full-frame data + /// movement — uploading the frame, copying it and building the known pyramid — which does not + /// shrink with the hole, while the relaxation itself does. The CPU's in-place Gauss–Seidel + /// over many cores is already competitive there, so the device only takes the stage where its + /// parallel search actually pays. (`membrane_tex` still runs on the device *inside* this + /// stage, where it is the coarsest level's start and never crosses back to the CPU.) + fn patch(&self, img: &Image, hole: &[bool], init: Option<&Image>, patch: usize, seed: u32) -> Image { + // The error scopes make an allocation or validation failure local to this frame: the + // CPU fills it instead, rather than an uncaptured error retiring the shared context. + if self.ctx.can_wait() + && let Ok(Some(out)) = self.within_memory(|| patch_image(&self.ctx, img, hole, init, patch, seed)) + { + return out; + } + effectcraft_raster::inpaint::patch_fill(img, hole, init, patch, seed) + } + + /// The block matching that the sequence's flow-guided propagation starts from, which is the + /// flow's dominant cost. The field is a small `cols × rows` grid (8 px tiles, so 240 × 135 at + /// 1080p), and the search is texture-load-bound: at 1920×1080 the CPU's `block_flow` is + /// ≈ 65–70 ms against ≈ 55 ms here, most of it the coarsest level's window (`radius 6` over + /// up to six seeds) and the finest level's 8×8 blocks. That is a smaller margin than + /// [`patch`](Self::patch), but the flow runs twice per frame pair, so on a four-frame 1080p + /// sequence it is worth ≈ 0.15× of the fill's total (1.42× against 1.27× with the CPU flow). + /// + /// The field is exhaustive, not heuristic: it matches the CPU's to a float rounding, so the + /// propagation sees the same motion. Both A's block per candidate (hoisted into registers) + /// and the CPU's `Plane::luma` (computed here by [`effectcraft_raster::flow::luma_pair`] and + /// uploaded, because a device may contract `a·b + c` into a multiply-add and round the luma + /// differently) are why it is that close. + /// + /// The completion — marking the tiles that touch either hole and interpolating the surrounding + /// motion over them — stays on the CPU ([`complete_flow`]). It is a membrane solve over that + /// same tiny grid whose cost is the mask, not the solve, so there is little for the device to + /// take; keeping it here also keeps the completed field the CPU's to within the search's + /// rounding. + fn flow(&self, a: &Image, b: &Image, hole_a: &[bool], hole_b: &[bool]) -> Flow { + // Frame-local error scopes, as for the patch: a failed GPU search falls back to the CPU. + if self.ctx.can_wait() + && let Ok(Some(f)) = self.within_memory(|| flow_image(&self.ctx, a, b)) + { + return complete_flow(f, hole_a, hole_b, a.width as usize, a.height as usize); + } + complete_flow(block_flow(a, b, 8, 6), hole_a, hole_b, a.width as usize, a.height as usize) + } +} + +// ------------------------------------------------------------------ block flow + +/// `effectcraft_raster::flow::block_flow(a, b, 8, 6)` on the device: the two frames' luma +/// pyramids, the coarse-to-fine search, the sub-pixel fit and the 3×3 median, returning the raw +/// field ([`complete_flow`] does the hole completion). `None` where the device cannot run it and +/// the caller falls back to the CPU. +pub(crate) fn flow_image(g: &GpuContext, a: &Image, b: &Image) -> Option { + let (w, h) = (a.width, a.height); + if w == 0 || h == 0 || w != b.width || h != b.height || !g.fits(w, h) { + return None; + } + let (tile, win, radius) = (8u32, 8u32, 6u32); + let cols = (w as usize).div_ceil(tile as usize).max(1); + let rows = (h as usize).div_ceil(tile as usize).max(1); + // block_flow_window's early out: a frame too small to hold one block and window. + if w < tile || h < tile || w < win || h < win { + return Some(Flow { width: w, height: h, block: tile, cols, rows, v: vec![[0.0; 2]; cols * rows] }); + } + let mut e = Enc::new(g); + // The luma pyramid: `Plane::luma` (packed .x/.y), halved while both sides stay worth matching + // (`Plane::half`), at most four levels. + let mut dims = vec![(w, h)]; + while dims.len() < 4 { + let (cw, ch) = *dims.last()?; + if cw / 2 < 16 || ch / 2 < 16 { + break; + } + dims.push(((cw / 2).max(1), (ch / 2).max(1))); + } + // The luma is the CPU's (`luma_pair`), so the field matches `block_flow`'s bit for bit; the + // halves below are exact additions and a multiply by a power of two. + let pair = g.upload_image(&effectcraft_raster::flow::luma_pair(a, b))?; + let mut luma: Vec = Vec::with_capacity(dims.len()); + luma.push(pair); + for k in 1..dims.len() { + let (pw, ph) = dims[k - 1]; + let (cw, ch) = dims[k]; + let dst = e.scratch(cw, ch); + e.pixels("imp_fl_half", ¶ms([pw, ph, cw, ch], [0; 4], [0; 4], [0.0; 4]), luma.get(k - 1)?, None, &dst, None); + luma.push(dst); + } + // The coarse-to-fine search: the coarsest level from the zero vector, each finer level from + // the doubled vectors of its tile and its four neighbours. + let top = dims.len() - 1; + let groups = (cols.div_ceil(16) as u32, rows.div_ceil(16) as u32); + let mut field: Vec = Vec::with_capacity(dims.len()); + for _ in 0..dims.len() { + field.push(e.scratch(cols as u32, rows as u32)); + } + for k in (0..dims.len()).rev() { + let (pw, ph) = dims[k]; + let coarser = (k != top).then(|| field.get(k + 1)).flatten(); + let p = params([pw, ph, cols as u32, rows as u32], [tile, win, 1u32 << k, if k == top { radius } else { 1 }], [u32::from(k != top), 0, 0, 0], [0.0; 4]); + e.dispatch("imp_fl_search", &p, KernelImages { src: luma.get(k)?, aux: coarser, out: field.get(k)? }, None, groups); + } + // The full-resolution sub-pixel fit, then the vector median. + let fine = e.scratch(cols as u32, rows as u32); + e.dispatch( + "imp_fl_subpix", + ¶ms([w, h, cols as u32, rows as u32], [tile, win, 0, 0], [0; 4], [0.0; 4]), + KernelImages { src: luma.first()?, aux: Some(field.first()?), out: &fine }, + None, + groups, + ); + let out = e.scratch(cols as u32, rows as u32); + e.dispatch( + "imp_fl_median", + ¶ms([w, h, cols as u32, rows as u32], [0; 4], [0; 4], [0.0; 4]), + KernelImages { src: &fine, aux: None, out: &out }, + None, + groups, + ); + if e.take_missing() { + return None; + } + // Only the field crosses back to the CPU: `cols × rows` entries. + let rect = e.read_rect(&out.texture, 0, 0, cols as u32, rows as u32, 16)?; + let mut v = Vec::with_capacity(cols * rows); + for i in 0..cols * rows { + let b = rect.get(i * 16..i * 16 + 16)?; + let f = |o: usize| b.get(o).copied().unwrap_or(0); + v.push([f32::from_le_bytes([f(0), f(1), f(2), f(3)]), f32::from_le_bytes([f(4), f(5), f(6), f(7)])]); + } + Some(Flow { width: w, height: h, block: tile, cols, rows, v }) +} + +// ------------------------------------------------------------------ PatchMatch + +/// One PatchMatch level's plan. +struct Lv { + w: u32, + h: u32, + r: u32, + /// The target centres: pixels whose patch overlaps the hole (`near`). + near: [u32; 4], + /// The hole's own bounding box (the vote's dispatch). + holes: [u32; 4], + /// Where the nearest-neighbour field must be defined: the vote reads it `r` pixels out. + field: [u32; 4], + flags: wgpu::Buffer, +} + +/// The multi-scale PatchMatch synthesis of `patch_fill`: `init` is the starting guess (the +/// previous frame's fill warped onto this one), which refines the finest level only. +pub(crate) fn patch_image(g: &GpuContext, img: &Image, hole: &[bool], init: Option<&Image>, patch: usize, seed: u32) -> Option { + let (w, h) = (img.width, img.height); + if w == 0 || h == 0 || !g.fits(w, h) { + return None; + } + let (uw, uh) = (w as usize, h as usize); + let n = uw * uh; + if hole.len() < n || !hole.iter().take(n).any(|b| *b) { + return Some(img.clone()); + } + let hole: Vec = hole.iter().take(n).copied().collect(); + let patch = patch.max(3) | 1; + let mut e = Enc::new(g); + + // A refinement of one frame (`patch_fill`'s `init`): the previous frame's fill seeds the + // hole and only the finest level runs, from 2 EM rounds of 4 iterations. + if let Some(init) = init { + if init.width != w || init.height != h { + return None; + } + let mut start = img.clone(); + for (i, v) in start.data.iter_mut().enumerate() { + if hole.get(i).copied().unwrap_or(false) { + *v = *init.data.get(i)?; + } + } + let lv = plan_level(&mut e, &hole, w, h, patch)?; + let image = g.upload_image(&start)?; + let spare = e.scratch(w, h); + let nnf = e.image(w, h); + let (out, _) = pm_level(&mut e, &lv, PmRun { seed, em: 2, iters: 4 }, &image, &spare, &nnf)?; + if e.take_missing() { + return None; + } + let bb = bbox(&hole, uw, uh)?; + let rect = e.read_rect(&out.texture, bb[0], bb[1], bb[2] - bb[0], bb[3] - bb[1], 16)?; + return Some(patch_back(img, &hole, bb, &rect)); + } + + // The pyramid: halve while the level is still a few patches across. + let mut dims = vec![(w, h)]; + let mut holes = vec![hole.clone()]; + while dims.len() < 6 { + let &(lw, lh) = dims.last()?; + let hl = holes.last()?; + if lw.min(lh) / 2 < patch as u32 * 4 || !hl.iter().any(|b| *b) { + break; + } + let (cw, ch) = (lw.div_ceil(2) as usize, lh.div_ceil(2) as usize); + let mut nh = vec![false; cw * ch]; + for y in 0..ch { + for x in 0..cw { + for (dx, dy) in [(0usize, 0usize), (1, 0), (0, 1), (1, 1)] { + let (fx, fy) = ((2 * x + dx).min(lw as usize - 1), (2 * y + dy).min(lh as usize - 1)); + if hl[fy * lw as usize + fx] { + nh[y * cw + x] = true; + break; + } + } + } + } + if !nh.iter().any(|b| !*b) { + break; + } + dims.push((cw as u32, ch as u32)); + holes.push(nh); + } + let levels = dims.len(); + let top = levels - 1; + + let mut plan = Vec::with_capacity(levels); + for k in 0..levels { + plan.push(plan_level(&mut e, holes.get(k)?, dims[k].0, dims[k].1, patch)?); + } + + // One image per level (the coarse-to-fine steps and the vote ping-pong in fresh textures). + let mut ima: Vec = Vec::with_capacity(levels); + ima.push(g.upload_image(img)?); + for k in 1..levels { + let (cw, ch) = dims[k]; + ima.push(e.scratch(cw, ch)); + } + for k in 1..levels { + let (pw, ph) = dims[k - 1]; + let (cw, ch) = dims[k]; + e.pixels("imp_pm_down", ¶ms([pw, ph, cw, ch], [0; 4], [0; 4], [0.0; 4]), &ima[k - 1], None, &ima[k], Some(&plan[k - 1].flags)); + } + // The coarsest level starts as its membrane fill. + let (cw, ch) = dims[top]; + let spare = e.scratch(cw, ch); + ima[top] = membrane_tex(&mut e, &ima[top], &spare, &holes[top], cw, ch)?; + + // Coarse to fine: upsample, then the EM rounds of propagation and voting. + let mut nnf: Option = None; + for k in (0..levels).rev() { + let lv = &plan[k]; + nnf = Some(match nnf { + Some(coarse) => { + let (cw, ch) = dims[k + 1]; + let fine = e.scratch(lv.w, lv.h); + e.pixels("imp_nnf_up", ¶ms([lv.w, lv.h, cw, ch], [0; 4], [0; 4], [0.0; 4]), &coarse, None, &fine, None); + fine + } + None => e.image(lv.w, lv.h), + }); + if k < top { + let (cw, ch) = dims[k + 1]; + let dst = e.scratch(lv.w, lv.h); + e.copy_into(&ima[k], &dst, 0, 0); + e.pixels("imp_pm_up", ¶ms([lv.w, lv.h, cw, ch], [0; 4], [0; 4], [0.0; 4]), &ima[k + 1], None, &dst, Some(&lv.flags)); + ima[k] = dst; + } + let (em, iters) = if k == top { (8, 5) } else { (4, 5) }; + let spare = e.scratch(lv.w, lv.h); + let run = PmRun { seed: seed.wrapping_add(k as u32 * 101), em, iters }; + let (out, out_nnf) = pm_level(&mut e, lv, run, &ima[k], &spare, nnf.as_ref()?)?; + ima[k] = out; + nnf = Some(out_nnf); + } + if e.take_missing() { + return None; + } + // Only the hole changes, so only the hole's bounds cross back to the CPU. + let bb = bbox(&hole, uw, uh)?; + let rect = e.read_rect(&ima[0].texture, bb[0], bb[1], bb[2] - bb[0], bb[3] - bb[1], 16)?; + Some(patch_back(img, &hole, bb, &rect)) +} + +/// The plan of one level: its flags (`bit 0` hole, `bit 1` a target centre) and the dispatch +/// bounds of its targets, holes and field. +fn plan_level(e: &mut Enc, hole: &[bool], w: u32, h: u32, patch: usize) -> Option { + let (iw, ih) = (w as usize, h as usize); + let r = radius(patch, w, h); + let near = dilate(hole, iw, ih, r as usize); + let near_bb = bbox(&near, iw, ih)?; + let flags: Vec = (0..iw * ih).map(|i| (u32::from(hole[i]) + 2 * u32::from(near[i])) as f32).collect(); + Some(Lv { w, h, r, near: near_bb, holes: bbox(hole, iw, ih)?, field: grow(near_bb, r + 1, w, h), flags: e.data(&flags) }) +} + +/// The PatchMatch run at one level: the random stream's seed and the EM round and iteration +/// counts. +struct PmRun { + seed: u32, + em: usize, + iters: usize, +} + +/// `pm_level`: `run.em` EM rounds of propagation and voting at one level. `img` holds the level's +/// image and `img_spare` is a free texture of the same size; the nearest-neighbour field starts +/// at `nnf`. Returns the level's filled image and its field. +fn pm_level(e: &mut Enc, lv: &Lv, run: PmRun, img: &GpuImage, img_spare: &GpuImage, nnf: &GpuImage) -> Option<(GpuImage, GpuImage)> { + let (seed, em, iters) = (run.seed, run.em, run.iters); + let (r, w, h) = (lv.r, lv.w, lv.h); + let field_a = e.scratch(w, h); + let field_b = e.scratch(w, h); + let hist = e.scratch(BINS as u32, 1); + let (img_a, img_b) = (img.clone(), img_spare.clone()); + let (mut field_nnf, mut nnf_spare) = (nnf.clone(), e.scratch(w, h)); + // The propagation and random search write the targets and the one-pixel ring around them, so + // every neighbour they read has a defined alpha (`grow(near, 1)`). The dispatch origin must + // be the ring's, not the targets': `over` sizes the dispatch from the box it is given. + let ring = grow(lv.near, 1, w, h); + let mut in_a = true; + for round in 0..em { + // The round's initial field: the inherited source where it is still valid, else random. + // Every pixel of the field's bounds is written, so the previous round's field is free. + let img = if in_a { &img_a } else { &img_b }; + let p = params([w, h, r, 0], [0; 4], [seed, 0, 0, round as u32], [0.0; 4]); + over(e, "imp_pm_init", &p, Pass { src: &field_nnf, aux: Some(img), out: &field_a, data: Some(&lv.flags), bb: lv.field }); + let mut f_in_a = true; + for it in 0..iters { + for parity in 0..2u32 { + let (s, o) = if f_in_a { (&field_a, &field_b) } else { (&field_b, &field_a) }; + let p = params([w, h, r, parity], [0; 4], [seed, it as u32, 0, round as u32], [0.0; 4]); + over(e, "imp_pm_prop", &p, Pass { src: s, aux: Some(img), out: o, data: Some(&lv.flags), bb: ring }); + f_in_a = !f_in_a; + } + let (s, o) = if f_in_a { (&field_a, &field_b) } else { (&field_b, &field_a) }; + let p = params([w, h, r, 0], [0; 4], [seed, it as u32, 0, round as u32], [0.0; 4]); + over(e, "imp_pm_rand", &p, Pass { src: s, aux: Some(img), out: o, data: Some(&lv.flags), bb: ring }); + f_in_a = !f_in_a; + } + let field = if f_in_a { &field_a } else { &field_b }; + // The vote's sigma²: the 75th percentile of the targets' match distances (`pm_level`). + let sigma2 = sigma2(e, field, lv, &hist)?; + let next_img = if in_a { &img_b } else { &img_a }; + e.copy_into(img, next_img, 0, 0); + let p = params([w, h, r, 0], [0; 4], [0; 4], [sigma2, 0.0, 0.0, 0.0]); + over(e, "imp_pm_vote", &p, Pass { src: field, aux: Some(img), out: next_img, data: Some(&lv.flags), bb: lv.holes }); + // Carry the round's sources into the nearest-neighbour field (targets only). + e.pixels("imp_nnf_store", ¶ms([w, h, 0, 0], [0; 4], [0; 4], [0.0; 4]), field, Some(&field_nnf), &nnf_spare, Some(&lv.flags)); + std::mem::swap(&mut field_nnf, &mut nnf_spare); + in_a = !in_a; + } + Some((if in_a { img_a } else { img_b }, field_nnf)) +} + +/// The 75th percentile of a 1024-bin distance histogram (`imp_hist`'s readback) covering +/// `0..range`, and the histogram's total. The bin's centre is the estimate: within half a bin +/// (`range/2048`) of the CPU's exact percentile. +fn percentile(bytes: &[u8], range: f32) -> (f32, f64) { + let bin = |i: usize| -> f64 { + match bytes.get(i * 16..i * 16 + 4) { + Some(b) => f32::from_le_bytes([b[0], b[1], b[2], b[3]]) as f64, + None => 0.0, + } + }; + let total: f64 = (0..BINS).map(bin).sum(); + if total < 1.0 { + return (1.0, total); + } + let rank = total * 3.0 / 4.0; + let mut cum = 0.0f64; + let mut at = BINS - 1; + for i in 0..BINS { + cum += bin(i); + if cum > rank { + at = i; + break; + } + } + let centre = ((at as f32) + 0.5) * range / BINS as f32; + (centre.max(1e-5), total) +} + +/// The largest finite distance `imp_hist` saw, from bin 0's `.g` channel. +fn max_distance(bytes: &[u8]) -> f32 { + match bytes.get(4..8) { + Some(b) => f32::from_le_bytes([b[0], b[1], b[2], b[3]]), + None => 0.0, + } +} + +/// `imp_hist` over `bb`, the targets' bounds: the 75th percentile of the finite distances and the +/// sample count. SDR input is binned over the fixed [`SDR_DISTANCE`] range; a frame with values +/// above 1 (32-bpc HDR) has larger distances, so the second dispatch bins over the observed +/// maximum instead. Clamping those to the SDR range would report a sigma² of ≈ 4 for a distance +/// of 1600, and `imp_pm_vote`'s weights would round to zero and keep the prior image. +/// +/// One 16 KB readback per EM round (≈ 28 per 1080p patch) on the SDR path: a few percent of the +/// stage, and the obvious next thing to remove if the patch has to be faster (a device-side +/// percentile would need a reduction instead). HDR frames pay one more dispatch and readback. +fn histogram(e: &mut Enc, field: &GpuImage, flags: &wgpu::Buffer, w: u32, h: u32, bb: [u32; 4], hist: &GpuImage) -> Option<(f32, f64)> { + let (bw, bh) = (bb[2] - bb[0], bb[3] - bb[1]); + let mut p = params([w, h, 0, 0], [bb[0], bb[1], bw, bh], [0; 4], [0.0; 4]); + e.dispatch("imp_hist", &p, KernelImages { src: field, aux: None, out: hist }, Some(flags), (1, 1)); + let bytes = e.read_texture(&hist.texture, BINS as u32, 1, 16)?; + let max = max_distance(&bytes); + if max <= SDR_DISTANCE { + return Some(percentile(&bytes, SDR_DISTANCE)); + } + // The observed maximum itself has to be in the last bin, so widen the range by a hair rather + // than binning it to 1024. + let range = max * (1.0 + 1e-6); + p.f[0] = [range, 0.0, 0.0, 0.0]; + e.dispatch("imp_hist", &p, KernelImages { src: field, aux: None, out: hist }, Some(flags), (1, 1)); + let bytes = e.read_texture(&hist.texture, BINS as u32, 1, 16)?; + Some(percentile(&bytes, range)) +} + +/// The vote's sigma² at a level (`pm_level`): the 75th percentile of the targets' distances. +fn sigma2(e: &mut Enc, field: &GpuImage, lv: &Lv, hist: &GpuImage) -> Option { + Some(histogram(e, field, &lv.flags, lv.w, lv.h, lv.near, hist)?.0) +} + +/// [`imp_hist`] on its own: `field`'s match distance (`.b`, `.a` > 0.5 where set) over the +/// pixels `near` marks, returning the 75th percentile and the sample count. Used by the test that +/// checks the histogram counts the whole target population. +#[cfg(test)] +pub(crate) fn distance_histogram(g: &GpuContext, field: &Image, near: &[bool], w: u32, h: u32) -> Option<(f32, f64)> { + let flags: Vec = near.iter().map(|b| if *b { 2.0 } else { 0.0 }).collect(); + let mut e = Enc::new(g); + let data = e.data(&flags); + let img = g.upload_image(field)?; + let hist = e.scratch(BINS as u32, 1); + let out = histogram(&mut e, &img, &data, w, h, [0, 0, w, h], &hist); + if e.take_missing() { + return None; + } + out +} diff --git a/crates/gpu/src/lib.rs b/crates/gpu/src/lib.rs index 1be095c3f..0844a4335 100644 --- a/crates/gpu/src/lib.rs +++ b/crates/gpu/src/lib.rs @@ -79,6 +79,7 @@ mod fx_tone; mod fx_transition; mod fx_vr; mod fx_warp; +mod inpaint; mod memcap; mod ops; mod particles; @@ -416,6 +417,17 @@ impl Accelerator for Gpu { self.ctx.can_readback().then_some(self as &dyn effectcraft_effects::psim::ParticleSim) } + fn inpaint(&self) -> Option<&dyn effectcraft_raster::inpaint::FillBackend> { + self.ctx.check_health().ok()?; + // Readback is synchronous-only, and a kernel that failed to build would make every stage + // fall back to the CPU *after* paying the upload and a failed dispatch (and the fill + // would still report the GPU path). Neither is a backend at all. + if !self.ctx.can_wait() || self.ctx.unavailable_kernels().iter().any(|(e, _)| crate::inpaint::KERNELS.contains(e)) { + return None; + } + Some(self as &dyn effectcraft_raster::inpaint::FillBackend) + } + fn auto_pick(&self) -> Option<&effectcraft_render::AutoPick> { Some(&self.auto) } diff --git a/crates/gpu/src/shaders/inpaint.wgsl b/crates/gpu/src/shaders/inpaint.wgsl new file mode 100644 index 000000000..56c468f85 --- /dev/null +++ b/crates/gpu/src/shaders/inpaint.wgsl @@ -0,0 +1,796 @@ +// GPU Content-Aware Fill: the membrane (harmonic) fill and PatchMatch synthesis of +// effectcraft_raster::inpaint, operation for operation (the CPU reference is the oracle; see +// crates/gpu/src/tests/inpaint.rs). +// +// Bindings are the shared ones (common.wgsl): `src` / `aux` are the input textures (an image, a +// coarser level, or the nearest-neighbour field), `out` is the storage texture and `data` is a +// per-pixel float table — the level's flags: +// * membrane kernels: data[i] = 1 where the pixel is known (the hole's complement); +// * PatchMatch kernels: data[i] = bit 0 hole, bit 1 near (a pixel whose patch touches the hole). +// +// Parameter block (P.u = u32, P.f = f32): +// membrane kernels: u[0] = (w, h, coarseW, coarseH), u[1] = (x0, y0, 0, 0) (dispatch +// origin), u[3].x = this level's offset into `data` +// down kernels: u[0] = (srcW, srcH, dstW, dstH) +// PatchMatch kernels: u[0] = (w, h, r, parity), u[1] = (x0, y0, 0, 0) (dispatch origin), +// u[2] = (seed, iteration, 0, EM round), f[0] = (sigma², 0, 0, 0) +// histogram: u[0] = (w, h, 0, 0), u[1] = (x0, y0, w, h) (target bounds), +// f[0].x = the distance range (0: the SDR range 0…4) +// The nearest-neighbour field texture holds one entry per pixel: r, g = source centre ("px"), +// b = SSD distance (inf where unset), a = 1 where the entry is set. + +const IMP_INF: f32 = 1e30; + +fn imp_flags(i: i32) -> u32 { + return u32(max(data[i], 0.0)); +} + +fn imp_hole(i: i32) -> bool { + return (imp_flags(i) & 1u) != 0u; +} + +fn imp_near(i: i32) -> bool { + return (imp_flags(i) & 2u) != 0u; +} + +// A source centre whose patch is fully known and inside the image (Pm::valid_source). +fn imp_valid_source(x: i32, y: i32, w: i32, h: i32, r: i32) -> bool { + if (x < r || y < r || x + r >= w || y + r >= h) { + return false; + } + return !imp_near(y * w + x); +} + +// PatchMatch's hashing rnd (raster::inpaint::rnd), bit for bit: the CPU's random stream is a +// contract of the algorithm, not an implementation detail. +fn imp_rnd(seed: u32, a: u32, b: u32) -> u32 { + var h = (seed * 0x9E3779B9u) ^ (a * 0x85EBCA6Bu) ^ (b * 0xC2B2AE35u); + h = h ^ (h >> 16u); + h = h * 0x7FEB352Du; + h = h ^ (h >> 15u); + h = h * 0x846CA68Bu; + return h ^ (h >> 16u); +} + +// SSD of the patch at target (tx, ty) (clipped to the image) against source (sx, sy) +// (Pm::dist). `cap` short-circuits a match that is already worse. +fn imp_dist(tx: i32, ty: i32, sx: i32, sy: i32, w: i32, h: i32, r: i32, cap: f32) -> f32 { + var d = 0.0; + var n = 0.0; + for (var dy = -r; dy <= r; dy = dy + 1) { + let ty2 = ty + dy; + if (ty2 < 0 || ty2 >= h) { + continue; + } + let sy2 = sy + dy; + for (var dx = -r; dx <= r; dx = dx + 1) { + let tx2 = tx + dx; + if (tx2 < 0 || tx2 >= w) { + continue; + } + let a = textureLoad(aux, vec2(tx2, ty2), 0); + let b = textureLoad(aux, vec2(sx + dx, sy2), 0); + let e = a - b; + d = d + e.x * e.x + e.y * e.y + e.z * e.z + e.w * e.w; + n = n + 1.0; + } + if (d > cap * max(n, 1.0)) { + return IMP_INF; + } + } + return d / max(n, 1.0); +} + +// ---------------------------------------------------------------- membrane + +// membrane_core's coarse level: the known-weighted 2×2 average. +@compute @workgroup_size(16, 16) +fn imp_down(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let cw = i32(P.u[0].z); + let ch = i32(P.u[0].w); + let cx = i32(gid.x); + let cy = i32(gid.y); + if (cx >= cw || cy >= ch) { + return; + } + var s = vec4(0.0); + var n = 0.0; + for (var dy = 0; dy < 2; dy = dy + 1) { + for (var dx = 0; dx < 2; dx = dx + 1) { + let fx = 2 * cx + dx; + let fy = 2 * cy + dy; + if (fx < w && fy < h && data[i32(P.u[3].x) + fy * w + fx] > 0.5) { + s = s + textureLoad(src, vec2(fx, fy), 0); + n = n + 1.0; + } + } + } + var v = vec4(0.0); + if (n > 0.0) { + v = s / n; + } + textureStore(out, vec2(cx, cy), v); +} + +// membrane_core's coarse-to-fine step: unknown pixels take the nearest coarse value. The fine +// image is copied into `out` first, so writing only the unknown pixels keeps the known ones. +// Dispatched over the unknown pixels' bounds (`P.u[1]`); the coarse level is valid everywhere. +@compute @workgroup_size(16, 16) +fn imp_up(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let cw = i32(P.u[0].z); + let ch = i32(P.u[0].w); + let x = i32(P.u[1].x) + i32(gid.x); + let y = i32(P.u[1].y) + i32(gid.y); + if (x >= w || y >= h || data[i32(P.u[3].x) + y * w + x] > 0.5) { + return; + } + let cx = min(x / 2, cw - 1); + let cy = min(y / 2, ch - 1); + textureStore(out, vec2(x, y), textureLoad(src, vec2(cx, cy), 0)); +} + +// One relaxation of the unknown pixels (membrane_core's Jacobi sweep). `src` is the current +// state and `aux` is unused. Dispatched over the unknown pixels' bounds (`P.u[1]`): known +// pixels never move, so nothing outside them needs writing. Both `src` and `out` are valid +// across the image (the fine level's, a copy of it, or the full downsample). +@compute @workgroup_size(16, 16) +fn imp_jacobi(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let x = i32(P.u[1].x) + i32(gid.x); + let y = i32(P.u[1].y) + i32(gid.y); + if (x >= w || y >= h) { + return; + } + var v = textureLoad(src, vec2(x, y), 0); + if (data[i32(P.u[3].x) + y * w + x] <= 0.5) { + var s = vec4(0.0); + var n = 0.0; + if (x > 0) { + s = s + textureLoad(src, vec2(x - 1, y), 0); + n = n + 1.0; + } + if (x + 1 < w) { + s = s + textureLoad(src, vec2(x + 1, y), 0); + n = n + 1.0; + } + if (y > 0) { + s = s + textureLoad(src, vec2(x, y - 1), 0); + n = n + 1.0; + } + if (y + 1 < h) { + s = s + textureLoad(src, vec2(x, y + 1), 0); + n = n + 1.0; + } + if (n > 0.0) { + v = s / n; + } + } + textureStore(out, vec2(x, y), v); +} + +// ---------------------------------------------------------------- PatchMatch + +// PatchMatch's pyramid downsample (downsample): the level's image, hole pixels excluded from the +// average, over `out`'s size. +@compute @workgroup_size(16, 16) +fn imp_pm_down(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let cw = i32(P.u[0].z); + let ch = i32(P.u[0].w); + let cx = i32(gid.x); + let cy = i32(gid.y); + if (cx >= cw || cy >= ch) { + return; + } + var s = vec4(0.0); + var n = 0.0; + for (var dy = 0; dy < 2; dy = dy + 1) { + for (var dx = 0; dx < 2; dx = dx + 1) { + let fx = min(2 * cx + dx, w - 1); + let fy = min(2 * cy + dy, h - 1); + if (!imp_hole(fy * w + fx)) { + s = s + textureLoad(src, vec2(fx, fy), 0); + n = n + 1.0; + } + } + } + var v = vec4(0.0); + if (n > 0.0) { + v = s / n; + } + textureStore(out, vec2(cx, cy), v); +} + +// PatchMatch's coarse-to-fine step: bilinear upsample of the coarse image into this level's hole +// (patch_fill's upsample loop). The fine image is copied into `out` first. +@compute @workgroup_size(16, 16) +fn imp_pm_up(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let cw = i32(P.u[0].z); + let ch = i32(P.u[0].w); + let x = i32(P.u[1].x) + i32(gid.x); + let y = i32(P.u[1].y) + i32(gid.y); + if (x >= w || y >= h || !imp_hole(y * w + x)) { + return; + } + let fx = clamp((f32(x) + 0.5) / 2.0 - 0.5, 0.0, f32(cw - 1)); + let fy = clamp((f32(y) + 0.5) / 2.0 - 0.5, 0.0, f32(ch - 1)); + let x0 = i32(floor(fx)); + let y0 = i32(floor(fy)); + let x1 = min(x0 + 1, cw - 1); + let y1 = min(y0 + 1, ch - 1); + let tx = fx - f32(x0); + let ty = fy - f32(y0); + let a = textureLoad(src, vec2(x0, y0), 0); + let b = textureLoad(src, vec2(x1, y0), 0); + let c = textureLoad(src, vec2(x0, y1), 0); + let d = textureLoad(src, vec2(x1, y1), 0); + let top = a + (b - a) * tx; + let bot = c + (d - c) * tx; + textureStore(out, vec2(x, y), top + (bot - top) * ty); +} + +// The NNF's coarse-to-fine step: ×2 (patch_fill's nnf upsample). +@compute @workgroup_size(16, 16) +fn imp_nnf_up(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let cw = i32(P.u[0].z); + let ch = i32(P.u[0].w); + let x = i32(gid.x); + let y = i32(gid.y); + if (x >= w || y >= h) { + return; + } + let cx = min(x / 2, cw - 1); + let cy = min(y / 2, ch - 1); + let c = textureLoad(src, vec2(cx, cy), 0); + var o = vec4(0.0); + if (c.a > 0.5) { + o = vec4(c.r * 2.0 + f32(x % 2), c.g * 2.0 + f32(y % 2), 0.0, 1.0); + } + textureStore(out, vec2(x, y), o); +} + +// A random valid source centre, by rejection (the CPU draws from its compacted source list; the +// distribution over the valid centres is the same, the stream differs from the list's order). +fn imp_any_source(seed: u32, i: u32, k: u32, w: i32, h: i32, r: i32) -> vec2 { + for (var n = 0u; n < 32u; n = n + 1u) { + let a = imp_rnd(seed, i, k + n); + let x = i32(a % u32(w)); + let y = i32((a >> 16u) % u32(h)); + if (imp_valid_source(x, y, w, h, r)) { + return vec2(x, y); + } + } + return vec2(-1, -1); +} + +// A round's initial field: the inherited source where it is still valid, a random one otherwise +// (pm_level's initialise/validate step). +@compute @workgroup_size(16, 16) +fn imp_pm_init(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let r = i32(P.u[0].z); + let x = i32(P.u[1].x) + i32(gid.x); + let y = i32(P.u[1].y) + i32(gid.y); + if (x >= w || y >= h) { + return; + } + let i = y * w + x; + var o = vec4(0.0, 0.0, IMP_INF, 0.0); + if (imp_near(i)) { + let prev = textureLoad(src, vec2(x, y), 0); + var s = vec2(i32(prev.r), i32(prev.g)); + if (prev.a <= 0.5 || !imp_valid_source(s.x, s.y, w, h, r)) { + // rnd(seed ^ e, i, 0), the round's random source (see imp_any_source). + s = imp_any_source(P.u[2].x ^ P.u[2].w, u32(i), 0u, w, h, r); + } + if (s.x >= 0) { + let d = imp_dist(x, y, s.x, s.y, w, h, r, IMP_INF); + o = vec4(f32(s.x), f32(s.y), d, 1.0); + } + } + textureStore(out, vec2(x, y), o); +} + +// One propagation pass: every target pixel tries its four neighbours' sources shifted by the +// step that would move that neighbour's match onto this pixel. The checkerboard parity in P.u[0].w +// makes the pass race-free (a pixel of one parity only reads the other's). +@compute @workgroup_size(16, 16) +fn imp_pm_prop(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let r = i32(P.u[0].z); + let x = i32(P.u[1].x) + i32(gid.x); + let y = i32(P.u[1].y) + i32(gid.y); + if (x >= w || y >= h) { + return; + } + let cur = textureLoad(src, vec2(x, y), 0); + let i = y * w + x; + if (!imp_near(i) || ((x + y) & 1) != i32(P.u[0].w)) { + textureStore(out, vec2(x, y), cur); + return; + } + var best = vec2(i32(cur.r), i32(cur.g)); + var bd = cur.b; + if (cur.a <= 0.5) { + bd = IMP_INF; + } + for (var k = 0; k < 4; k = k + 1) { + var nx = x; + var ny = y; + var sx = 0; + var sy = 0; + if (k == 0) { + nx = x - 1; + sx = 1; + } else if (k == 1) { + nx = x + 1; + sx = -1; + } else if (k == 2) { + ny = y - 1; + sy = 1; + } else { + ny = y + 1; + sy = -1; + } + if (nx < 0 || ny < 0 || nx >= w || ny >= h) { + continue; + } + let n = textureLoad(src, vec2(nx, ny), 0); + if (n.a <= 0.5) { + continue; + } + let c = vec2(i32(n.r) + sx, i32(n.g) + sy); + if (c.x == best.x && c.y == best.y) { + continue; + } + if (!imp_valid_source(c.x, c.y, w, h, r)) { + continue; + } + let d = imp_dist(x, y, c.x, c.y, w, h, r, bd); + if (d < bd) { + bd = d; + best = c; + } + } + // alpha = the entry is set: `bd` stays INF only if no match was found. + textureStore(out, vec2(x, y), vec4(f32(best.x), f32(best.y), bd, select(0.0, 1.0, bd < IMP_INF))); +} + +// The random search (halving radius) and one random source anywhere, after propagation. +@compute @workgroup_size(16, 16) +fn imp_pm_rand(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let r = i32(P.u[0].z); + let x = i32(P.u[1].x) + i32(gid.x); + let y = i32(P.u[1].y) + i32(gid.y); + if (x >= w || y >= h) { + return; + } + let cur = textureLoad(src, vec2(x, y), 0); + let i = y * w + x; + if (!imp_near(i) || cur.a <= 0.5) { + textureStore(out, vec2(x, y), cur); + return; + } + var best = vec2(i32(cur.r), i32(cur.g)); + var bd = cur.b; + let seed = P.u[2].x; + let it = P.u[2].y; + let em = P.u[2].w; + let iu = u32(i); + var rad = max(w, h); + var k = 0u; + while (rad >= 1) { + let a = imp_rnd(seed + it * 977u + em * 31u, iu, k); + let ox = i32(a % u32(2 * rad + 1)) - rad; + let oy = i32((a >> 16u) % u32(2 * rad + 1)) - rad; + let c = vec2(best.x + ox, best.y + oy); + if (imp_valid_source(c.x, c.y, w, h, r)) { + let d = imp_dist(x, y, c.x, c.y, w, h, r, bd); + if (d < bd) { + bd = d; + best = c; + } + } + rad = rad / 2; + k = k + 1u; + } + let any = imp_any_source(seed ^ 0xABCDu, iu, it + em * 64u, w, h, r); + if (any.x >= 0 && (any.x != best.x || any.y != best.y)) { + let d = imp_dist(x, y, any.x, any.y, w, h, r, bd); + if (d < bd) { + bd = d; + best = any; + } + } + textureStore(out, vec2(x, y), vec4(f32(best.x), f32(best.y), bd, select(0.0, 1.0, bd < IMP_INF))); +} + +// Carry a round's field back into the nearest-neighbour field (only targets, as on the CPU). +@compute @workgroup_size(16, 16) +fn imp_nnf_store(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let x = i32(gid.x); + let y = i32(gid.y); + if (x >= w || y >= h) { + return; + } + let i = y * w + x; + var o = textureLoad(aux, vec2(x, y), 0); + if (imp_near(i)) { + let f = textureLoad(src, vec2(x, y), 0); + o = vec4(f.r, f.g, 0.0, 1.0); + } + textureStore(out, vec2(x, y), o); +} + +// The vote: every hole pixel becomes the distance-weighted mean of the samples its overlapping +// patches carry (pm_level's vote step). `aux` is the current image; the image is copied into +// `out` first, so only hole pixels are written. +@compute @workgroup_size(16, 16) +fn imp_pm_vote(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let r = i32(P.u[0].z); + let x = i32(P.u[1].x) + i32(gid.x); + let y = i32(P.u[1].y) + i32(gid.y); + if (x >= w || y >= h) { + return; + } + let i = y * w + x; + if (!imp_hole(i)) { + return; + } + let sigma2 = P.f[0].x; + var acc = vec4(0.0); + var wsum = 0.0; + for (var dy = -r; dy <= r; dy = dy + 1) { + let py = y + dy; + if (py < 0 || py >= h) { + continue; + } + for (var dx = -r; dx <= r; dx = dx + 1) { + let px = x + dx; + if (px < 0 || px >= w) { + continue; + } + let f = textureLoad(src, vec2(px, py), 0); + if (f.a <= 0.5 || f.b >= IMP_INF) { + continue; + } + let s = vec2(i32(f.r) - dx, i32(f.g) - dy); + let wt = exp(-f.b / (2.0 * sigma2)); + acc = acc + textureLoad(aux, vec2(s.x, s.y), 0) * wt; + wsum = wsum + wt; + } + } + if (wsum > 0.0) { + textureStore(out, vec2(x, y), acc / wsum); + } else { + textureStore(out, vec2(x, y), textureLoad(aux, vec2(x, y), 0)); + } +} + +// The workgroup's 1024 distance bins (`imp_hist`); the name is unique across the concatenated +// shader sources because `var` is module scope. +var imp_hist_bins: array, 1024>; +// The largest finite distance this dispatch saw (as its f32 bit pattern: positive floats are +// ordered by their bits). `sigma2` uses it to learn whether the frame needs a wider range. +var imp_hist_max: atomic; + +// One workgroup's histogram of the targets' match distances, for the vote's sigma² (the CPU takes +// the 75th percentile of the finite distances). The single workgroup accumulates the whole +// population through the shared atomic bins above. +// +// `P.f[0].x` is the range the bins cover; 0 selects the SDR range 0…4, where an image whose +// channels are all ≤ 1 lives: each channel's squared difference is ≤ 4, so the mean SSD is too. +// Frames with values above 1 (32-bpc HDR) have larger distances, so `sigma2` re-dispatches with +// the observed maximum as the range instead of clamping them to bin 1023, which would report a +// sigma² of ≈ 4 for a distance of 1600 and round the vote's weights to zero. The observed maximum +// rides in bin 0's unused `.g` channel on every dispatch (its `.x` stays the count). +@compute @workgroup_size(256, 1, 1) +fn imp_hist(@builtin(global_invocation_id) gid: vec3) { + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let x0 = i32(P.u[1].x); + let y0 = i32(P.u[1].y); + let bw = i32(P.u[1].z); + let bh = i32(P.u[1].w); + let t = gid.x; + let range = select(4.0, P.f[0].x, P.f[0].x > 0.0); + let scale = 1024.0 / range; + if (t == 0u) { + atomicStore(&imp_hist_max, 0u); + } + // One workgroup, so the whole population is counted through shared atomics: a per-thread + // strided pass would keep only the pixels whose bin selected that thread (a biased ~1/256 + // subsample, correlated with the pixel index). + for (var i = t; i < 1024u; i = i + 256u) { + atomicStore(&imp_hist_bins[i], 0u); + } + workgroupBarrier(); + let total = u32(max(bw, 0) * max(bh, 0)); + for (var idx = t; idx < total; idx = idx + 256u) { + let x = x0 + i32(idx) % bw; + let y = y0 + i32(idx) / bw; + let i = y * w + x; + if (!imp_near(i)) { + continue; + } + let f = textureLoad(src, vec2(x, y), 0); + if (f.a <= 0.5 || f.b >= IMP_INF) { + continue; + } + atomicMax(&imp_hist_max, bitcast(f.b)); + // Clamp in f32 before the conversion: a distance far outside the range (a mismatched patch + // in an HDR frame) would overflow the i32 conversion, which is indeterminate in WGSL. + let b = u32(clamp(f.b * scale, 0.0, 1023.0)); + atomicAdd(&imp_hist_bins[b], 1u); + } + workgroupBarrier(); + // 256 threads × 4 bins cover the histogram. + for (var j = 0u; j < 4u; j = j + 1u) { + let bin = t * 4u + j; + textureStore(out, vec2(i32(bin), 0), vec4(f32(atomicLoad(&imp_hist_bins[bin])), 0.0, 0.0, 0.0)); + } + // Thread 0 is bin 0's only writer, so this store follows the one above. `.g` carries the + // observed maximum: host readbacks keep the row 1024 texels wide (a 16 KB, 256-aligned copy). + if (t == 0u) { + textureStore(out, vec2(0, 0), vec4(f32(atomicLoad(&imp_hist_bins[0])), bitcast(atomicLoad(&imp_hist_max)), 0.0, 0.0)); + } +} + +// ---------------------------------------------------------------- block flow (flow completion) + +// The block matcher of effectcraft_raster::flow::block_flow(a, b, 8, 6), which the fill uses to +// compute the motion between neighbouring membrane-filled frames before completing it inside the +// holes (`FillBackend::flow`). `src` packs the two frames' luma planes at the current level: +// .x is frame A's, .y is frame B's (`effectcraft_raster::flow::Plane::luma`). The kernels mirror +// the CPU operation for operation, so the field matches the CPU's (crates/gpu/src/tests/inpaint.rs). + +// Plane::half of the packed pair: the clamped 2×2 average of each channel. +@compute @workgroup_size(16, 16) +fn imp_fl_half(@builtin(global_invocation_id) gid: vec3) { + let cx = i32(gid.x); + let cy = i32(gid.y); + if (cx >= i32(P.u[0].z) || cy >= i32(P.u[0].w)) { + return; + } + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + var sx = 0.0; + var sy = 0.0; + for (var dy = 0; dy < 2; dy = dy + 1) { + for (var dx = 0; dx < 2; dx = dx + 1) { + let v = textureLoad(src, vec2(min(2 * cx + dx, w - 1), min(2 * cy + dy, h - 1)), 0); + sx = sx + v.x; + sy = sy + v.y; + } + } + textureStore(out, vec2(cx, cy), vec4(sx * 0.25, sy * 0.25, 0.0, 0.0)); +} + +// Plane::get: the packed luma at (x, y), clamped to the plane. +fn imp_fl_tex(x: i32, y: i32, w: i32, h: i32) -> vec2 { + let v = textureLoad(src, clamp(vec2(x, y), vec2(0, 0), vec2(w - 1, h - 1)), 0); + return vec2(v.x, v.y); +} + +// The block match's SAD with A's block already loaded (`la`, row-major): B's luma is the packed +// .y. The sub-pixel kernel's five fits share the window's block, so it is read once; the search +// hoists its own 4×4/8×8 block the same way. The loops' bounds are literals so the backend can +// unroll them and keep the loads in flight — a runtime-bound loop runs an order of magnitude +// slower. +fn imp_fl_sad8h(bx: i32, by: i32, dx: i32, dy: i32, w: i32, h: i32, la: array) -> f32 { + var s = 0.0; + for (var yy = 0; yy < 8; yy = yy + 1) { + for (var xx = 0; xx < 8; xx = xx + 1) { + s = s + abs(la[yy * 8 + xx] - imp_fl_tex(bx + xx + dx, by + yy + dy, w, h).y); + } + } + return s; +} + +// block_flow_window's coarse-to-fine search at one level: one integer vector per tile, seeded by +// the previous (coarser) level's vector at the tile and its four neighbours. `aux` is that field +// (`P.u[2].x` = 0 at the coarsest level, where the only seed is the zero vector). +@compute @workgroup_size(16, 16) +fn imp_fl_search(@builtin(global_invocation_id) gid: vec3) { + let cols = i32(P.u[0].z); + let rows = i32(P.u[0].w); + let c = i32(gid.x); + let r = i32(gid.y); + if (c >= cols || r >= rows) { + return; + } + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let bs = i32(P.u[1].x); + let k = i32(P.u[1].z); + let rad = i32(P.u[1].w); + let lbs = max(i32(P.u[1].y) / k, 4); + let bx = c * bs / k; + let by = r * bs / k; + // The zero vector plus up to five neighbour seeds (the tile and its four neighbours). + var seeds: array, 6>; + var ns = 0; + seeds[ns] = vec2(0, 0); + ns = ns + 1; + if (P.u[2].x != 0u) { + var nbr = array, 5>(vec2(0, 0), vec2(-1, 0), vec2(1, 0), vec2(0, -1), vec2(0, 1)); + for (var j = 0; j < 5; j = j + 1) { + let nc = c + nbr[j].x; + let nr = r + nbr[j].y; + if (nc >= 0 && nr >= 0 && nc < cols && nr < rows) { + let pv = textureLoad(aux, vec2(nc, nr), 0); + seeds[ns] = vec2(i32(round(pv.x)) * 2, i32(round(pv.y)) * 2); + ns = ns + 1; + } + } + } + // sort_unstable + dedup: distinct seeds in lexicographic (x, then y) order. + for (var i = 1; i < ns; i = i + 1) { + let key = seeds[i]; + var j = i - 1; + while (j >= 0 && (seeds[j].x > key.x || (seeds[j].x == key.x && seeds[j].y > key.y))) { + seeds[j + 1] = seeds[j]; + j = j - 1; + } + seeds[j + 1] = key; + } + var m = 0; + for (var i = 0; i < ns; i = i + 1) { + if (m == 0 || seeds[i].x != seeds[m - 1].x || seeds[i].y != seeds[m - 1].y) { + seeds[m] = seeds[i]; + m = m + 1; + } + } + ns = m; + // A's block is the same for every candidate, so load it once (the unrolled loops keep it in + // registers) and the search only reads B's shifted block. + var la4: array; + var la8: array; + if (lbs >= 8) { + for (var yy = 0; yy < 8; yy = yy + 1) { + for (var xx = 0; xx < 8; xx = xx + 1) { + la8[yy * 8 + xx] = imp_fl_tex(bx + xx, by + yy, w, h).x; + } + } + } else { + for (var yy = 0; yy < 4; yy = yy + 1) { + for (var xx = 0; xx < 4; xx = xx + 1) { + la4[yy * 4 + xx] = imp_fl_tex(bx + xx, by + yy, w, h).x; + } + } + } + var best = vec2(0, 0); + var best_e = 1e30; + for (var si = 0; si < ns; si = si + 1) { + let s = seeds[si]; + for (var dy = -rad; dy <= rad; dy = dy + 1) { + for (var dx = -rad; dx <= rad; dx = dx + 1) { + let d0 = s.x + dx; + let d1 = s.y + dy; + var acc = 0.0; + if (lbs >= 8) { + for (var yy = 0; yy < 8; yy = yy + 1) { + for (var xx = 0; xx < 8; xx = xx + 1) { + acc = acc + abs(la8[yy * 8 + xx] - imp_fl_tex(bx + xx + d0, by + yy + d1, w, h).y); + } + } + } else { + for (var yy = 0; yy < 4; yy = yy + 1) { + for (var xx = 0; xx < 4; xx = xx + 1) { + acc = acc + abs(la4[yy * 4 + xx] - imp_fl_tex(bx + xx + d0, by + yy + d1, w, h).y); + } + } + } + let e = acc + 1e-4 * f32(abs(d0) + abs(d1)); + if (e < best_e) { + best_e = e; + best = vec2(d0, d1); + } + } + } + } + textureStore(out, vec2(c, r), vec4(f32(best.x), f32(best.y), 0.0, 0.0)); +} + +// The parabola fit of the SAD around the best vector (block_flow_window's sub-pixel step). +fn imp_fl_fit(em: f32, ep: f32, e0: f32) -> f32 { + let den = em - 2.0 * e0 + ep; + if (den > 1e-6) { + return clamp(0.5 * (em - ep) / den, -0.5, 0.5); + } + return 0.0; +} + +// Sub-pixel refinement at full resolution: `aux` is the integer field of level 0, `src` the +// level-0 packed luma (`P.u[1]` = (tile, window, 0, 0)). +@compute @workgroup_size(16, 16) +fn imp_fl_subpix(@builtin(global_invocation_id) gid: vec3) { + let cols = i32(P.u[0].z); + let rows = i32(P.u[0].w); + let c = i32(gid.x); + let r = i32(gid.y); + if (c >= cols || r >= rows) { + return; + } + let w = i32(P.u[0].x); + let h = i32(P.u[0].y); + let bs = i32(P.u[1].x); + let win = i32(P.u[1].y); + let bx = c * bs; + let by = r * bs; + let f = textureLoad(aux, vec2(c, r), 0); + let d = vec2(i32(round(f.x)), i32(round(f.y))); + // The window's A block is the same for the five SADs: load it once (the loop is unrolled). + var la: array; + for (var yy = 0; yy < 8; yy = yy + 1) { + for (var xx = 0; xx < 8; xx = xx + 1) { + la[yy * 8 + xx] = imp_fl_tex(bx + xx, by + yy, w, h).x; + } + } + let e0 = imp_fl_sad8h(bx, by, d.x, d.y, w, h, la); + let sx = imp_fl_fit(imp_fl_sad8h(bx, by, d.x - 1, d.y, w, h, la), imp_fl_sad8h(bx, by, d.x + 1, d.y, w, h, la), e0); + let sy = imp_fl_fit(imp_fl_sad8h(bx, by, d.x, d.y - 1, w, h, la), imp_fl_sad8h(bx, by, d.x, d.y + 1, w, h, la), e0); + textureStore(out, vec2(c, r), vec4(f32(d.x) + sx, f32(d.y) + sy, 0.0, 0.0)); +} + +// The 3×3 component-wise vector median against outliers (each component sorted on its own). +@compute @workgroup_size(16, 16) +fn imp_fl_median(@builtin(global_invocation_id) gid: vec3) { + let cols = i32(P.u[0].z); + let rows = i32(P.u[0].w); + let c = i32(gid.x); + let r = i32(gid.y); + if (c >= cols || r >= rows) { + return; + } + var xs: array; + var ys: array; + var n = 0; + for (var dr = -1; dr <= 1; dr = dr + 1) { + for (var dc = -1; dc <= 1; dc = dc + 1) { + let nc = c + dc; + let nr = r + dr; + if (nc >= 0 && nr >= 0 && nc < cols && nr < rows) { + let v = textureLoad(src, vec2(nc, nr), 0); + xs[n] = v.x; + ys[n] = v.y; + n = n + 1; + } + } + } + for (var i = 1; i < n; i = i + 1) { + let kx = xs[i]; + var j = i - 1; + while (j >= 0 && xs[j] > kx) { + xs[j + 1] = xs[j]; + j = j - 1; + } + xs[j + 1] = kx; + let ky = ys[i]; + j = i - 1; + while (j >= 0 && ys[j] > ky) { + ys[j + 1] = ys[j]; + j = j - 1; + } + ys[j + 1] = ky; + } + textureStore(out, vec2(c, r), vec4(xs[n / 2], ys[n / 2], 0.0, 0.0)); +} diff --git a/crates/gpu/src/tests.rs b/crates/gpu/src/tests.rs index ef8673b5f..3f8e158f6 100644 --- a/crates/gpu/src/tests.rs +++ b/crates/gpu/src/tests.rs @@ -52,6 +52,7 @@ mod fx_vr; mod fx_warp; #[cfg(all(test, not(target_arch = "wasm32")))] mod init; +mod inpaint; mod particles; mod three_d; mod wgsl_grammar; diff --git a/crates/gpu/src/tests/inpaint.rs b/crates/gpu/src/tests/inpaint.rs new file mode 100644 index 000000000..e00d3cc18 --- /dev/null +++ b/crates/gpu/src/tests/inpaint.rs @@ -0,0 +1,310 @@ +//! GPU Content-Aware Fill vs the CPU reference: the CPU's fill is the oracle. +//! +//! Parity here is a quality contract, not bit-exactness (see `inpaint.rs`): PatchMatch's +//! scan-order propagation becomes a two-pass checkerboard. So the tests compare both +//! syntheses against the ground truth they are reconstructing, and check the GPU is not the +//! worse of the two. + +use effectcraft_raster::Image; +use effectcraft_raster::inpaint::{FillBackend, FillInput, FillOpts, fill_sequence_with, membrane_fill, patch_fill, psnr}; + +use crate::tests::gpu; + +/// A textured background (periodic stripes + blocky noise), shifted by `dx` (the CPU tests' +/// scene, so the numbers here are comparable with theirs). +fn background(w: u32, h: u32, dx: f32) -> Image { + let mut img = Image::new(w, h); + for y in 0..h { + for x in 0..w { + let fx = x as f32 + dx; + let s = 0.5 + 0.3 * (fx * std::f32::consts::TAU / 10.0).sin(); + let n = effectcraft_raster::hash_noise(((fx.floor() as i64).rem_euclid(20) / 4) as u32, (y % 20) / 4, 3); + img.set(x, y, [s, 0.3 + 0.4 * n, 0.6 - 0.3 * s * n, 1.0]); + } + } + img +} + +fn cut(img: &Image, hole: &[bool]) -> Image { + let mut o = img.clone(); + for (p, h) in o.data.iter_mut().zip(hole) { + if *h { + *p = [0.0; 4]; + } + } + o +} + +fn rect_hole(w: u32, h: u32, x0: u32, y0: u32, s: u32) -> Vec { + (0..w * h).map(|i| (x0..x0 + s).contains(&(i % w)) && (y0..y0 + s).contains(&(i / w))).collect() +} + +/// The kernels must exist: a shader that fails to build falls back to the CPU silently, which +/// would make every quality test below pass while testing nothing. +#[test] +fn the_inpaint_kernels_are_available() { + let Some(g) = gpu() else { return }; + let bad: Vec<&str> = g.ctx.unavailable_kernels().iter().filter(|(e, _)| crate::inpaint::KERNELS.contains(e)).map(|(e, _)| *e).collect(); + assert!(bad.is_empty(), "inpaint kernels did not build on {}: {bad:?}", g.ctx.name); +} + +/// The device path (not the CPU fallback) is what the quality tests measure. +#[test] +fn the_gpu_runs_the_fill_itself() { + let Some(g) = gpu() else { return }; + let (w, h) = (96u32, 72u32); + let gt = background(w, h, 0.0); + let hole = rect_hole(w, h, 40, 28, 16); + let src = cut(>, &hole); + assert!(crate::inpaint::patch_image(&g.ctx, &src, &hole, None, 7, 1).is_some(), "PatchMatch fell back to the CPU"); +} + +#[test] +fn patchmatch_beats_the_membrane_on_texture() { + let Some(g) = gpu() else { return }; + let (w, h) = (96u32, 72u32); + let gt = background(w, h, 0.0); + let hole = rect_hole(w, h, 40, 28, 16); + let src = cut(>, &hole); + let mem = membrane_fill(&src, &hole); + let pm = g.patch(&src, &hole, None, 7, 1); + let (pm_db, mem_db) = (psnr(&pm, >, Some(&hole)), psnr(&mem, >, Some(&hole))); + assert!(pm_db > mem_db + 3.0 && pm_db > 20.0, "PatchMatch {pm_db:.1} dB vs membrane {mem_db:.1} dB"); + for i in 0..hole.len() { + if !hole[i] { + assert_eq!(pm.data[i], gt.data[i]); + } + } +} + +#[test] +fn patchmatch_agrees_with_the_cpu() { + let Some(g) = gpu() else { return }; + let (w, h) = (96u32, 72u32); + let gt = background(w, h, 0.0); + let hole = rect_hole(w, h, 40, 28, 16); + let src = cut(>, &hole); + let (cpu, gpu_pm) = (patch_fill(&src, &hole, None, 7, 1), g.patch(&src, &hole, None, 7, 1)); + // Both are syntheses of the same hole: neither is the other's ground truth, so they are + // compared as two reconstructions of the same thing. On this periodic background both are + // near-exact, so a relative dB check at machine precision says nothing; the device must be a + // good reconstruction *and* agree with the CPU (a wrong fill that still looks plausible would + // pass the first and fail the second). + let gpu_db = psnr(&gpu_pm, >, Some(&hole)); + assert!(gpu_db > 25.0, "GPU {gpu_db:.1} dB"); + assert!(psnr(&gpu_pm, &cpu, Some(&hole)) > 30.0, "the two syntheses disagree"); +} + +/// The `init` path (a frame refined from the previous frame's warped fill) is the finest level +/// only, from two rounds of four iterations (`patch_fill`). +#[test] +fn patchmatch_refines_an_init_frame() { + let Some(g) = gpu() else { return }; + let (w, h) = (96u32, 72u32); + let gt = background(w, h, 0.0); + let hole = rect_hole(w, h, 40, 28, 16); + let src = cut(>, &hole); + let init = gt.clone(); + let (cpu, gpu_pm) = (patch_fill(&src, &hole, Some(&init), 7, 1), g.patch(&src, &hole, Some(&init), 7, 1)); + assert!(psnr(&gpu_pm, >, Some(&hole)) > 25.0, "{:.1} dB", psnr(&gpu_pm, >, Some(&hole))); + assert!(psnr(&gpu_pm, &cpu, Some(&hole)) > 20.0, "the two syntheses disagree"); +} + +/// A hole across the whole multi-scale pyramid (the coarse levels' membrane start included): +/// every level must come back finite and the pyramid must beat a single-level solve. +#[test] +fn patchmatch_survives_a_large_hole_across_levels() { + let Some(g) = gpu() else { return }; + let (w, h) = (256u32, 192u32); + let gt = background(w, h, 0.0); + let hole = rect_hole(w, h, 90, 70, 80); + let src = cut(>, &hole); + let pm = g.patch(&src, &hole, None, 7, 1); + assert!(pm.data.iter().flatten().all(|v| v.is_finite()), "the synthesis left a non-finite value"); + let db = psnr(&pm, >, Some(&hole)); + assert!(db > 20.0, "{db:.1} dB"); + for i in 0..hole.len() { + if !hole[i] { + assert_eq!(pm.data[i], gt.data[i], "known pixel {i}"); + } + } +} + +/// A moving object removed from a (static or panning) textured background, filled frame by +/// frame through the device backend the sequence driver uses. +#[test] +fn object_fill_recovers_background_and_is_stable() { + let Some(g) = gpu() else { return }; + let (w, h) = (96u32, 64u32); + for pan in [0.0f32, 1.0] { + let mut frames = vec![]; + let mut holes = vec![]; + let mut gts = vec![]; + for t in 0..8u32 { + let gt = background(w, h, t as f32 * pan); + let hole = rect_hole(w, h, 10 + 8 * t, 22, 18); + frames.push(cut(>, &hole)); + holes.push(hole); + gts.push(gt); + } + let input = FillInput { frames, holes, references: vec![] }; + let out = fill_sequence_with(&input, &FillOpts::default(), g, &mut |_, _| true).expect("fill"); + for t in 0..out.len() { + let db = psnr(&out[t], >s[t], Some(&input.holes[t])); + assert!(db > 26.0, "frame {t} pan {pan}: {db:.1} dB"); + assert_eq!(out[t].data[0][3], 1.0); + } + if pan == 0.0 { + for t in 1..out.len() { + let both: Vec = input.holes[t].iter().zip(&input.holes[t - 1]).map(|(a, b)| *a || *b).collect(); + let db = psnr(&out[t], &out[t - 1], Some(&both)); + assert!(db > 30.0, "frames {}–{t} differ: {db:.1} dB", t - 1); + } + } + } +} + +/// The membrane stays on the CPU (see [`crate::inpaint`]): Edge Blend must still get the CPU's +/// result through the backend, not a device one. +#[test] +fn edge_blend_takes_the_cpu_membrane() { + let Some(g) = gpu() else { return }; + let (w, h) = (80u32, 60u32); + let gt = background(w, h, 0.0); + let hole = rect_hole(w, h, 30, 20, 14); + let input = FillInput { frames: vec![cut(>, &hole); 4], holes: vec![hole.clone(); 4], references: vec![] }; + let opts = FillOpts { method: effectcraft_raster::inpaint::FillMethod::EdgeBlend, ..Default::default() }; + let edge = fill_sequence_with(&input, &opts, g, &mut |_, _| true).expect("fill"); + for t in 1..4 { + assert!(psnr(&edge[t], &edge[0], Some(&hole)) > 40.0, "frame {t}"); + } + assert!(psnr(&edge[0], &membrane_fill(&input.frames[0], &hole), Some(&hole)) > 60.0); +} + +/// The dilation the level plans are built from (kept honest: the flags' target bit is `near`). +#[test] +fn the_target_band_is_the_dilated_hole() { + let (w, h) = (24u32, 18u32); + let hole = rect_hole(w, h, 8, 6, 4); + let near = effectcraft_raster::inpaint::dilate(&hole, w as usize, h as usize, 3); + assert!(near.iter().filter(|b| **b).count() > hole.iter().filter(|b| **b).count()); + assert!(near.iter().zip(&hole).all(|(n, h)| !*h || *n)); +} + +/// The flow's block matching runs on the device and finds the same motion the CPU does: the two +/// frames are the membrane-filled neighbours the sequence driver builds, with a hole in each. +#[test] +fn flow_agrees_with_the_cpu() { + let Some(g) = gpu() else { return }; + let (w, h) = (192u32, 128u32); + let a_gt = background(w, h, 0.0); + let b_gt = background(w, h, 5.0); + let hole_a = rect_hole(w, h, 30, 40, 40); + let mut hole_b = vec![false; (w * h) as usize]; + for y in 50..90 { + for x in 120..160 { + hole_b[(y * w + x) as usize] = true; + } + } + let a = membrane_fill(&cut(&a_gt, &hole_a), &hole_a); + let b = membrane_fill(&cut(&b_gt, &hole_b), &hole_b); + assert!(crate::inpaint::flow_image(&g.ctx, &a, &b).is_some(), "the flow fell back to the CPU"); + let cpu = effectcraft_raster::inpaint::CpuFill.flow(&a, &b, &hole_a, &hole_b); + let dev = g.flow(&a, &b, &hole_a, &hole_b); + assert_eq!((dev.cols, dev.rows, dev.block), (cpu.cols, cpu.rows, cpu.block)); + // The field is the CPU's to a float rounding: the luma is the CPU's (`luma_pair`), the + // pyramid is exact additions, and the search enumerates candidates in the CPU's order. + let diff = cpu.v.iter().zip(&dev.v).map(|(c, d)| (c[0] - d[0]).abs().max((c[1] - d[1]).abs())).fold(0.0f32, f32::max); + assert!(diff <= 1e-6, "the flow fields differ by {diff}"); + // The motion is real: the panning background's 5 px shift survives the completion. `b`'s + // content sits 5 px left of `a`'s (`background` adds `dx` inside the sample), so a → b is −5. + let m = dev.at(96.0, 64.0); + assert!((m[0] + 5.0).abs() < 0.5 && m[1].abs() < 0.5, "flow at centre: {m:?}"); +} + +/// The vote's σ² comes from a histogram that must count *every* target: the old per-thread +/// strided pass kept only the pixels whose bin selected that thread (a biased ~1/256 subsample), +/// which left small holes with a `total < 1` fallback and near-uniform vote weights. +#[test] +fn the_distance_histogram_counts_every_target() { + let Some(g) = gpu() else { return }; + let (w, h) = (48u32, 32u32); + // Every target has the same distance, so the histogram is one spike. 0.5 is the SDR range + // (bin 128); 1600 is above it, which is reachable from a 32-bpc frame (the reviewer's example: + // clamping it to bin 1023 reported σ² ≈ 4, and the vote's weights then rounded to zero). + for (d, want) in [(0.5f32, 128.5f32 / 256.0), (1600.0, 1600.0)] { + let (mut field, near) = (Image::new(w, h), vec![true; (w * h) as usize]); + for px in field.data.iter_mut() { + *px = [0.0, 0.0, d, 1.0]; + } + let (sigma2, total) = crate::inpaint::distance_histogram(&g.ctx, &field, &near, w, h).expect("histogram"); + assert_eq!(total as usize, (w * h) as usize, "distance {d}: the histogram kept {total} of {} samples", w * h); + let tol = (want * 1e-3).max(1e-3); + assert!((sigma2 - want).abs() <= tol, "distance {d}: sigma2 {sigma2}, want {want}"); + } +} + +/// HDR frames (channels above 1, as a 32-bpc image has) must still fill like the CPU. Their patch +/// distances are far past the SDR 0…4 range the histogram used to clamp at, so this covers the +/// rescaled binning end to end; `the_distance_histogram_counts_every_target` pins the estimator +/// itself (the old clamp read σ² ≈ 4 for a distance of 1600). +#[test] +fn hdr_frames_fill_like_the_cpu() { + let Some(g) = gpu() else { return }; + let (w, h) = (96u32, 72u32); + let mut gt = background(w, h, 0.0); + for px in &mut gt.data { + // A 32-bpc frame's values are unbounded: at 24× the patch distances run into the thousands. + // The old fixed 0…4 histogram reported σ² ≈ 4 for those, `exp(-d/(2σ²))` rounded to zero, + // and the vote kept the prior image instead of synthesising. + for c in 0..3 { + px[c] *= 24.0; + } + } + let hole = rect_hole(w, h, 40, 28, 16); + let src = cut(>, &hole); + let (cpu, gpu_pm) = (patch_fill(&src, &hole, None, 7, 1), g.patch(&src, &hole, None, 7, 1)); + let gpu_db = psnr(&gpu_pm, >, Some(&hole)); + assert!(gpu_db > 25.0, "GPU {gpu_db:.1} dB"); + assert!(psnr(&gpu_pm, &cpu, Some(&hole)) > 30.0, "the two syntheses disagree"); +} + +/// A frame too small for one block is the CPU's zero field, on the device path too. +#[test] +fn small_frames_get_a_zero_flow() { + let Some(g) = gpu() else { return }; + let (w, h) = (6u32, 6u32); + let a = background(w, h, 0.0); + let b = background(w, h, 2.0); + let none = vec![false; (w * h) as usize]; + let f = g.flow(&a, &b, &none, &none); + assert_eq!((f.cols, f.rows), (1, 1)); + assert_eq!(f.v, vec![[0.0, 0.0]]); +} + +/// The flow the sequence driver uses (through the `Gpu` backend) keeps the fill stable: the same +/// frames filled with the CPU backend and the device backend agree, so the propagation sees the +/// same motion. +#[test] +fn sequence_fill_with_the_device_flow_matches_the_cpu() { + let Some(g) = gpu() else { return }; + let (w, h) = (96u32, 64u32); + let cpu = effectcraft_raster::inpaint::CpuFill; + for pan in [0.0f32, 2.0] { + let mut frames = vec![]; + let mut holes = vec![]; + for t in 0..6u32 { + let gt = background(w, h, t as f32 * pan); + let hole = rect_hole(w, h, 12 + 6 * t, 22, 18); + frames.push(cut(>, &hole)); + holes.push(hole); + } + let input = FillInput { frames, holes, references: vec![] }; + let dev = fill_sequence_with(&input, &FillOpts::default(), g, &mut |_, _| true).expect("gpu fill"); + let ref_cpu = fill_sequence_with(&input, &FillOpts::default(), &cpu, &mut |_, _| true).expect("cpu fill"); + for t in 0..dev.len() { + let db = psnr(&dev[t], &ref_cpu[t], Some(&input.holes[t])); + assert!(db > 40.0, "frame {t} pan {pan}: device vs CPU {db:.1} dB"); + } + } +} diff --git a/crates/raster/src/flow.rs b/crates/raster/src/flow.rs index feb033d65..180b235bf 100644 --- a/crates/raster/src/flow.rs +++ b/crates/raster/src/flow.rs @@ -51,6 +51,21 @@ impl Flow { } } +/// The `Plane::luma` of `a` and `b` packed into one image (`a` in the red channel, `b` in the +/// green; the rest zero). A device block matcher takes this as its input rather than computing +/// the luma itself: the weighted sum is not reproducible on a device, which may contract +/// `a·b + c` into a fused multiply-add, and the search's near-tie choices would then differ. +pub fn luma_pair(a: &Image, b: &Image) -> Image { + let luma = |p: &[f32; 4]| 0.2126 * p[0] + 0.7152 * p[1] + 0.0722 * p[2] + 0.25 * p[3]; + let mut out = Image::new(a.width, a.height); + out.data.par_iter_mut().enumerate().for_each(|(i, q)| { + let la = a.data.get(i).map(luma).unwrap_or(0.0); + let lb = b.data.get(i).map(luma).unwrap_or(0.0); + *q = [la, lb, 0.0, 0.0]; + }); + out +} + /// A single-channel plane. struct Plane { w: usize, diff --git a/crates/raster/src/inpaint.rs b/crates/raster/src/inpaint.rs index c64de3375..592ecd221 100644 --- a/crates/raster/src/inpaint.rs +++ b/crates/raster/src/inpaint.rs @@ -562,23 +562,31 @@ pub fn patch_fill(img: &Image, hole: &[bool], init: Option<&Image>, patch: usize // ------------------------------------------------------------------ flow-guided propagation -/// Block flow from `a` to `b` with the tiles that touch either hole replaced by the membrane -/// interpolation of the surrounding motion (flow completion). +/// Block flow `a` → `b` (tile 8, search radius 6) with the tiles that touch either hole replaced +/// by the membrane interpolation of the surrounding motion (flow completion). fn completed_flow(a: &Image, b: &Image, hole_a: &[bool], hole_b: &[bool]) -> Flow { - let mut f = block_flow(a, b, 8, 6); + let f = block_flow(a, b, 8, 6); + complete_flow(f, hole_a, hole_b, a.width as usize, a.height as usize) +} + +/// The completion half of `completed_flow`: the tiles of `f` (a `block_flow` of `w × h`) that +/// touch either hole become the harmonic interpolation of the surrounding motion. Split out so a +/// hardware backend can compute the block matching itself and hand the raw field here; the +/// completion is cheap (the field is a `cols × rows` grid) and stays on the CPU. +pub fn complete_flow(mut f: Flow, hole_a: &[bool], hole_b: &[bool], w: usize, h: usize) -> Flow { if f.v.is_empty() { return f; } - let (w, bs) = (a.width as usize, f.block as usize); + let bs = f.block as usize; let mut known = vec![true; f.v.len()]; for (i, k) in known.iter_mut().enumerate() { let (c, r) = (i % f.cols, i / f.cols); // A tile is unreliable when any pixel within it (and a block's margin) is a hole. let (x0, y0) = ((c * bs).saturating_sub(bs / 2), (r * bs).saturating_sub(bs / 2)); - let (x1, y1) = (((c + 1) * bs + bs / 2).min(w), ((r + 1) * bs + bs / 2).min(a.height as usize)); + let (x1, y1) = (((c + 1) * bs + bs / 2).min(w), ((r + 1) * bs + bs / 2).min(h)); 'tile: for y in y0..y1 { for x in x0..x1 { - if hole_a[y * w + x] || hole_b[y * w + x] { + if hole_a.get(y * w + x).copied().unwrap_or(true) || hole_b.get(y * w + x).copied().unwrap_or(true) { *k = false; break 'tile; } @@ -660,8 +668,54 @@ fn trace(context: &TraceContext<'_>, start: TraceStart) -> Option<([f32; 4], usi None } -/// Fill a range of frames. `progress(done, total)` returns `false` to cancel (→ `None`). +/// A hardware backend for the fill's stages (the GPU compositor implements it; see +/// `effectcraft_render::Accelerator::inpaint`). Every method has the CPU implementation as its +/// default, so a backend that cannot take a stage returns the CPU's work and the fill stays +/// correct — only slower. +/// +/// The orchestration ([`fill_sequence_with`]) stays on the CPU: the sequence's frame-to-frame +/// coupling is sequential, and only the per-frame and per-pair work is worth offloading. +/// +/// **A backend must be safe to call concurrently.** [`fill_sequence_with`] runs the stages of +/// independent frames and frame pairs in parallel (`rayon`), so `membrane`, `flow` and `patch` +/// may all be entered from several threads at once — a backend that owns a device must serialise +/// its use of it (the GPU compositor's per-call encoder does; a backend whose device cannot be +/// used from several threads must lock). +pub trait FillBackend: Send + Sync { + /// [`membrane_fill`]: the `hole` pixels take the harmonic interpolation of their border. + fn membrane(&self, img: &Image, hole: &[bool]) -> Image { + membrane_fill(img, hole) + } + /// `membrane_core`: harmonic interpolation of `data`'s unknown pixels in place (`known` is + /// true where the pixel is known). + fn membrane_known(&self, data: &mut [[f32; 4]], known: &[bool], w: usize, h: usize) { + membrane_core(data, known, w, h) + } + /// `completed_flow`: block flow `a` → `b` with the tiles touching either hole replaced by + /// the membrane interpolation of the surrounding motion. + fn flow(&self, a: &Image, b: &Image, hole_a: &[bool], hole_b: &[bool]) -> Flow { + completed_flow(a, b, hole_a, hole_b) + } + /// [`patch_fill`]: PatchMatch synthesis of `img`'s `hole` (`init`: a starting guess, e.g. the + /// previous frame's fill warped here). + fn patch(&self, img: &Image, hole: &[bool], init: Option<&Image>, patch: usize, seed: u32) -> Image { + patch_fill(img, hole, init, patch, seed) + } +} + +/// The CPU backend: every stage is the reference implementation. +pub struct CpuFill; + +impl FillBackend for CpuFill {} + +/// Fill a range of frames on the CPU. `progress(done, total)` returns `false` to cancel +/// (→ `None`). pub fn fill_sequence(input: &FillInput, o: &FillOpts, progress: &mut dyn FnMut(usize, usize) -> bool) -> Option> { + fill_sequence_with(input, o, &CpuFill, progress) +} + +/// [`fill_sequence`] with a hardware backend for the stages it can take (see [`FillBackend`]). +pub fn fill_sequence_with(input: &FillInput, o: &FillOpts, backend: &dyn FillBackend, progress: &mut dyn FnMut(usize, usize) -> bool) -> Option> { let n = input.frames.len(); if n == 0 { return Some(vec![]); @@ -672,7 +726,7 @@ pub fn fill_sequence(input: &FillInput, o: &FillOpts, progress: &mut dyn FnMut(u if o.method == FillMethod::EdgeBlend { let mut out = Vec::with_capacity(n); for (i, (f, hl)) in input.frames.iter().zip(&input.holes).enumerate() { - out.push(opaque_fill(membrane_fill(f, hl), hl)); + out.push(opaque_fill(backend.membrane(f, hl), hl)); if !progress(i + 1, n) { return None; } @@ -691,10 +745,10 @@ pub fn fill_sequence(input: &FillInput, o: &FillOpts, progress: &mut dyn FnMut(u } } // 1. Flows between membrane-filled neighbours, completed inside the holes. - let pre: Vec = frames.par_iter().zip(&holes).map(|(f, hl)| membrane_fill(f, hl)).collect(); + let pre: Vec = frames.par_iter().zip(&holes).map(|(f, hl)| backend.membrane(f, hl)).collect(); let pairs: Vec<(Flow, Flow)> = (0..n.saturating_sub(1)) .into_par_iter() - .map(|i| (completed_flow(&pre[i], &pre[i + 1], &holes[i], &holes[i + 1]), completed_flow(&pre[i + 1], &pre[i], &holes[i + 1], &holes[i]))) + .map(|i| (backend.flow(&pre[i], &pre[i + 1], &holes[i], &holes[i + 1]), backend.flow(&pre[i + 1], &pre[i], &holes[i + 1], &holes[i]))) .collect(); let (fwd, bwd): (Vec, Vec) = pairs.into_iter().unzip(); if !progress(1, total) { @@ -746,7 +800,7 @@ pub fn fill_sequence(input: &FillInput, o: &FillOpts, progress: &mut dyn FnMut(u diff[i] = std::array::from_fn(|c| if c < 3 { a[c] - v[c] } else { 0.0 }); } } - membrane_core(&mut diff, &dk, w, h); + backend.membrane_known(&mut diff, &dk, w, h); for (i, _) in &filled { for c in 0..3 { img.data[*i][c] += o.lighting * diff[*i][c]; @@ -770,7 +824,7 @@ pub fn fill_sequence(input: &FillInput, o: &FillOpts, progress: &mut dyn FnMut(u *slot = Some(propagated[t].0.clone()); } let mut done = 2; - out[t0] = Some(patch_fill(&propagated[t0].0, &propagated[t0].1, None, o.patch, o.seed)); + out[t0] = Some(backend.patch(&propagated[t0].0, &propagated[t0].1, None, o.patch, o.seed)); done += 1; if !progress(done, total) { return None; @@ -783,7 +837,7 @@ pub fn fill_sequence(input: &FillInput, o: &FillOpts, progress: &mut dyn FnMut(u // Warp the previous result here along the backward flow (t → t − 1). // Frames after `t0` are filled in order, so the previous one is set. let Some(prev) = out[t - 1].as_ref() else { - out[t] = Some(patch_fill(img, left, None, o.patch, o.seed)); + out[t] = Some(backend.patch(img, left, None, o.patch, o.seed)); continue; }; let mut warped = img.clone(); @@ -796,7 +850,7 @@ pub fn fill_sequence(input: &FillInput, o: &FillOpts, progress: &mut dyn FnMut(u } out[t] = Some(match o.method { FillMethod::Surface => warped, - _ => patch_fill(img, left, Some(&warped), o.patch, o.seed), + _ => backend.patch(img, left, Some(&warped), o.patch, o.seed), }); } done += 1; diff --git a/crates/render/src/accel.rs b/crates/render/src/accel.rs index 0d5fe1b6b..fe5a8f3fe 100644 --- a/crates/render/src/accel.rs +++ b/crates/render/src/accel.rs @@ -76,6 +76,16 @@ pub trait Accelerator: Send + Sync { fn particles(&self) -> Option<&dyn effectcraft_effects::psim::ParticleSim> { None } + /// A hardware backend for Content-Aware Fill's stages (the GPU compositor's kernels), with + /// the CPU fill's semantics. `None` = the fill runs on the CPU. + /// + /// The fill's input frames are read back to the CPU before it runs (analyses render on the + /// CPU by design, see `effectcraft_engine::background::offload::job_accel`), so a backend + /// here is about the fill's own work — PatchMatch and the flow's block matching today, one + /// stage at a time as each pays (`effectcraft_gpu::inpaint`) — not about compositing. + fn inpaint(&self) -> Option<&dyn effectcraft_raster::inpaint::FillBackend> { + None + } /// Timing history [`Backend::Auto`] uses to send each comp's top-level frames to the /// faster compositor. `None` = Auto always tries the accelerator first. fn auto_pick(&self) -> Option<&AutoPick> { diff --git a/docs/hardware-parity.md b/docs/hardware-parity.md index c25da4928..998fe4453 100644 --- a/docs/hardware-parity.md +++ b/docs/hardware-parity.md @@ -1,6 +1,6 @@ # Hardware parity -> **Last reviewed:** 2026-10-12 · **Last updated:** 2026-10-12 · **Change:** minor (closed issue references re-audited against `upstream/main` 470ea7a3) · **Target:** Adobe After Effects 2026 +> **Last reviewed:** 2026-10-29 · **Last updated:** 2026-10-29 · **Change:** minor (Content-Aware Fill's PatchMatch and flow search on the GPU) · **Target:** Adobe After Effects 2026 Each piece of hardware After Effects uses, per platform, against EffectCraft. Ours comes from the code at `origin/main` 1446e27 (`crates/gpu`, `crates/render`, `crates/ui-egui`) and the issue @@ -14,6 +14,7 @@ monitors are missing, and two GPU start-up failures are open. | Feature | After Effects 2026 | EffectCraft macOS | Windows | Linux | Web | Notes / evidence | |---|---|---|---|---|---|---| | GPU compositing and effects | Mercury GPU (Metal; CUDA / OpenCL / DirectX on Windows) | yes, Metal via wgpu | Vulkan yes; **Direct3D 12 fails** (FXC can't compile the `pointwise` kernel, #613) | Vulkan | WebGPU, in frame workers | 280 of 306 effects on the GPU, checked against the CPU renderer; Fractal and the analysis/tool effects stay on the CPU ([effects-parity.md](effects-parity.md)) | +| GPU analysis | GPU-accelerated Content-Aware Fill | Content-Aware Fill's PatchMatch **and** flow block matching on the GPU (`Gpu::inpaint`); every other stage on the CPU | same | same | no (needs synchronous readback) | On an Intel Arc iGPU a four-frame 1920 × 1080 fill is ≈ 1.4× faster than the CPU's (PatchMatch is the dominant win, the flow search a smaller one), and even a 16 × 16 hole at 1920 × 1080 is ≈ 1.9× faster — the CPU's own stages are O(frame), not O(hole). The device's flow field matches `block_flow` to a float rounding. The membrane and the flow's hole completion stay on the CPU: both are bandwidth-bound, and the device's full-frame transfers cost more than they save — see `effectcraft_gpu::inpaint` | | CPU fallback and device-loss recovery | yes | yes | yes; Windows ARM64 Adreno panic after device loss (#591, closed) | yes | yes (#225 fixed) | Out of video memory falls back to the CPU (#106) | | GPU detection and choice | GPU info, renderer choice | Backend Auto per comp | "GPU isn't being detected" (#596, closed); Intel HD 5500 fix unverified (#198) | yes | yes | No settings page listing adapters yet | | Multi-core rendering | Multi-Frame Rendering | rayon pools; preview scheduler sized to the render pool | same | same | Web Workers, one engine per worker | No benchmark against Multi-Frame Rendering | @@ -39,5 +40,6 @@ the Adreno and Intel cases. | Date | Change | Summary | |---|---|---| +| 2026-10-29 | minor | Content-Aware Fill's PatchMatch and flow block matching run on the GPU (`FillBackend`/`Accelerator::inpaint`, `effectcraft_gpu::inpaint`) | | 2026-10-12 | minor | #591, #596 and #554 marked closed; #498 answer and open PR #819 | | 2026-10-10 | major | First hardware parity document |