dr-gpu: the merge pass — warp, accumulate, resolve, chunk by chunk

merge.wgsl warps one camera-space tile into one output chunk — output
pixel to direction (the projection maths of dr_pano::projection, verbatim),
direction to the frame's camera, camera to source pixel, bilinear by hand
from four textureLoads because rgba32float is not filterable — and adds it
into a storage-buffer accumulator weighted by its distance from the
frame's edge. A resolve pass divides by the weights and packs sixteen-bit
samples at the sensor's scale with a coverage bit.

MergePass::merge drives it: bands of rows, chunks across a band, and for
each chunk only the frames whose footprint meets it, each rendered as the
source rectangle the chunk needs and nothing more. The working set is one
chunk, one tile and one band (FR-MRG-11); the frame textures are the
caller's to cache. Feathered, not seamed; gain a scalar per frame — the
blend quality is panorama.md §10's step 5, after the path writes a file.
This commit is contained in:
2026-09-19 15:24:12 +02:00
parent acab0d7abb
commit 44ea763c61
4 changed files with 698 additions and 0 deletions
+539
View File
@@ -0,0 +1,539 @@
//! TRACES: FR-MRG-10 | FR-MRG-11
//! The merge: source frames warped into an output surface, chunk by chunk.
//!
//! The per-pixel half of a panorama (FR-MRG-10), on the GPU: the warp of a
//! source tile into an output chunk, the weighted accumulation across
//! frames, and the resolve to sixteen-bit samples. The geometry it is
//! given — rotations, focal length, projection — is `dr-pano`'s, solved on
//! proxies before any full-resolution pixel exists (panorama.md §5), and
//! that is what makes this simple: every output pixel's source coordinates
//! are a closed-form function, so a chunk can be produced from the source
//! tiles that project into it and nothing else.
//!
//! # The loop
//!
//! ```text
//! for each band of rows of the output:
//! for each chunk across the band:
//! zero the accumulator
//! for each frame whose footprint meets the chunk:
//! the source rectangle the chunk needs, from the geometry
//! render it camera-linear through the pipeline (the tile)
//! warp the tile into the chunk, accumulate ← GPU
//! resolve the chunk to u16 ← GPU
//! copy it into the band
//! hand the band to the writer (one DNG strip)
//! ```
//!
//! No stage holds the composite (FR-MRG-11): the working set is one
//! chunk's accumulator, one tile, one band of u16 rows. The frame textures
//! are the caller's to provide and cache — `source` is asked for frame `k`
//! as it is needed, and a caller short of memory may demosaic on demand.
//!
//! # What is not here yet
//!
//! A feathered blend, not seams and a Laplacian pyramid: the weight is the
//! distance to the frame's edge, which hides exposure steps and small
//! misalignments and does not hide parallax. Gain is a scalar per frame
//! the caller supplies. Both are panorama.md §10's step 5, after the path
//! writes a file end to end.
use std::sync::Arc;
use dr_pano::bundle::Cameras;
use dr_pano::projection::{Bounds, Projection};
use wgpu::util::DeviceExt;
use crate::readback::await_mapping;
use crate::{AdjustPass, DemosaicedImage, GpuContext, GpuError};
/// One frame's part in the merge.
pub struct MergeFrame {
/// The frame's edit, for its lens corrections — the only part of an
/// edit the camera-space tap uses (FR-MRG-2).
pub graph: Arc<dr_pipeline::EditGraph>,
/// Multiplies the frame's samples, to bring its exposure to the
/// reference frame's. 1.0 for no correction.
pub gain: f32,
}
/// The output the merge produces.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct MergeOutput {
pub projection: Projection,
/// The projection's scale in output pixels: the cylinder's radius, the
/// plane's distance. The source focal length at full resolution gives
/// output pixels the size of source pixels at the centre.
pub scale: f64,
/// The rectangle of the projection to produce, centred coordinates.
pub bounds: Bounds,
/// Pixels over which a frame's weight ramps up from its edge.
pub feather: f32,
/// Chunk size: the unit of GPU work and of memory.
pub chunk: (u32, u32),
/// Multiplies a normalised sample (1.0 = white) to the sensor's scale.
pub sample_scale: f32,
}
impl MergeOutput {
pub fn width(&self) -> u32 {
self.bounds.width().ceil().max(1.0) as u32
}
pub fn height(&self) -> u32 {
self.bounds.height().ceil().max(1.0) as u32
}
}
/// A band of finished rows: `rows × width × 3` RGB `u16`, plus a coverage
/// mask (`true` where any frame reached the pixel).
pub struct Band<'a> {
pub first_row: u32,
pub rows: u32,
pub rgb: &'a [u16],
pub covered: &'a [bool],
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct WarpParams {
chunk_origin: [f32; 2],
chunk_size: [u32; 2],
projection: u32,
proj_scale: f32,
focal: f32,
gain: f32,
r0: [f32; 4],
r1: [f32; 4],
r2: [f32; 4],
frame_size: [f32; 2],
tile_origin: [f32; 2],
tile_size: [u32; 2],
feather: f32,
_pad: f32,
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct ResolveParams {
chunk_size: [u32; 2],
scale: f32,
_pad: f32,
}
/// The two pipelines and the chunk buffers.
pub struct MergePass {
ctx: GpuContext,
warp: wgpu::ComputePipeline,
warp_layout: wgpu::BindGroupLayout,
resolve: wgpu::ComputePipeline,
resolve_layout: wgpu::BindGroupLayout,
/// Accumulator and packed output for the current chunk size.
buffers: Option<(wgpu::Buffer, wgpu::Buffer, wgpu::Buffer, (u32, u32))>,
}
impl MergePass {
pub fn new(ctx: &GpuContext) -> Result<Self, GpuError> {
let module = ctx
.device
.create_shader_module(wgpu::ShaderModuleDescriptor {
label: Some("merge"),
source: wgpu::ShaderSource::Wgsl(include_str!("shaders/merge.wgsl").into()),
});
let uniform = |binding| wgpu::BindGroupLayoutEntry {
binding,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Uniform,
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
};
let storage = |binding, read_only| wgpu::BindGroupLayoutEntry {
binding,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Buffer {
ty: wgpu::BufferBindingType::Storage { read_only },
has_dynamic_offset: false,
min_binding_size: None,
},
count: None,
};
let warp_layout = ctx
.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("merge-warp-bgl"),
entries: &[
uniform(0),
wgpu::BindGroupLayoutEntry {
binding: 1,
visibility: wgpu::ShaderStages::COMPUTE,
ty: wgpu::BindingType::Texture {
// Unfilterable: rgba32float, loaded by hand.
sample_type: wgpu::TextureSampleType::Float { filterable: false },
view_dimension: wgpu::TextureViewDimension::D2,
multisampled: false,
},
count: None,
},
storage(2, false),
],
});
let resolve_layout = ctx
.device
.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
label: Some("merge-resolve-bgl"),
entries: &[uniform(0), storage(1, true), storage(2, false)],
});
let pipeline = |name: &str, layout: &wgpu::BindGroupLayout| {
let pl = ctx
.device
.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
label: Some(name),
bind_group_layouts: &[Some(layout)],
immediate_size: 0,
});
ctx.device
.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
label: Some(name),
layout: Some(&pl),
module: &module,
entry_point: Some(name),
compilation_options: Default::default(),
cache: None,
})
};
Ok(MergePass {
ctx: ctx.clone(),
warp: pipeline("warp", &warp_layout),
warp_layout,
resolve: pipeline("resolve", &resolve_layout),
resolve_layout,
buffers: None,
})
}
/// Allocate the chunk buffers for this size if the last ones differ.
fn ensure_buffers(&mut self, chunk: (u32, u32)) {
if self.buffers.as_ref().is_none_or(|b| b.3 != chunk) {
let n = u64::from(chunk.0) * u64::from(chunk.1);
let acc = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("merge-acc"),
size: n * 16,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
mapped_at_creation: false,
});
let out = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("merge-out"),
size: n * 8,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let read = self.ctx.device.create_buffer(&wgpu::BufferDescriptor {
label: Some("merge-read"),
size: n * 8,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
});
self.buffers = Some((acc, out, read, chunk));
}
}
fn chunk_buffers(&self) -> (&wgpu::Buffer, &wgpu::Buffer, &wgpu::Buffer) {
let b = self.buffers.as_ref().expect("ensured by the caller");
(&b.0, &b.1, &b.2)
}
/// Produce the whole output, band by band, handing each finished band
/// to `sink`.
///
/// `cameras` are in **full-resolution source pixels** (`frame_size`),
/// with frame `k` corresponding to `frames[k]` and `source(k)`. `source`
/// supplies the demosaiced frame on demand and may cache as it sees fit.
#[allow(clippy::too_many_arguments)]
pub fn merge<S, F>(
&mut self,
adjust: &mut AdjustPass,
frames: &[MergeFrame],
cameras: &Cameras,
frame_size: (u32, u32),
output: &MergeOutput,
mut source: S,
mut sink: F,
mut cancelled: impl FnMut() -> bool,
) -> Result<(), GpuError>
where
S: FnMut(usize) -> Result<Arc<DemosaicedImage>, GpuError>,
F: FnMut(Band<'_>) -> Result<(), GpuError>,
{
let (out_w, out_h) = (output.width(), output.height());
let (cw, ch) = (output.chunk.0.max(8), output.chunk.1.max(8));
let (fw, fh) = (frame_size.0 as f64, frame_size.1 as f64);
let mut band_rgb = vec![0u16; (out_w * ch * 3) as usize];
let mut band_cov = vec![false; (out_w * ch) as usize];
let mut chunk_px: Vec<u32> = Vec::new();
let mut y = 0u32;
while y < out_h {
let rows = ch.min(out_h - y);
band_rgb.iter_mut().for_each(|v| *v = 0);
band_cov.iter_mut().for_each(|v| *v = false);
let mut x = 0u32;
while x < out_w {
if cancelled() {
return Err(GpuError::Readback("merge cancelled".into()));
}
let cols = cw.min(out_w - x);
let origin = (
output.bounds.min_u + f64::from(x),
output.bounds.min_v + f64::from(y),
);
self.zero_accumulator((cols, rows));
for (k, frame) in frames.iter().enumerate() {
let Some(rect) = source_rect(
output.projection,
output.scale,
cameras,
k,
origin,
(cols, rows),
(fw, fh),
) else {
continue;
};
let image = source(k)?;
// The tile: that rectangle of the frame, camera-linear,
// at 1:1.
let view = dr_pipeline::CropRect {
x: (rect.0 as f32) / fw as f32,
y: (rect.1 as f32) / fh as f32,
width: (rect.2 as f32) / fw as f32,
height: (rect.3 as f32) / fh as f32,
};
let shader = frame.graph.compose_camera_linear(view);
let tile = adjust.render_camera_linear(&image, &shader, rect.2, rect.3)?;
let r = cameras.rotations[k].transpose();
let params = WarpParams {
chunk_origin: [origin.0 as f32, origin.1 as f32],
chunk_size: [cols, rows],
projection: match output.projection {
Projection::Perspective => 0,
Projection::Cylindrical => 1,
Projection::Spherical => 2,
},
proj_scale: output.scale as f32,
focal: cameras.focal as f32,
gain: frame.gain,
r0: [r.0[0][0] as f32, r.0[0][1] as f32, r.0[0][2] as f32, 0.0],
r1: [r.0[1][0] as f32, r.0[1][1] as f32, r.0[1][2] as f32, 0.0],
r2: [r.0[2][0] as f32, r.0[2][1] as f32, r.0[2][2] as f32, 0.0],
frame_size: [fw as f32, fh as f32],
tile_origin: [rect.0 as f32, rect.1 as f32],
tile_size: [rect.2, rect.3],
feather: output.feather,
_pad: 0.0,
};
self.accumulate(&params, tile);
}
self.resolve_chunk((cols, rows), output.sample_scale, &mut chunk_px)?;
// Into the band.
for row in 0..rows as usize {
for col in 0..cols as usize {
let px = chunk_px[(row * cols as usize + col) * 2..][..2].to_vec();
let i = row * out_w as usize + (x as usize + col);
band_rgb[i * 3] = (px[0] & 0xFFFF) as u16;
band_rgb[i * 3 + 1] = (px[0] >> 16) as u16;
band_rgb[i * 3 + 2] = (px[1] & 0xFFFF) as u16;
band_cov[i] = (px[1] >> 16) != 0;
}
}
x += cols;
}
sink(Band {
first_row: y,
rows,
rgb: &band_rgb[..(out_w * rows * 3) as usize],
covered: &band_cov[..(out_w * rows) as usize],
})?;
y += rows;
}
Ok(())
}
fn zero_accumulator(&mut self, chunk: (u32, u32)) {
self.ensure_buffers(chunk);
let (acc, _, _) = self.chunk_buffers();
let n = u64::from(chunk.0) * u64::from(chunk.1) * 16;
let mut enc = self.ctx.device.create_command_encoder(&Default::default());
enc.clear_buffer(acc, 0, Some(n));
self.ctx.queue.submit(Some(enc.finish()));
}
fn accumulate(&mut self, params: &WarpParams, tile: &wgpu::Texture) {
let chunk = (params.chunk_size[0], params.chunk_size[1]);
let uniforms = self
.ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("merge-warp-params"),
contents: bytemuck::bytes_of(params),
usage: wgpu::BufferUsages::UNIFORM,
});
let view = tile.create_view(&Default::default());
self.ensure_buffers(chunk);
let (acc, _, _) = self.chunk_buffers();
let bind = self
.ctx
.device
.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("merge-warp-bg"),
layout: &self.warp_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: uniforms.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: wgpu::BindingResource::TextureView(&view),
},
wgpu::BindGroupEntry {
binding: 2,
resource: acc.as_entire_binding(),
},
],
});
let mut enc = self.ctx.device.create_command_encoder(&Default::default());
{
let mut pass = enc.begin_compute_pass(&Default::default());
pass.set_pipeline(&self.warp);
pass.set_bind_group(0, &bind, &[]);
pass.dispatch_workgroups(chunk.0.div_ceil(8), chunk.1.div_ceil(8), 1);
}
self.ctx.queue.submit(Some(enc.finish()));
}
fn resolve_chunk(&mut self, chunk: (u32, u32), scale: f32, out: &mut Vec<u32>) -> Result<(), GpuError> {
let params = ResolveParams {
chunk_size: [chunk.0, chunk.1],
scale,
_pad: 0.0,
};
let uniforms = self
.ctx
.device
.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("merge-resolve-params"),
contents: bytemuck::bytes_of(&params),
usage: wgpu::BufferUsages::UNIFORM,
});
let n = u64::from(chunk.0) * u64::from(chunk.1);
self.ensure_buffers(chunk);
let (acc, packed, read) = self.chunk_buffers();
let bind = self
.ctx
.device
.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("merge-resolve-bg"),
layout: &self.resolve_layout,
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: uniforms.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: acc.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: packed.as_entire_binding(),
},
],
});
let mut enc = self.ctx.device.create_command_encoder(&Default::default());
{
let mut pass = enc.begin_compute_pass(&Default::default());
pass.set_pipeline(&self.resolve);
pass.set_bind_group(0, &bind, &[]);
pass.dispatch_workgroups(chunk.0.div_ceil(8), chunk.1.div_ceil(8), 1);
}
enc.copy_buffer_to_buffer(packed, 0, read, 0, n * 8);
self.ctx.queue.submit(Some(enc.finish()));
let slice = read.slice(..n * 8);
let (tx, rx) = std::sync::mpsc::channel();
slice.map_async(wgpu::MapMode::Read, move |r| {
let _ = tx.send(r);
});
await_mapping(&self.ctx, &rx)?;
{
let data = slice.get_mapped_range();
out.clear();
out.extend_from_slice(bytemuck::cast_slice::<u8, u32>(&data));
}
read.unmap();
Ok(())
}
}
/// The rectangle of frame `k` (x, y, w, h in source pixels) a chunk reads,
/// or `None` if the chunk sees nothing of the frame.
///
/// Walks the chunk's border, projects each point into the frame, and takes
/// the bounding box with a two-pixel margin for the bilinear fetch. The
/// border rather than the corners because under a cylinder or sphere the
/// extreme of a footprint is not at a corner.
fn source_rect(
projection: Projection,
scale: f64,
cameras: &Cameras,
k: usize,
origin: (f64, f64),
size: (u32, u32),
frame: (f64, f64),
) -> Option<(u32, u32, u32, u32)> {
let (w, h) = (f64::from(size.0), f64::from(size.1));
let steps = 16;
let mut min = (f64::MAX, f64::MAX);
let mut max = (f64::MIN, f64::MIN);
let mut any = false;
let mut visit = |u: f64, v: f64| {
let d = projection.to_direction(scale, u, v);
if let Some((x, y)) = cameras.project(k, d) {
let (x, y) = (x + frame.0 / 2.0, y + frame.1 / 2.0);
min = (min.0.min(x), min.1.min(y));
max = (max.0.max(x), max.1.max(y));
any = true;
}
};
for s in 0..=steps {
let t = f64::from(s) / f64::from(steps);
visit(origin.0 + w * t, origin.1);
visit(origin.0 + w * t, origin.1 + h);
visit(origin.0, origin.1 + h * t);
visit(origin.0 + w, origin.1 + h * t);
}
// The interior too, coarsely: a chunk can contain a frame entirely.
for i in 1..4 {
for j in 1..4 {
visit(origin.0 + w * f64::from(i) / 4.0, origin.1 + h * f64::from(j) / 4.0);
}
}
if !any {
return None;
}
let x0 = (min.0.floor() - 2.0).max(0.0);
let y0 = (min.1.floor() - 2.0).max(0.0);
let x1 = (max.0.ceil() + 2.0).min(frame.0);
let y1 = (max.1.ceil() + 2.0).min(frame.1);
if x1 <= x0 || y1 <= y0 {
return None;
}
Some((x0 as u32, y0 as u32, (x1 - x0) as u32, (y1 - y0) as u32))
}