mirror of
https://github.com/storytold/effectcraft.git
synced 2026-11-05 18:18:49 +00:00
gpu: run Content-Aware Fill's PatchMatch and flow on the GPU (#1856)
* gpu: run Content-Aware Fill's PatchMatch on the device Content-Aware Fill's synthesis is the stage that makes the tool slow: at 1920×1080 with a 300×200 hole, PatchMatch costs ≈ 1.1 s per frame on the CPU, and it is what Object, Surface and every refining frame spend their time in. Add a `FillBackend` trait to `effectcraft-raster` whose stages (the membrane solve, the flow fields, PatchMatch) default to the CPU reference, and `Accelerator::inpaint()` on the compositor, which is how the engine asks for a device backend. `effectcraft_gpu::inpaint` implements PatchMatch: the multi-scale pyramid, the EM rounds, the propagation (a two-pass checkerboard instead of the CPU's scan order), the random search, the SSD, the histogram-based σ² and the weighted vote, as compute kernels in `shaders/inpaint.wgsl`. The fill's frames are read back to the CPU before it runs (analyses composite on the CPU by design), so the backend is about the fill's own work; the orchestration, the flow-guided propagation and the per-frame flow control stay on the CPU. Only the hole's bounds cross back to the CPU, which is what makes the offload pay. Measured on an Intel Arc (MTL) iGPU, 1920×1080, 300×200 hole: CPU 1087 ms GPU 397 ms (2.7×, 41.0 dB PSNR against the CPU's synthesis) The membrane (harmonic) solve deliberately stays on the CPU: the same measurement has it at ≈ 90 ms on the device against ≈ 55 ms on the CPU, because a fill's full-frame data movement (upload, the known pyramid, a copy) does not shrink with the hole, while the relaxation does. The device takes only the stage where its parallel search pays; `membrane_tex` still runs on the device inside PatchMatch, where it is the coarsest level's start and never crosses back. Closes #1839. * gpu: fix the fill's vote histogram and run its flow on the device Review follow-ups on the Content-Aware Fill GPU port, plus the fill's flow. Fixes: * imp_hist kept about one sample in 256: each of the 256 threads scanned a strided subset of the pixels and retained only those whose bin selected it, so the vote's sigma² came from a sample correlated with the pixel index (and a small hole usually fell back to 1.0, making the weights near-uniform). One workgroup now accumulates the whole population through shared atomic bins, with a test that pins the sample count. * imp_pm_prop and imp_pm_rand were dispatched over grow(near, 1) but with near's origin, so the ring's left column and top row were never written (NaN under test, stale pooled data in production) while a strip past it was. `over` now derives the shader's origin from the box it sizes the dispatch from, so a dispatch's region and its origin can no longer disagree. * Gpu::inpaint() requires its kernels to have built. Otherwise every stage paid an upload and a failed dispatch to fall back to the CPU, while the fill still reported the GPU path. * imp_rnd's XOR chain parenthesises its multiplicative terms: WGSL's grammar makes each XOR operand a unary expression, so `a * b ^ c * d` is rejected by spec-compliant parsers (naga accepts it). * FillBackend documents that its stages may be called concurrently. * The coarsest membrane's sweep count (64, against the CPU's 200) is now justified: matching 200 measured ~6% slower over the fill with no change in the result. And the flow: Gpu::flow runs block_flow's coarse-to-fine search, sub-pixel fit and vector median on the device. The luma is the CPU's (flow::luma_pair), so the field matches block_flow's to a float rounding and the propagation sees the same motion; the hole completion stays on the CPU with the membrane. A four-frame 1920x1080 fill goes from 1.27x to 1.42x against the CPU, and the device path is covered for holes from 16x16 to 300x200. * gpu: keep the fill vote's histogram range on HDR input imp_hist binned every distance at a fixed 256 bins per unit and clamped to 1023, so the vote's sigma2 topped out at 3.998. A 32-bpc frame's channels can exceed 1, and masked_frame forwards them unclamped: a patch distance of 1600 then reported sigma2 ~ 4, exp(-1600/(2*sigma2)) rounded to zero in f32, and imp_pm_vote kept the prior image instead of synthesising. The histogram now reports the largest finite distance it saw (atomicMax on the float bits, carried in bin 0's unused .g channel). SDR input keeps the fixed 0..4 range with identical bins; HDR input is re-binned over the observed range and percentile() maps the bin centre back through it. Nothing clamps the input. Also group over/sweeps/pm_level's arguments into small structs: the workspace denies too_many_arguments and the lint policy (#2496) does not allow new #[allow]s. No behaviour change on those paths. * engine: don't report Edge Blend fills as device fills Edge Blend calls only the membrane, which the GPU backend leaves to the CPU reference, so run_plan's "Filling (GPU)" status was wrong for it whenever a backend was present. Object and Surface still run the flow and PatchMatch on the device and keep the label. CodeRabbit's #1856 finding. * gpu: run the fill's device stages inside error scopes An allocation or validation failure inside patch_image/flow_image left an invalid texture whose later use raised an uncaptured error, retiring the shared compositor instead of falling back to the CPU for that frame. Wrap both stages in Gpu::within_memory, whose scopes are native thread-local, so the fill's rayon workers do not share them. --------- Co-authored-by: FredApps <153173854+FredApps@users.noreply.github.com>
This commit is contained in:
gmrmad1andFredApps authored and 
GitHub committed
1 parent
6df7eabe05
commit
389894e69b
12 files changed
+1999
-24
No files matched your search
@@ -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 |_| {},
|
||||
)
|
||||
|
||||
@@ -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<Str
|
||||
})
|
||||
}
|
||||
|
||||
/// The status line shown while the frames are filled. Edge Blend calls only the membrane, which
|
||||
/// the device backend leaves to the CPU reference, so it is not a device fill even when a `fill`
|
||||
/// backend is present; Object and Surface run the flow and PatchMatch on the device.
|
||||
fn filling_status(on_device: bool, method: FillMethod) -> &'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<Vec<String>, 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<Value> {
|
||||
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<Value> {
|
||||
&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<CommandSpec> {
|
||||
|
||||
#[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]
|
||||
|
||||
@@ -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<Vec<u8>> {
|
||||
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<Vec<u8>> {
|
||||
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 },
|
||||
);
|
||||
|
||||
@@ -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<GpuImage> {
|
||||
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<f32>> = 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<f32> = Vec::with_capacity(knowns.iter().map(|k| k.len()).sum());
|
||||
let mut offs: Vec<u32> = 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<Option<[u32; 4]>> = 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<GpuImage> = 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<Flow> {
|
||||
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<GpuImage> = 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<GpuImage> = 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<Image> {
|
||||
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<bool> = 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<GpuImage> = 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<GpuImage> = 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<Lv> {
|
||||
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<f32> = (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<f32> {
|
||||
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<f32> = 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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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<i32>(tx2, ty2), 0);
|
||||
let b = textureLoad(aux, vec2<i32>(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<u32>) {
|
||||
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<f32>(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<i32>(fx, fy), 0);
|
||||
n = n + 1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
var v = vec4<f32>(0.0);
|
||||
if (n > 0.0) {
|
||||
v = s / n;
|
||||
}
|
||||
textureStore(out, vec2<i32>(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<u32>) {
|
||||
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<i32>(x, y), textureLoad(src, vec2<i32>(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<u32>) {
|
||||
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<i32>(x, y), 0);
|
||||
if (data[i32(P.u[3].x) + y * w + x] <= 0.5) {
|
||||
var s = vec4<f32>(0.0);
|
||||
var n = 0.0;
|
||||
if (x > 0) {
|
||||
s = s + textureLoad(src, vec2<i32>(x - 1, y), 0);
|
||||
n = n + 1.0;
|
||||
}
|
||||
if (x + 1 < w) {
|
||||
s = s + textureLoad(src, vec2<i32>(x + 1, y), 0);
|
||||
n = n + 1.0;
|
||||
}
|
||||
if (y > 0) {
|
||||
s = s + textureLoad(src, vec2<i32>(x, y - 1), 0);
|
||||
n = n + 1.0;
|
||||
}
|
||||
if (y + 1 < h) {
|
||||
s = s + textureLoad(src, vec2<i32>(x, y + 1), 0);
|
||||
n = n + 1.0;
|
||||
}
|
||||
if (n > 0.0) {
|
||||
v = s / n;
|
||||
}
|
||||
}
|
||||
textureStore(out, vec2<i32>(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<u32>) {
|
||||
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<f32>(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<i32>(fx, fy), 0);
|
||||
n = n + 1.0;
|
||||
}
|
||||
}
|
||||
}
|
||||
var v = vec4<f32>(0.0);
|
||||
if (n > 0.0) {
|
||||
v = s / n;
|
||||
}
|
||||
textureStore(out, vec2<i32>(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<u32>) {
|
||||
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<i32>(x0, y0), 0);
|
||||
let b = textureLoad(src, vec2<i32>(x1, y0), 0);
|
||||
let c = textureLoad(src, vec2<i32>(x0, y1), 0);
|
||||
let d = textureLoad(src, vec2<i32>(x1, y1), 0);
|
||||
let top = a + (b - a) * tx;
|
||||
let bot = c + (d - c) * tx;
|
||||
textureStore(out, vec2<i32>(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<u32>) {
|
||||
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<i32>(cx, cy), 0);
|
||||
var o = vec4<f32>(0.0);
|
||||
if (c.a > 0.5) {
|
||||
o = vec4<f32>(c.r * 2.0 + f32(x % 2), c.g * 2.0 + f32(y % 2), 0.0, 1.0);
|
||||
}
|
||||
textureStore(out, vec2<i32>(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<i32> {
|
||||
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<i32>(x, y);
|
||||
}
|
||||
}
|
||||
return vec2<i32>(-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<u32>) {
|
||||
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<f32>(0.0, 0.0, IMP_INF, 0.0);
|
||||
if (imp_near(i)) {
|
||||
let prev = textureLoad(src, vec2<i32>(x, y), 0);
|
||||
var s = vec2<i32>(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>(f32(s.x), f32(s.y), d, 1.0);
|
||||
}
|
||||
}
|
||||
textureStore(out, vec2<i32>(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<u32>) {
|
||||
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<i32>(x, y), 0);
|
||||
let i = y * w + x;
|
||||
if (!imp_near(i) || ((x + y) & 1) != i32(P.u[0].w)) {
|
||||
textureStore(out, vec2<i32>(x, y), cur);
|
||||
return;
|
||||
}
|
||||
var best = vec2<i32>(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<i32>(nx, ny), 0);
|
||||
if (n.a <= 0.5) {
|
||||
continue;
|
||||
}
|
||||
let c = vec2<i32>(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<i32>(x, y), vec4<f32>(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<u32>) {
|
||||
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<i32>(x, y), 0);
|
||||
let i = y * w + x;
|
||||
if (!imp_near(i) || cur.a <= 0.5) {
|
||||
textureStore(out, vec2<i32>(x, y), cur);
|
||||
return;
|
||||
}
|
||||
var best = vec2<i32>(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<i32>(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<i32>(x, y), vec4<f32>(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<u32>) {
|
||||
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<i32>(x, y), 0);
|
||||
if (imp_near(i)) {
|
||||
let f = textureLoad(src, vec2<i32>(x, y), 0);
|
||||
o = vec4<f32>(f.r, f.g, 0.0, 1.0);
|
||||
}
|
||||
textureStore(out, vec2<i32>(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<u32>) {
|
||||
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<f32>(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<i32>(px, py), 0);
|
||||
if (f.a <= 0.5 || f.b >= IMP_INF) {
|
||||
continue;
|
||||
}
|
||||
let s = vec2<i32>(i32(f.r) - dx, i32(f.g) - dy);
|
||||
let wt = exp(-f.b / (2.0 * sigma2));
|
||||
acc = acc + textureLoad(aux, vec2<i32>(s.x, s.y), 0) * wt;
|
||||
wsum = wsum + wt;
|
||||
}
|
||||
}
|
||||
if (wsum > 0.0) {
|
||||
textureStore(out, vec2<i32>(x, y), acc / wsum);
|
||||
} else {
|
||||
textureStore(out, vec2<i32>(x, y), textureLoad(aux, vec2<i32>(x, y), 0));
|
||||
}
|
||||
}
|
||||
|
||||
// The workgroup's 1024 distance bins (`imp_hist`); the name is unique across the concatenated
|
||||
// shader sources because `var<workgroup>` is module scope.
|
||||
var<workgroup> imp_hist_bins: array<atomic<u32>, 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<workgroup> imp_hist_max: atomic<u32>;
|
||||
|
||||
// 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<u32>) {
|
||||
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<i32>(x, y), 0);
|
||||
if (f.a <= 0.5 || f.b >= IMP_INF) {
|
||||
continue;
|
||||
}
|
||||
atomicMax(&imp_hist_max, bitcast<u32>(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>(i32(bin), 0), vec4<f32>(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<i32>(0, 0), vec4<f32>(f32(atomicLoad(&imp_hist_bins[0])), bitcast<f32>(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<u32>) {
|
||||
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<i32>(min(2 * cx + dx, w - 1), min(2 * cy + dy, h - 1)), 0);
|
||||
sx = sx + v.x;
|
||||
sy = sy + v.y;
|
||||
}
|
||||
}
|
||||
textureStore(out, vec2<i32>(cx, cy), vec4<f32>(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<f32> {
|
||||
let v = textureLoad(src, clamp(vec2<i32>(x, y), vec2<i32>(0, 0), vec2<i32>(w - 1, h - 1)), 0);
|
||||
return vec2<f32>(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, 64>) -> 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<u32>) {
|
||||
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<vec2<i32>, 6>;
|
||||
var ns = 0;
|
||||
seeds[ns] = vec2<i32>(0, 0);
|
||||
ns = ns + 1;
|
||||
if (P.u[2].x != 0u) {
|
||||
var nbr = array<vec2<i32>, 5>(vec2<i32>(0, 0), vec2<i32>(-1, 0), vec2<i32>(1, 0), vec2<i32>(0, -1), vec2<i32>(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<i32>(nc, nr), 0);
|
||||
seeds[ns] = vec2<i32>(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<f32, 16>;
|
||||
var la8: array<f32, 64>;
|
||||
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<i32>(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<i32>(d0, d1);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
textureStore(out, vec2<i32>(c, r), vec4<f32>(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<u32>) {
|
||||
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<i32>(c, r), 0);
|
||||
let d = vec2<i32>(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<f32, 64>;
|
||||
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<i32>(c, r), vec4<f32>(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<u32>) {
|
||||
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<f32, 9>;
|
||||
var ys: array<f32, 9>;
|
||||
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<i32>(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<i32>(c, r), vec4<f32>(xs[n / 2], ys[n / 2], 0.0, 0.0));
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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<bool> {
|
||||
(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<bool> = 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");
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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<Vec<Image>> {
|
||||
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<Vec<Image>> {
|
||||
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<Image> = frames.par_iter().zip(&holes).map(|(f, hl)| membrane_fill(f, hl)).collect();
|
||||
let pre: Vec<Image> = 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<Flow>, Vec<Flow>) = 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;
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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 |
|
||||
Reference in new issue
Block a user