use std::borrow::Cow;
use image::RgbaImage;
use wgpu::util::DeviceExt;
use super::GpuContext;
use crate::error::StitchError;
const SLAB_BYTE_BUDGET: u64 = 64 * 1024 * 1024;
const RESULT_CANDIDATE_BUDGET: u64 = 128 * 1024 * 1024;
const MAX_DISPATCH_THREADS: u32 = 65_535 * 16;
const WORKGROUP_DIM: u32 = 16;
pub(crate) struct SearchParams {
pub transpose: bool,
pub edges: bool,
pub window: u32,
pub crop: u32,
pub skip: u32,
pub offset_start: i32,
pub offset_end: i32,
}
#[derive(Debug, Clone, Copy)]
pub(crate) struct BestMatch {
pub score: u32,
pub anchor: u32,
pub offset: i32,
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct PreprocessParams {
in_width: u32,
in_height: u32,
out_width: u32,
out_height: u32,
halo_x: u32,
halo_y: u32,
mode: u32,
transpose: u32,
}
#[repr(C)]
#[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
struct SadParams {
width1: u32,
width2: u32,
window: u32,
anchor_count: u32,
offset_start: i32,
offset_count: u32,
_pad0: u32,
_pad1: u32,
}
struct ProcessedRegion {
buffer: wgpu::Buffer,
width: u32,
}
fn match_space_dims(image: &RgbaImage, transpose: bool) -> (u32, u32) {
match transpose {
false => (image.width(), image.height()),
true => (image.height(), image.width()),
}
}
fn preprocess_region(
context: &GpuContext,
image: &RgbaImage,
params: &SearchParams,
row_start: u32,
row_end: u32,
) -> ProcessedRegion {
let halo = if params.edges { 1 } else { 0 };
let rows = row_end - row_start;
let (bytes, uniform) = match params.transpose {
false => {
let width = image.width();
let first = row_start.saturating_sub(halo);
let last = (row_end + halo).min(image.height());
let byte_range = (first * width * 4) as usize..(last * width * 4) as usize;
(
Cow::Borrowed(&image.as_raw()[byte_range]),
PreprocessParams {
in_width: width,
in_height: last - first,
out_width: width,
out_height: rows,
halo_x: 0,
halo_y: row_start - first,
mode: params.edges as u32,
transpose: 0,
},
)
}
true => {
let width = image.width();
let first = row_start.saturating_sub(halo);
let last = (row_end + halo).min(width);
let raw = image.as_raw();
let mut bytes = Vec::with_capacity(((last - first) * image.height() * 4) as usize);
for y in 0..image.height() {
let row_base = (y * width) as usize;
bytes.extend_from_slice(&raw[(row_base + first as usize) * 4..(row_base + last as usize) * 4]);
}
(
Cow::Owned(bytes),
PreprocessParams {
in_width: last - first,
in_height: image.height(),
out_width: rows,
out_height: image.height(),
halo_x: row_start - first,
halo_y: 0,
mode: params.edges as u32,
transpose: 1,
},
)
}
};
let device = &context.device;
let input = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("preprocess-input"),
contents: &bytes,
usage: wgpu::BufferUsages::STORAGE,
});
let sad_width = match params.transpose {
false => uniform.out_width,
true => uniform.out_height,
};
let output = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("preprocess-output"),
size: rows as u64 * sad_width as u64 * 4,
usage: wgpu::BufferUsages::STORAGE,
mapped_at_creation: false,
});
let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("preprocess-params"),
contents: bytemuck::bytes_of(&uniform),
usage: wgpu::BufferUsages::UNIFORM,
});
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("preprocess-bind-group"),
layout: &context.preprocess_pipeline.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: uniform_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: input.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: output.as_entire_binding(),
},
],
});
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("preprocess-encoder"),
});
{
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
pass.set_pipeline(&context.preprocess_pipeline);
pass.set_bind_group(0, &bind_group, &[]);
pass.dispatch_workgroups(
uniform.out_width.div_ceil(WORKGROUP_DIM),
uniform.out_height.div_ceil(WORKGROUP_DIM),
1,
);
}
context.queue.submit([encoder.finish()]);
ProcessedRegion {
buffer: output,
width: sad_width,
}
}
async fn run_sad(
context: &GpuContext,
part1: &ProcessedRegion,
part2: &ProcessedRegion,
window: u32,
anchor_count: u32,
offset_start: i32,
offset_count: u32,
) -> Result<Vec<u32>, StitchError> {
let device = &context.device;
let workgroups_x = offset_count.div_ceil(WORKGROUP_DIM);
let workgroups_y = anchor_count.div_ceil(WORKGROUP_DIM);
let result_size = workgroups_x as u64 * workgroups_y as u64 * 2 * 4;
let uniform = SadParams {
width1: part1.width,
width2: part2.width,
window,
anchor_count,
offset_start,
offset_count,
_pad0: 0,
_pad1: 0,
};
let uniform_buffer = device.create_buffer_init(&wgpu::util::BufferInitDescriptor {
label: Some("sad-params"),
contents: bytemuck::bytes_of(&uniform),
usage: wgpu::BufferUsages::UNIFORM,
});
let results = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("sad-results"),
size: result_size,
usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
mapped_at_creation: false,
});
let staging = device.create_buffer(&wgpu::BufferDescriptor {
label: Some("sad-staging"),
size: result_size,
usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
mapped_at_creation: false,
});
let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
label: Some("sad-bind-group"),
layout: &context.sad_pipeline.get_bind_group_layout(0),
entries: &[
wgpu::BindGroupEntry {
binding: 0,
resource: uniform_buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 1,
resource: part1.buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 2,
resource: part2.buffer.as_entire_binding(),
},
wgpu::BindGroupEntry {
binding: 3,
resource: results.as_entire_binding(),
},
],
});
let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
label: Some("sad-encoder"),
});
{
let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
pass.set_pipeline(&context.sad_pipeline);
pass.set_bind_group(0, &bind_group, &[]);
pass.dispatch_workgroups(workgroups_x, workgroups_y, 1);
}
encoder.copy_buffer_to_buffer(&results, 0, &staging, 0, result_size);
context.queue.submit([encoder.finish()]);
context.read_buffer_u32(&staging).await
}
pub(crate) async fn find_best_overlap(
context: &GpuContext,
part1: &RgbaImage,
part2: &RgbaImage,
params: &SearchParams,
) -> Result<BestMatch, StitchError> {
find_best_overlap_slabbed(context, part1, part2, params, u32::MAX).await
}
async fn find_best_overlap_slabbed(
context: &GpuContext,
part1: &RgbaImage,
part2: &RgbaImage,
params: &SearchParams,
slab_cap: u32,
) -> Result<BestMatch, StitchError> {
let (width1, height1) = match_space_dims(part1, params.transpose);
let (width2, height2) = match_space_dims(part2, params.transpose);
let max_dim = width1.max(height1).max(width2).max(height2);
if max_dim > MAX_DISPATCH_THREADS {
return Err(StitchError::Gpu(format!(
"image dimension {max_dim} exceeds the supported maximum of {MAX_DISPATCH_THREADS}"
)));
}
let anchor_min = params.skip;
let anchor_max = (height1 - params.crop)
.checked_sub(params.window)
.ok_or(StitchError::NoCandidates)?;
if anchor_min > anchor_max {
return Err(StitchError::NoCandidates);
}
let offset_count = (params.offset_end - params.offset_start + 1).max(1) as u32;
let part2_processed =
preprocess_region(context, part2, params, params.crop, params.crop + params.window);
let budget_rows = (SLAB_BYTE_BUDGET / 4 / width1 as u64).min(u32::MAX as u64) as u32;
let anchors_per_slab = (budget_rows.saturating_sub(params.window) + 1)
.min((RESULT_CANDIDATE_BUDGET / offset_count as u64).min(u32::MAX as u64) as u32)
.min(MAX_DISPATCH_THREADS)
.min(slab_cap)
.max(1);
let mut best: Option<BestMatch> = None;
let mut slab_start = anchor_min;
while slab_start <= anchor_max {
let slab_end = slab_start
.saturating_add(anchors_per_slab - 1)
.min(anchor_max);
let anchor_count = slab_end - slab_start + 1;
let part1_processed = preprocess_region(
context,
part1,
params,
slab_start,
slab_end + params.window,
);
let results = run_sad(
context,
&part1_processed,
&part2_processed,
params.window,
anchor_count,
params.offset_start,
offset_count,
)
.await?;
for pair in results.chunks_exact(2) {
let (score, index) = (pair[0], pair[1]);
if score == u32::MAX {
continue;
}
let candidate = BestMatch {
score,
anchor: slab_start + index / offset_count,
offset: params.offset_start + (index % offset_count) as i32,
};
let is_better = match &best {
None => true,
Some(current) => {
(score, candidate.anchor, candidate.offset)
< (current.score, current.anchor, current.offset)
}
};
if is_better {
best = Some(candidate);
}
}
slab_start = slab_end + 1;
}
best.ok_or(StitchError::NoCandidates)
}
#[cfg(test)]
mod tests {
use super::*;
use image::Rgba;
fn lcg(seed: &mut u64) -> u8 {
*seed = seed
.wrapping_mul(6364136223846793005)
.wrapping_add(1442695040888963407);
(*seed >> 33) as u8
}
fn random_image(width: u32, height: u32, seed: u64) -> RgbaImage {
let mut state = seed;
RgbaImage::from_fn(width, height, |_, _| {
Rgba([lcg(&mut state), lcg(&mut state), lcg(&mut state), 255])
})
}
fn cpu_match_values(image: &RgbaImage, transpose: bool, edges: bool) -> (Vec<u32>, u32, u32) {
let (width, height) = (image.width() as i32, image.height() as i32);
let luma = |x: i32, y: i32| -> i32 {
let pixel = image.get_pixel(x.clamp(0, width - 1) as u32, y.clamp(0, height - 1) as u32);
((2126 * pixel[0] as u32 + 7152 * pixel[1] as u32 + 722 * pixel[2] as u32) / 10000)
as i32
};
let value = |x: i32, y: i32| -> u32 {
if !edges {
let pixel = image.get_pixel(x as u32, y as u32);
(pixel[0] as u32 + pixel[1] as u32 + pixel[2] as u32) / 3
} else {
let gx = (luma(x + 1, y - 1) + 2 * luma(x + 1, y) + luma(x + 1, y + 1))
- (luma(x - 1, y - 1) + 2 * luma(x - 1, y) + luma(x - 1, y + 1));
let gy = (luma(x - 1, y + 1) + 2 * luma(x, y + 1) + luma(x + 1, y + 1))
- (luma(x - 1, y - 1) + 2 * luma(x, y - 1) + luma(x + 1, y - 1));
(gx.abs() + gy.abs()) as u32 / 8
}
};
let (match_width, match_height) = match_space_dims(image, transpose);
let mut values = vec![0u32; (match_width * match_height) as usize];
for match_y in 0..match_height {
for match_x in 0..match_width {
let (x, y) = match transpose {
false => (match_x, match_y),
true => (match_y, match_x),
};
let mut pixel_value = value(x as i32, y as i32);
if image.get_pixel(x, y)[3] == 0 {
pixel_value |= 0x100;
}
values[(match_y * match_width + match_x) as usize] = pixel_value;
}
}
(values, match_width, match_height)
}
fn margin_cost(raw: u32) -> u32 {
match raw & 0x100 {
0 => raw & 0xff,
_ => 0,
}
}
fn diff_cost(a: u32, b: u32) -> u32 {
match (a | b) & 0x100 {
0 => (a & 0xff).abs_diff(b & 0xff),
_ => 255,
}
}
#[allow(clippy::too_many_arguments)]
fn cpu_sad(
values1: &[u32],
width1: u32,
values2: &[u32],
width2: u32,
window: u32,
crop: u32,
anchor: u32,
offset: i32,
) -> u32 {
let mut score = 0u32;
for i in 0..window {
let row1 = &values1[((anchor + i) * width1) as usize..][..width1 as usize];
let row2 = &values2[((crop + i) * width2) as usize..][..width2 as usize];
if offset >= 0 {
let margin1 = (offset as u32).min(width1);
let margin2 = (offset as u32).min(width2);
let overlap = (width1 - margin1).min(width2);
score += row1[..margin1 as usize].iter().map(|v| margin_cost(*v)).sum::<u32>();
score += row2[(width2 - margin2) as usize..]
.iter()
.map(|v| margin_cost(*v))
.sum::<u32>();
for x in 0..overlap as usize {
score += diff_cost(row1[offset as usize + x], row2[x]);
}
} else {
let margin1 = (offset.unsigned_abs()).min(width1);
let margin2 = (offset.unsigned_abs()).min(width2);
let overlap = width1.min(width2 - margin2);
score += row1[(width1 - margin1) as usize..]
.iter()
.map(|v| margin_cost(*v))
.sum::<u32>();
score += row2[..margin2 as usize].iter().map(|v| margin_cost(*v)).sum::<u32>();
for x in 0..overlap as usize {
score += diff_cost(row1[x], row2[offset.unsigned_abs() as usize + x]);
}
}
}
score
}
fn cpu_best(part1: &RgbaImage, part2: &RgbaImage, params: &SearchParams) -> BestMatch {
let (values1, width1, height1) = cpu_match_values(part1, params.transpose, params.edges);
let (values2, width2, _) = cpu_match_values(part2, params.transpose, params.edges);
let anchor_max = height1 - params.crop - params.window;
let mut best: Option<BestMatch> = None;
for anchor in params.skip..=anchor_max {
for offset in params.offset_start..=params.offset_end {
let score = cpu_sad(
&values1,
width1,
&values2,
width2,
params.window,
params.crop,
anchor,
offset,
);
let candidate = BestMatch {
score,
anchor,
offset,
};
let is_better = match &best {
None => true,
Some(current) => {
(score, anchor, offset) < (current.score, current.anchor, current.offset)
}
};
if is_better {
best = Some(candidate);
}
}
}
best.unwrap()
}
fn gpu_context() -> Option<super::super::SharedGpuContext> {
match pollster::block_on(super::super::shared_context()) {
Ok(context) => Some(context),
Err(StitchError::NoAdapter(details)) => {
eprintln!("skipping GPU test, no adapter: {details}");
None
}
Err(err) => panic!("failed to create GPU context: {err}"),
}
}
#[test]
fn gpu_matches_cpu_reference() {
let Some(context) = gpu_context() else { return };
for (case, transpose) in [(0u64, false), (1, true)] {
for edges in [false, true] {
let part1 = random_image(24, 31, 1000 + case);
let part2 = random_image(19, 26, 2000 + case);
let params = SearchParams {
transpose,
edges,
window: 3,
crop: 2,
skip: 1,
offset_start: -6,
offset_end: 6,
};
let gpu = pollster::block_on(find_best_overlap(context, &part1, &part2, ¶ms))
.unwrap();
let cpu = cpu_best(&part1, &part2, ¶ms);
assert_eq!(
(gpu.score, gpu.anchor, gpu.offset),
(cpu.score, cpu.anchor, cpu.offset),
"transpose={transpose} edges={edges}"
);
}
}
}
#[test]
fn gpu_matches_cpu_reference_with_transparent_padding() {
let Some(context) = gpu_context() else { return };
let mut part1 = random_image(30, 36, 77);
let part2 = random_image(24, 28, 78);
for y in 0..14 {
for x in 0..9 {
part1.put_pixel(x, y, image::Rgba([0, 0, 0, 0]));
}
}
for y in 20..36 {
part1.put_pixel(29, y, image::Rgba([0, 0, 0, 0]));
}
for edges in [false, true] {
let params = SearchParams {
transpose: false,
edges,
window: 3,
crop: 1,
skip: 0,
offset_start: -8,
offset_end: 8,
};
let gpu =
pollster::block_on(find_best_overlap(context, &part1, &part2, ¶ms)).unwrap();
let cpu = cpu_best(&part1, &part2, ¶ms);
assert_eq!(
(gpu.score, gpu.anchor, gpu.offset),
(cpu.score, cpu.anchor, cpu.offset),
"edges={edges}"
);
}
}
#[test]
fn gpu_matches_cpu_reference_with_wide_offsets_and_slabs() {
let Some(context) = gpu_context() else { return };
let part1 = random_image(20, 40, 42);
let part2 = random_image(26, 33, 43);
let params = SearchParams {
transpose: false,
edges: true,
window: 4,
crop: 0,
skip: 0,
offset_start: -25,
offset_end: 25,
};
let gpu = pollster::block_on(find_best_overlap_slabbed(
context, &part1, &part2, ¶ms, 4,
))
.unwrap();
let cpu = cpu_best(&part1, &part2, ¶ms);
assert_eq!(
(gpu.score, gpu.anchor, gpu.offset),
(cpu.score, cpu.anchor, cpu.offset)
);
}
}