struct SpdParams {
src_size: vec2<u32>,
mip_count: u32,
_padding: u32,
}
@group(0) @binding(0) var src_texture: texture_2d<f32>;
@group(0) @binding(1) var dst_mip0: texture_storage_2d<r32float, write>;
@group(0) @binding(2) var dst_mip1: texture_storage_2d<r32float, write>;
@group(0) @binding(3) var dst_mip2: texture_storage_2d<r32float, write>;
@group(0) @binding(4) var dst_mip3: texture_storage_2d<r32float, write>;
@group(0) @binding(5) var<uniform> params: SpdParams;
var<workgroup> lds: array<f32, 4096>;
fn reduce4(v0: f32, v1: f32, v2: f32, v3: f32) -> f32 {
return min(min(v0, v1), min(v2, v3));
}
fn load_from_source(coord: vec2<u32>) -> f32 {
let clamped = min(coord, params.src_size - vec2(1u));
return textureLoad(src_texture, clamped, 0).r;
}
@compute @workgroup_size(256, 1, 1)
fn main(
@builtin(local_invocation_index) local_index: u32,
@builtin(workgroup_id) workgroup_id: vec3<u32>,
) {
let tile_x = workgroup_id.x;
let tile_y = workgroup_id.y;
let base_x = tile_x * 64u;
let base_y = tile_y * 64u;
for (var load_index = local_index; load_index < 4096u; load_index += 256u) {
let local_x = load_index % 64u;
let local_y = load_index / 64u;
let src_coord = vec2(base_x + local_x, base_y + local_y);
lds[load_index] = load_from_source(src_coord);
}
workgroupBarrier();
for (var store_index = local_index; store_index < 4096u; store_index += 256u) {
let local_x = store_index % 64u;
let local_y = store_index / 64u;
let dst_coord = vec2(base_x + local_x, base_y + local_y);
let mip0_size = params.src_size;
if dst_coord.x < mip0_size.x && dst_coord.y < mip0_size.y {
textureStore(dst_mip0, dst_coord, vec4(lds[store_index], 0.0, 0.0, 0.0));
}
}
if params.mip_count <= 1u {
return;
}
workgroupBarrier();
// Reduce mip0 -> mip1 into private memory first so that no thread writes
// to LDS while another thread is still reading the higher source indices.
// Writing in place during the loop would race: iter 0 reads indices up to
// ~1088 while iter 2 and iter 3 write indices 512..1024 with no barrier
// between them.
var mip1_reduced: array<f32, 4>;
for (var iter = 0u; iter < 4u; iter++) {
let reduce_index = local_index + iter * 256u;
if reduce_index < 1024u {
let local_x = reduce_index % 32u;
let local_y = reduce_index / 32u;
let src_x = local_x * 2u;
let src_y = local_y * 2u;
let idx00 = src_y * 64u + src_x;
let idx10 = src_y * 64u + src_x + 1u;
let idx01 = (src_y + 1u) * 64u + src_x;
let idx11 = (src_y + 1u) * 64u + src_x + 1u;
mip1_reduced[iter] = reduce4(lds[idx00], lds[idx10], lds[idx01], lds[idx11]);
}
}
workgroupBarrier();
for (var iter = 0u; iter < 4u; iter++) {
let reduce_index = local_index + iter * 256u;
if reduce_index < 1024u {
lds[reduce_index] = mip1_reduced[iter];
}
}
workgroupBarrier();
for (var store_index = local_index; store_index < 1024u; store_index += 256u) {
let local_x = store_index % 32u;
let local_y = store_index / 32u;
let dst_coord = vec2((base_x >> 1u) + local_x, (base_y >> 1u) + local_y);
let mip1_size = max(vec2(1u), params.src_size >> vec2(1u));
if dst_coord.x < mip1_size.x && dst_coord.y < mip1_size.y {
textureStore(dst_mip1, dst_coord, vec4(lds[store_index], 0.0, 0.0, 0.0));
}
}
if params.mip_count <= 2u {
return;
}
workgroupBarrier();
// Same race-free split for mip1 -> mip2: thread T reads up to lds[idx11]
// where idx11 may be in [0, 1024), and thread S writes to lds[S] where
// S in [0, 256), so reads of indices < 256 collide with writes from the
// very same pass without the staging.
var mip2_reduced: f32 = 0.0;
if local_index < 256u {
let local_x = local_index % 16u;
let local_y = local_index / 16u;
let src_x = local_x * 2u;
let src_y = local_y * 2u;
let idx00 = src_y * 32u + src_x;
let idx10 = src_y * 32u + src_x + 1u;
let idx01 = (src_y + 1u) * 32u + src_x;
let idx11 = (src_y + 1u) * 32u + src_x + 1u;
mip2_reduced = reduce4(lds[idx00], lds[idx10], lds[idx01], lds[idx11]);
}
workgroupBarrier();
if local_index < 256u {
lds[local_index] = mip2_reduced;
}
workgroupBarrier();
if local_index < 256u {
let local_x = local_index % 16u;
let local_y = local_index / 16u;
let dst_coord = vec2((base_x >> 2u) + local_x, (base_y >> 2u) + local_y);
let mip2_size = max(vec2(1u), params.src_size >> vec2(2u));
if dst_coord.x < mip2_size.x && dst_coord.y < mip2_size.y {
textureStore(dst_mip2, dst_coord, vec4(lds[local_index], 0.0, 0.0, 0.0));
}
}
if params.mip_count <= 3u {
return;
}
workgroupBarrier();
var mip3_reduced: f32 = 0.0;
if local_index < 64u {
let local_x = local_index % 8u;
let local_y = local_index / 8u;
let src_x = local_x * 2u;
let src_y = local_y * 2u;
let idx00 = src_y * 16u + src_x;
let idx10 = src_y * 16u + src_x + 1u;
let idx01 = (src_y + 1u) * 16u + src_x;
let idx11 = (src_y + 1u) * 16u + src_x + 1u;
mip3_reduced = reduce4(lds[idx00], lds[idx10], lds[idx01], lds[idx11]);
}
workgroupBarrier();
if local_index < 64u {
lds[local_index] = mip3_reduced;
}
workgroupBarrier();
if local_index < 64u {
let local_x = local_index % 8u;
let local_y = local_index / 8u;
let dst_coord = vec2((base_x >> 3u) + local_x, (base_y >> 3u) + local_y);
let mip3_size = max(vec2(1u), params.src_size >> vec2(3u));
if dst_coord.x < mip3_size.x && dst_coord.y < mip3_size.y {
textureStore(dst_mip3, dst_coord, vec4(lds[local_index], 0.0, 0.0, 0.0));
}
}
}