diff --git a/core/dr-denoise/examples/denoise_raw.rs b/core/dr-denoise/examples/denoise_raw.rs index 123b63c..3e62164 100644 --- a/core/dr-denoise/examples/denoise_raw.rs +++ b/core/dr-denoise/examples/denoise_raw.rs @@ -11,7 +11,9 @@ //! camera RGB — for comparison with the training repo's own path //! (`tools/compare_rust.py` in darkroom-denoise). `DARKROOM_ORT_DIR` points //! at an ONNX Runtime build; the engine's cache goes to `DR_ENGINE_CACHE` or -//! a temporary directory. +//! a temporary directory. The whole-frame network (`mosaic-best.onnx` beside +//! the fixed file) runs where the rung takes any size; `DR_PLAN=tiles` keeps +//! the 1408² tiles anyway, to compare the two. use std::path::PathBuf; use std::time::{Duration, Instant}; @@ -38,6 +40,12 @@ fn main() { let model = PathBuf::from(env!("CARGO_MANIFEST_DIR")) .join("../../models/denoise") .join(shipped.file); + let whole = model.with_file_name(shipped.whole); + let tiles_only = std::env::var("DR_PLAN").is_ok_and(|p| p == "tiles"); + let mut models = vec![(Role::Denoiser, model.clone())]; + if whole.is_file() && !tiles_only { + models.push((Role::WholeDenoiser, whole)); + } let cache = std::env::var_os("DR_ENGINE_CACHE") .map(PathBuf::from) .unwrap_or_else(|| std::env::temp_dir().join("dr-denoise-engines")); @@ -48,7 +56,7 @@ fn main() { .into_iter() .collect(), cache_dir: cache, - models: vec![(Role::Denoiser, model.clone())], + models, embedded: Vec::new(), ceiling: None, threads: 0, @@ -101,10 +109,20 @@ fn main() { noise.col ); - let mut net = OnnxNet::from_path(&model, shipped.halo).expect("model"); + let mut net = if tiles_only { + OnnxNet::open_tiled(&model, shipped) + } else { + OnnxNet::open(&model, shipped) + } + .expect("model"); println!( - "rung {}", - net.rung().map(|r| r.label()).unwrap_or("?") + "rung {} · {}", + net.rung().map(|r| r.label()).unwrap_or("?"), + if net.whole_frame() { + "whole frame" + } else { + "1408² tiles" + } ); let t = Instant::now(); let rgb = dr_denoise::denoise(&raw, &noise, &mut net, &mut |done, total| { diff --git a/core/dr-denoise/src/lib.rs b/core/dr-denoise/src/lib.rs index eb78e76..a977826 100644 --- a/core/dr-denoise/src/lib.rs +++ b/core/dr-denoise/src/lib.rs @@ -25,32 +25,37 @@ pub mod tile; use dr_decode::RawImage; pub use noise::{NoiseModel, Source}; -pub use tile::{TileNet, HALO}; +pub use tile::{Sizes, TileNet, HALO}; /// TRACES: FR-DEV-3g -/// A network the app ships in `models/denoise/`: its file, and the context -/// it needs past a tile's kept centre (docs/dev/denoise.md §13). +/// A network the app ships in `models/denoise/`: its fixed-tile file, the +/// same network with any height and width for a whole frame (§14), and the +/// context it needs past a tile's kept centre (§13). #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct Shipped { pub file: &'static str, + pub whole: &'static str, pub halo: usize, } /// The smallest student: 0.9 M parameters, 11 GMAC a megapixel. pub const FAST: Shipped = Shipped { file: "mosaic-fast-1408.onnx", + whole: "mosaic-fast.onnx", halo: HALO, }; /// A student of the mixture with the first release's shape: 3.2 M /// parameters, 48 GMAC a megapixel. pub const MEDIUM: Shipped = Shipped { file: "mosaic-medium-1408.onnx", + whole: "mosaic-medium.onnx", halo: HALO, }; /// The mixture: a flat expert, an edge expert and the gate that blends them. /// It reaches further than either, so it keeps a smaller centre of each tile. pub const BEST: Shipped = Shipped { file: "mosaic-best-1408.onnx", + whole: "mosaic-best.onnx", halo: 256, }; diff --git a/core/dr-denoise/src/onnx.rs b/core/dr-denoise/src/onnx.rs index 748478f..f7d0e14 100644 --- a/core/dr-denoise/src/onnx.rs +++ b/core/dr-denoise/src/onnx.rs @@ -10,30 +10,80 @@ //! lost 5–9 dB (docs/dev/inference.md §1.5). That sibling is the same network //! with the Bayer packing spelled `SpaceToDepth`, which QNN can hold and the //! 6-D reshape it replaces it cannot. +//! +//! Each network also ships with any height and width (`mosaic-best.onnx` +//! beside `mosaic-best-1408.onnx`, darkroom-denoise `tools/export_whole.py`, +//! identical to the fixed file at 1408²). Where the rung takes any size, the +//! frame runs whole instead of in tiles whose borders are thrown away — +//! Best computes 2.47 photosites for every one it keeps in 1408² tiles +//! (denoise.md §14). -use crate::tile::TileNet; -use crate::DenoiseError; -use dr_inference_engine::{Model, Role}; +use crate::tile::{Sizes, TileNet}; +use crate::{DenoiseError, Shipped}; +use dr_inference_engine::{Form, Model, Role}; -/// The edge of the tile the shipped export takes. +/// The edge of the tile the shipped fixed-shape export takes. pub const TILE: usize = 1408; +/// What a whole-frame input's sides must be multiples of: the networks pack +/// 2×2 and halve three times, so a side is a whole number of positions at +/// their coarsest level only in steps of 16. +pub const ALIGN: usize = 16; + pub struct OnnxNet { model: Model, - tile: usize, + sizes: Sizes, halo: usize, } impl OnnxNet { - /// The network at `path`, which needs `halo` photosites of context - /// ([`crate::Shipped::halo`]). - pub fn from_path(path: &std::path::Path, halo: usize) -> Result { + /// The network `shipped`, whose fixed-tile file is at `path`. + /// + /// On a rung that runs any input size (TensorRT, the CUDA provider — + /// [`dr_inference_engine::whole_frame_limit`]) and with the any-size + /// export installed beside it, the whole-frame network: the frame in one + /// call, or the fewest large tiles that fit (§14). Its output is the + /// fixed tiles' to rounding. Everywhere else, and if the whole-frame + /// model will not open, the 1408² tiles. + pub fn open(path: &std::path::Path, shipped: Shipped) -> Result { + if let Some(max) = dr_inference_engine::whole_frame_limit() { + let whole = path.with_file_name(shipped.whole); + if whole.is_file() { + let opened = std::fs::read(&whole) + .map_err(DenoiseError::from) + .and_then(|bytes| { + Ok(dr_inference_engine::open( + Role::WholeDenoiser, + Form::F32, + &bytes, + )?) + }); + match opened { + Ok(model) => { + return Ok(OnnxNet { + model, + sizes: Sizes::Any { align: ALIGN, max }, + halo: shipped.halo, + }) + } + Err(e) => log::warn!( + "learned denoise: {} will not open ({e}); running 1408² tiles", + whole.display() + ), + } + } + } + Self::open_tiled(path, shipped) + } + + /// The fixed-tile network at `path`, whatever the rung: 1408² tiles. + pub fn open_tiled(path: &std::path::Path, shipped: Shipped) -> Result { let (path, form) = dr_inference_engine::resolve_model(Role::Denoiser, path); let bytes = std::fs::read(&path)?; Ok(OnnxNet { model: dr_inference_engine::open(Role::Denoiser, form, &bytes)?, - tile: TILE, - halo, + sizes: Sizes::Square(TILE), + halo: shipped.halo, }) } @@ -41,11 +91,16 @@ impl OnnxNet { pub fn rung(&self) -> Result { Ok(self.model.acquire()?.rung()) } + + /// Whether this is the whole-frame network. + pub fn whole_frame(&self) -> bool { + matches!(self.sizes, Sizes::Any { .. }) + } } impl TileNet for OnnxNet { - fn tile(&self) -> usize { - self.tile + fn sizes(&self) -> Sizes { + self.sizes } fn halo(&self) -> usize { @@ -54,12 +109,13 @@ impl TileNet for OnnxNet { fn run( &mut self, + rows: usize, + cols: usize, mosaic: Vec, sigma: Vec, write: &mut dyn FnMut(&[f32]), ) -> Result<(), DenoiseError> { - let n = self.tile; - let shape = ndarray::IxDyn(&[1, 1, n, n]); + let shape = ndarray::IxDyn(&[1, 1, rows, cols]); // The vectors become the tensors: no copy on the way in. let m = ort::value::Tensor::from_array( ndarray::Array::from_shape_vec(shape.clone(), mosaic) @@ -74,9 +130,9 @@ impl TileNet for OnnxNet { let outputs = session.run(ort::inputs!["mosaic" => m, "sigma" => s])?; let (shape, data) = outputs[0].try_extract_tensor::()?; let dims: Vec = shape.iter().copied().collect(); - if dims != [1, 3, n as i64, n as i64] { + if dims != [1, 3, rows as i64, cols as i64] { return Err(DenoiseError::Model(format!( - "output is {dims:?}, expected [1, 3, {n}, {n}]" + "output is {dims:?}, expected [1, 3, {rows}, {cols}]" ))); } // And none on the way out: the frame is written from the runtime's buffer. diff --git a/core/dr-denoise/src/tile.rs b/core/dr-denoise/src/tile.rs index 1cc9344..35ee40f 100644 --- a/core/dr-denoise/src/tile.rs +++ b/core/dr-denoise/src/tile.rs @@ -1,12 +1,18 @@ //! TRACES: FR-DEV-3g -//! A whole frame through a fixed-shape network, exactly (denoise.md §3.4). +//! A whole frame through a network, in tiles, exactly (denoise.md §3.4, §14). //! -//! The network sees `TILE_IN`² photosites and its output is exact in the -//! central `TILE_IN − 2·HALO`: the halo is wider than its receptive field -//! (185 photosites, counted from the layers), so a tile's centre equals the -//! whole frame's at the same place. The frame is extended by reflection -//! about its edge photosites, which keeps every photosite's CFA colour, so -//! edge tiles see real context too. +//! A tile's output is exact in its centre: past a halo wider than the +//! network's receptive field (185 photosites for a single network, more for +//! the mixture), a tile's centre equals the whole frame's at the same place. +//! The frame is extended by reflection about its edge photosites, which +//! keeps every photosite's CFA colour, so edge tiles see real context too. +//! +//! **Tile sizes.** A fixed-shape network takes one square ([`Sizes::Square`], +//! 1408², of which Best keeps 896²). A network exported with any height and +//! width ([`Sizes::Any`]) takes the frame whole when it is small enough, and +//! otherwise the fewest equal tiles that are: [`plan`] picks the grid that +//! computes the fewest photosites. If the first tile of a plan fails — a +//! GPU out of memory — the limit is halved and the frame planned again. //! //! **Phase.** The network was trained on RGGB. A frame whose pattern starts //! on another colour is read from one photosite up and/or left — the @@ -20,16 +26,26 @@ use dr_decode::CfaPattern; /// [`TileNet::halo`]. pub const HALO: usize = 192; -/// A fixed-shape network: `mosaic` and `sigma`, `n×n` RGGB, in; `3×n×n` +/// The tiles a network takes. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Sizes { + /// One square, `n` photosites a side. + Square(usize), + /// Any rectangle whose sides are multiples of `align`, at most `max` + /// (rows, columns). + Any { align: usize, max: (usize, usize) }, +} + +/// A network: `mosaic` and `sigma`, `rows×cols` RGGB, in; `3×rows×cols` /// planar linear camera RGB out. /// /// The inputs are handed over, and the output is lent to `write` rather than -/// returned: a 1408² tile is 24 MB of output, and copying it out of the -/// runtime's buffer and back into the frame was a measurable share of a -/// frame's time. +/// returned: a 1408² tile is 24 MB of output and a whole frame 300 MB, and +/// copying it out of the runtime's buffer and back into the frame was a +/// measurable share of a frame's time. pub trait TileNet { - /// The edge `n` of the square tile the network takes. - fn tile(&self) -> usize; + /// The tile sizes it takes. + fn sizes(&self) -> Sizes; /// Photosites of context it needs past a tile's kept centre: at least /// its receptive field. [`HALO`] unless the network says otherwise. fn halo(&self) -> usize { @@ -37,12 +53,87 @@ pub trait TileNet { } fn run( &mut self, + rows: usize, + cols: usize, mosaic: Vec, sigma: Vec, write: &mut dyn FnMut(&[f32]), ) -> Result<(), crate::DenoiseError>; } +/// How a frame is cut: every tile `rows × cols` in, keeping its centre +/// `core.0 × core.1` past the halo, on a `grid.0 × grid.1` grid. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct Plan { + pub rows: usize, + pub cols: usize, + pub core: (usize, usize), + pub grid: (usize, usize), +} + +impl Plan { + /// Photosites the network computes for the frame. + pub fn work(&self) -> usize { + self.grid.0 * self.grid.1 * self.rows * self.cols + } +} + +/// The tiles for an `uh × uw` frame (in the network's phase) at `sizes`, or +/// `None` when no tile fits. +/// +/// Square tiles are today's grid. Any-size tiles are equal on each axis, so +/// one call shape serves the frame — TensorRT's profile tunes for one, and +/// the CUDA provider searches its algorithms once per shape — and the grid +/// is the one with the least work: one tile whenever the frame and its +/// halo fit under `max`. +pub fn plan(uh: usize, uw: usize, halo: usize, sizes: Sizes) -> Option { + match sizes { + Sizes::Square(n) => { + if n <= 2 * halo || !(n - 2 * halo).is_multiple_of(2) { + return None; + } + let core = n - 2 * halo; + Some(Plan { + rows: n, + cols: n, + core: (core, core), + grid: (uh.div_ceil(core), uw.div_ceil(core)), + }) + } + Sizes::Any { align, max } => { + // An even align keeps every tile origin on an even photosite, + // so every tile starts on red. + let align = align.max(2).next_multiple_of(2); + let axis = |extent: usize, tiles: usize, limit: usize| { + let size = (extent.div_ceil(tiles) + 2 * halo).next_multiple_of(align); + let core = size.checked_sub(2 * halo)?; + (size <= limit && core > 0 && core.is_multiple_of(2)).then_some((size, core)) + }; + let mut best: Option = None; + for gy in 1..=16 { + let Some((rows, cy)) = axis(uh, gy, max.0) else { + continue; + }; + for gx in 1..=16 { + let Some((cols, cx)) = axis(uw, gx, max.1) else { + continue; + }; + let p = Plan { + rows, + cols, + core: (cy, cx), + grid: (uh.div_ceil(cy), uw.div_ceil(cx)), + }; + if best.is_none_or(|b| p.work() < b.work()) { + best = Some(p); + } + } + } + best + } + } +} + /// Index into `0..n` by reflection about the end photosites, any distance /// out: …2 1 [0 1 2 … n−1] n−2 n−3…, period `2(n−1)`. Parity is kept, which /// is what keeps a CFA colour. @@ -71,7 +162,9 @@ pub fn rggb_offset(p: CfaPattern) -> Option<(usize, usize)> { /// `sigma(colour, value)`, and return `h×w` interleaved RGB. /// /// `progress(done, total)` is called after each tile and stops the run by -/// returning `false`, in which case the result is `Ok(None)`. +/// returning `false`, in which case the result is `Ok(None)`. An any-size +/// network whose first tile fails is planned again with tiles half that +/// size, until a tile would keep no centre; then the failure is returned. #[allow(clippy::too_many_arguments)] pub fn run_tiled( net: &mut dyn TileNet, @@ -85,40 +178,103 @@ pub fn run_tiled( let (dy, dx) = rggb_offset(pattern).ok_or_else(|| { crate::DenoiseError::Unsupported(format!("{pattern:?} is not a Bayer pattern")) })?; - let (n, halo) = (net.tile(), net.halo()); - if n <= 2 * halo || !(n - 2 * halo).is_multiple_of(2) { - return Err(crate::DenoiseError::Model(format!( - "tile {n} leaves no even centre past a {halo} halo" - ))); + let halo = net.halo(); + let (uh, uw) = (h + dy, w + dx); + let mut sizes = net.sizes(); + loop { + let plan = plan(uh, uw, halo, sizes).ok_or_else(|| { + crate::DenoiseError::Model(format!( + "no tile of {sizes:?} keeps a centre past a {halo} halo" + )) + })?; + match run_plan(net, plan, h, w, (dy, dx), halo, at, sigma, progress) { + Err(Failed { error, first: true }) => { + // The first call of a size is where a GPU runs out of + // memory. Halve the larger kept centre of the tile that + // failed — not the limit, which may be far above it, and not + // the tile, half of which may be all halo — and plan again, + // until no smaller tile keeps a centre. + let Sizes::Any { align, .. } = sizes else { + return Err(error); + }; + let (cr, cc) = plan.core; + let smaller = if cr >= cc { + (cr / 2 + 2 * halo, plan.cols) + } else { + (plan.rows, cc / 2 + 2 * halo) + }; + let next = Sizes::Any { + align, + max: smaller, + }; + if self::plan(uh, uw, halo, next).is_none() { + return Err(error); + } + log::warn!( + "learned denoise: a {}×{} tile failed ({error}); trying tiles up to {}×{}", + plan.rows, + plan.cols, + smaller.0, + smaller.1 + ); + sizes = next; + } + Err(Failed { error, .. }) => return Err(error), + Ok(done) => return Ok(done), + } } - let core = n - 2 * halo; +} + +/// A run that stopped on an error, and whether it was the plan's first call. +struct Failed { + error: crate::DenoiseError, + first: bool, +} + +#[allow(clippy::too_many_arguments)] +fn run_plan( + net: &mut dyn TileNet, + plan: Plan, + h: usize, + w: usize, + (dy, dx): (usize, usize), + halo: usize, + at: &(dyn Fn(usize, usize) -> f32 + Sync), + sigma: &(dyn Fn(usize, f32) -> f32 + Sync), + progress: &mut dyn FnMut(usize, usize) -> bool, +) -> Result>, Failed> { + let Plan { + rows: nr, + cols: nc, + core: (cr, cc), + grid: (ty, tx), + } = plan; // In unified coordinates the frame spans u ∈ [dy, dy + h), v ∈ [dx, dx + w). let (uh, uw) = (h + dy, w + dx); - let (ty, tx) = (uh.div_ceil(core), uw.div_ceil(core)); let total = ty * tx; let origins: Vec<(usize, usize)> = (0..ty) - .flat_map(|i| (0..tx).map(move |j| (i * core, j * core))) + .flat_map(|i| (0..tx).map(move |j| (i * cr, j * cc))) .collect(); let threads = std::thread::available_parallelism().map_or(1, |n| n.get()); // One tile's mosaic and σ, gathered on every core: rows are independent. let gather = |u0: usize, v0: usize| { - let mut mos = vec![0.0f32; n * n]; - let mut sig = vec![0.0f32; n * n]; - let rows_per = n.div_ceil(threads).max(1); + let mut mos = vec![0.0f32; nr * nc]; + let mut sig = vec![0.0f32; nr * nc]; + let rows_per = nr.div_ceil(threads).max(1); std::thread::scope(|scope| { for (chunk, (m, s)) in mos - .chunks_mut(rows_per * n) - .zip(sig.chunks_mut(rows_per * n)) + .chunks_mut(rows_per * nc) + .zip(sig.chunks_mut(rows_per * nc)) .enumerate() { scope.spawn(move || { - for (i, (mrow, srow)) in m.chunks_mut(n).zip(s.chunks_mut(n)).enumerate() { + for (i, (mrow, srow)) in m.chunks_mut(nc).zip(s.chunks_mut(nc)).enumerate() { let r = chunk * rows_per + i; // Unified row u = u0 + r − halo; frame row y = u − dy, reflected. let u = u0 as isize + r as isize - halo as isize; let y = reflect(u - dy as isize, h); - for c in 0..n { + for c in 0..nc { let v = v0 as isize + c as isize - halo as isize; let x = reflect(v - dx as isize, w); let val = at(y, x); @@ -138,7 +294,14 @@ pub fn run_tiled( // most two tiles' inputs alive. let mut out = vec![0.0f32; h * w * 3]; let stop = std::sync::atomic::AtomicBool::new(false); - std::thread::scope(|scope| -> Result, crate::DenoiseError> { + let fail = |error, k: usize, stop: &std::sync::atomic::AtomicBool| { + stop.store(true, std::sync::atomic::Ordering::Relaxed); + Failed { + error, + first: k == 0, + } + }; + std::thread::scope(|scope| -> Result, Failed> { let (tx_tiles, rx_tiles) = std::sync::mpsc::sync_channel(1); let (origins, stop, gather) = (&origins, &stop, &gather); scope.spawn(move || { @@ -156,18 +319,15 @@ pub fn run_tiled( break; }; let mut wrong = None; - let ran = net.run(mos, sig, &mut |rgb: &[f32]| { - if rgb.len() != 3 * n * n { + let ran = net.run(nr, nc, mos, sig, &mut |rgb: &[f32]| { + if rgb.len() != 3 * nr * nc { wrong = Some(rgb.len()); return; } // The tile's centre back into the frame: the frame rows it covers, // split across cores (each row is written by one thread only). - let (y_lo, y_hi) = ( - (u0 + dy.saturating_sub(u0)).max(dy) - dy, - (u0 + core).min(uh) - dy, - ); - let (x_lo, x_hi) = ((v0.max(dx)) - dx, (v0 + core).min(uw) - dx); + let (y_lo, y_hi) = (u0.max(dy) - dy, (u0 + cr).min(uh) - dy); + let (x_lo, x_hi) = (v0.max(dx) - dx, (v0 + cc).min(uw) - dx); if y_hi > y_lo && x_hi > x_lo { let rows = &mut out[y_lo * w * 3..y_hi * w * 3]; let per = (y_hi - y_lo).div_ceil(threads).max(1); @@ -181,7 +341,7 @@ pub fn run_tiled( for x in x_lo..x_hi { let c = x + dx + halo - v0; for ch in 0..3 { - row[x * 3 + ch] = rgb[ch * n * n + r * n + c]; + row[x * 3 + ch] = rgb[ch * nr * nc + r * nc + c]; } } } @@ -191,14 +351,18 @@ pub fn run_tiled( } }); if let Err(e) = ran { - stop.store(true, std::sync::atomic::Ordering::Relaxed); - return Err(e); + while rx_tiles.try_recv().is_ok() {} + return Err(fail(e, k, stop)); } if let Some(len) = wrong { - stop.store(true, std::sync::atomic::Ordering::Relaxed); - return Err(crate::DenoiseError::Model(format!( - "network returned {len} values for a {n}² tile" - ))); + while rx_tiles.try_recv().is_ok() {} + return Err(fail( + crate::DenoiseError::Model(format!( + "network returned {len} values for a {nr}×{nc} tile" + )), + k, + stop, + )); } if !progress(k + 1, total) { stop.store(true, std::sync::atomic::Ordering::Relaxed); @@ -236,37 +400,65 @@ mod tests { /// is its 2×2 quad's (R, mean G, B), averaged over the quads within /// `reach` quads. Purely a function of the tile, like the real one. struct BoxNet { - n: usize, + sizes: Sizes, reach: usize, + /// Fails any call with more photosites than this, as a GPU out of + /// memory does. + fails_above: usize, + calls: Vec<(usize, usize)>, + } + + fn square(n: usize, reach: usize) -> BoxNet { + BoxNet { + sizes: Sizes::Square(n), + reach, + fails_above: usize::MAX, + calls: Vec::new(), + } + } + + fn any(max: (usize, usize), reach: usize) -> BoxNet { + BoxNet { + sizes: Sizes::Any { align: 16, max }, + reach, + fails_above: usize::MAX, + calls: Vec::new(), + } } impl TileNet for BoxNet { - fn tile(&self) -> usize { - self.n + fn sizes(&self) -> Sizes { + self.sizes } fn run( &mut self, + rows: usize, + cols: usize, m: Vec, _s: Vec, write: &mut dyn FnMut(&[f32]), ) -> Result<(), crate::DenoiseError> { - let n = self.n; - let q = n / 2; + self.calls.push((rows, cols)); + if rows * cols > self.fails_above { + return Err(crate::DenoiseError::Model("out of memory".into())); + } + let (qr, qc) = (rows / 2, cols / 2); let quad = |qy: usize, qx: usize| { let (y, x) = (2 * qy, 2 * qx); [ - m[y * n + x], - 0.5 * (m[y * n + x + 1] + m[(y + 1) * n + x]), - m[(y + 1) * n + x + 1], + m[y * cols + x], + 0.5 * (m[y * cols + x + 1] + m[(y + 1) * cols + x]), + m[(y + 1) * cols + x + 1], ] }; - let mut out = vec![0.0; 3 * n * n]; - for qy in 0..q { - for qx in 0..q { + let plane = rows * cols; + let mut out = vec![0.0; 3 * plane]; + for qy in 0..qr { + for qx in 0..qc { let mut acc = [0.0f32; 3]; let mut cnt = 0.0; - for a in qy.saturating_sub(self.reach)..(qy + self.reach + 1).min(q) { - for b in qx.saturating_sub(self.reach)..(qx + self.reach + 1).min(q) { + for a in qy.saturating_sub(self.reach)..(qy + self.reach + 1).min(qr) { + for b in qx.saturating_sub(self.reach)..(qx + self.reach + 1).min(qc) { let v = quad(a, b); for c in 0..3 { acc[c] += v[c]; @@ -276,7 +468,7 @@ mod tests { } for (dy, dx) in [(0, 0), (0, 1), (1, 0), (1, 1)] { for c in 0..3 { - out[c * n * n + (2 * qy + dy) * n + 2 * qx + dx] = acc[c] / cnt; + out[c * plane + (2 * qy + dy) * cols + 2 * qx + dx] = acc[c] / cnt; } } } @@ -306,10 +498,7 @@ mod tests { ] { let (h, w) = (300, 410); let at = field(p); - let mut net = BoxNet { - n: 2 * HALO + 64, - reach: 0, - }; + let mut net = square(2 * HALO + 64, 0); let out = run_tiled(&mut net, h, w, p, &at, &|_, _| 0.01, &mut |_, _| true) .unwrap() .unwrap(); @@ -336,14 +525,8 @@ mod tests { let (h, w) = (230, 170); for p in [CfaPattern::Rggb, CfaPattern::Bggr] { let at = |y: usize, x: usize| ((y * 7919 + x * 104729) % 1000) as f32 / 1000.0; - let mut small = BoxNet { - n: 2 * HALO + 32, - reach: 20, - }; - let mut big = BoxNet { - n: 2 * HALO + 256, - reach: 20, - }; + let mut small = square(2 * HALO + 32, 20); + let mut big = square(2 * HALO + 256, 20); let a = run_tiled(&mut small, h, w, p, &at, &|_, _| 0.0, &mut |_, _| true) .unwrap() .unwrap(); @@ -359,12 +542,99 @@ mod tests { } } + /// The same frame through square tiles, one whole-frame call, a grid of + /// any-size tiles, and a network that runs out of memory on the whole + /// frame and is planned again: one answer. + #[test] + fn any_size_tiles_give_the_square_tiles_answer() { + let (h, w) = (230, 170); + for p in [ + CfaPattern::Rggb, + CfaPattern::Grbg, + CfaPattern::Gbrg, + CfaPattern::Bggr, + ] { + let at = |y: usize, x: usize| ((y * 7919 + x * 104729) % 1000) as f32 / 1000.0; + let run = |net: &mut BoxNet| { + run_tiled(net, h, w, p, &at, &|_, _| 0.0, &mut |_, _| true) + .unwrap() + .unwrap() + }; + // A reach of 6 quads is well inside the halo, and keeps a + // debug-build test of four phases short. + let want = run(&mut square(2 * HALO + 32, 6)); + + let mut whole = any((4096, 4096), 6); + let got = run(&mut whole); + assert_eq!(whole.calls.len(), 1, "the frame fits: one call"); + assert_eq!(got, want, "{p:?}: whole frame"); + + let mut grid = any((2 * HALO + 96, 2 * HALO + 64), 6); + let got = run(&mut grid); + assert!(grid.calls.len() > 1); + assert!( + grid.calls.windows(2).all(|c| c[0] == c[1]), + "one call shape" + ); + assert_eq!(got, want, "{p:?}: a grid of any-size tiles"); + + let mut tight = any((4096, 4096), 6); + tight.fails_above = (2 * HALO + 200) * (2 * HALO + 200); + let got = run(&mut tight); + assert_eq!(got, want, "{p:?}: planned again after a failure"); + assert!(tight.calls.len() > 2, "the whole frame failed, then tiles"); + } + } + + #[test] + fn the_plan_is_one_tile_when_the_frame_fits_and_the_least_work_when_not() { + // A 6D frame with Best's halo, under the whole-frame limit: one call. + let one = plan( + 3648, + 5472, + 256, + Sizes::Any { + align: 16, + max: (4608, 6656), + }, + ) + .unwrap(); + assert_eq!(one.grid, (1, 1)); + assert_eq!((one.rows, one.cols), (4160, 5984)); + assert!(one.core.0 >= 3648 && one.core.1 >= 5472); + // Too wide for one: the cheapest grid, every tile within the limit. + let two = plan( + 3648, + 8192, + 256, + Sizes::Any { + align: 16, + max: (4608, 6656), + }, + ) + .unwrap(); + assert!(two.cols <= 6656 && two.rows <= 4608); + assert_eq!(two.grid, (1, 2)); + // And always less work than today's 1408 squares. + let squares = plan(3648, 5472, 256, Sizes::Square(1408)).unwrap(); + assert_eq!(squares.grid, (5, 7)); + assert!(one.work() * 2 < squares.work()); + // A limit no tile fits under. + assert!(plan( + 3648, + 5472, + 256, + Sizes::Any { + align: 16, + max: (400, 400) + } + ) + .is_none()); + } + #[test] fn a_cancelled_run_returns_nothing() { - let mut net = BoxNet { - n: 2 * HALO + 32, - reach: 0, - }; + let mut net = square(2 * HALO + 32, 0); let r = run_tiled( &mut net, 100, @@ -389,11 +659,13 @@ mod timing { struct Null(usize, Vec); impl TileNet for Null { - fn tile(&self) -> usize { - self.0 + fn sizes(&self) -> Sizes { + Sizes::Square(self.0) } fn run( &mut self, + _rows: usize, + _cols: usize, m: Vec, _s: Vec, write: &mut dyn FnMut(&[f32]), diff --git a/ui/dr-ui/src/develop/denoise.rs b/ui/dr-ui/src/develop/denoise.rs index 51b7583..33c729c 100644 --- a/ui/dr-ui/src/develop/denoise.rs +++ b/ui/dr-ui/src/develop/denoise.rs @@ -179,7 +179,7 @@ impl DevelopSession { iso: self.denoise.iso, cache_key: self.denoise.file_hash.as_ref().map(|h| h.key(&model)), model, - halo: net.halo, + net, cancel, }) } @@ -316,8 +316,8 @@ struct Work { profile: Option>, iso: Option, model: std::path::PathBuf, - /// The context `model` needs past a tile's kept centre. - halo: usize, + /// The network `model` is: its context and its whole-frame sibling. + net: dr_denoise::Shipped, cancel: Arc, cache_key: Option, } @@ -351,8 +351,8 @@ impl Work { .map_err(|e| e.to_string())?; let noise = dr_denoise::noise::for_frame_with(&raw, self.profile.as_deref(), self.iso) .ok_or("this photograph gives no way to measure its noise")?; - let mut net = dr_denoise::onnx::OnnxNet::from_path(&self.model, self.halo) - .map_err(|e| e.to_string())?; + let mut net = + dr_denoise::onnx::OnnxNet::open(&self.model, self.net).map_err(|e| e.to_string())?; let rung = net .rung() .map(|r| r.label().to_string()) diff --git a/ui/dr-ui/src/inference.rs b/ui/dr-ui/src/inference.rs index 88f253a..0ab8123 100644 --- a/ui/dr-ui/src/inference.rs +++ b/ui/dr-ui/src/inference.rs @@ -35,9 +35,11 @@ pub fn init(runtime_dirs: Vec) { (Role::EyeClassifier, crate::library::SUNGLASSES_MODEL), (Role::Inpainter, crate::library::INPAINT_MODEL), ]); - wanted.extend( - [dr_denoise::FAST, dr_denoise::MEDIUM, dr_denoise::BEST].map(|n| (Role::Denoiser, n.file)), - ); + let denoisers = [dr_denoise::FAST, dr_denoise::MEDIUM, dr_denoise::BEST]; + wanted.extend(denoisers.map(|n| (Role::Denoiser, n.file))); + // Their any-size siblings, which the engine compiles only on a rung + // that runs whole frames (TensorRT; denoise.md §14). + wanted.extend(denoisers.map(|n| (Role::WholeDenoiser, n.whole))); let models: Vec<(Role, PathBuf)> = wanted .into_iter() .filter_map(|(role, name)| Some((role, crate::library::shared_model(name)?)))