From a4ff60233fce134e8f779ef001872dfd9a8f9923 Mon Sep 17 00:00:00 2001 From: skal Date: Sat, 21 Mar 2026 08:38:29 +0100 Subject: feat(cnn_v3): Phase 3 complete — WGSL U-Net inference shaders MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 5 compute shaders + cnn_v3/common snippet: enc0: Conv(20→4,3×3) + FiLM + ReLU full-res enc1: AvgPool + Conv(4→8,3×3) + FiLM + ReLU half-res bottleneck: AvgPool + Conv(8→8,1×1) + ReLU quarter-res dec1: NearestUp + cat(enc1) + Conv(16→4) + FiLM half-res dec0: NearestUp + cat(enc0) + Conv(8→4) + FiLM + Sigmoid full-res Parity rules: zero-pad conv, AvgPool down, NearestUp, FiLM after conv+bias, skip=concat, OIHW weights+bias layout. Matches PyTorch train_cnn_v3.py forward() exactly. Registered in workspaces/main/assets.txt + src/effects/shaders.cc. Weight layout + Params struct documented in cnn_v3/docs/HOWTO.md §7. Next: Phase 4 — C++ CNNv3Effect + FiLM uniform upload. Co-Authored-By: Claude Sonnet 4.6 --- cnn_v3/shaders/cnn_v3_bottleneck.wgsl | 73 +++++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 cnn_v3/shaders/cnn_v3_bottleneck.wgsl (limited to 'cnn_v3/shaders/cnn_v3_bottleneck.wgsl') diff --git a/cnn_v3/shaders/cnn_v3_bottleneck.wgsl b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl new file mode 100644 index 0000000..909fd41 --- /dev/null +++ b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl @@ -0,0 +1,73 @@ +// CNN v3 — Bottleneck +// AvgPool2x2(enc1) + Conv(8->8, 1x1) + ReLU (no FiLM) +// +// Input: enc1_tex (rgba32uint, 8xf16) half-res +// Output: bottleneck_out (rgba32uint, 8xf16) quarter-res (dispatch at quarter-res dims) +// +// Weight layout (f16, OIHW + bias): +// [0 .. 8*8*1) conv: w[out][in] (1x1 kernel) +// [64 .. +8) bias: b[out] + +#include "cnn_v3/common" + +const BN_IN: u32 = 8u; +const BN_OUT: u32 = 8u; + +struct Params { + weight_offset: u32, + _pad0: u32, + _pad1: u32, + _pad2: u32, +} + +@group(0) @binding(0) var enc1_tex: texture_2d; +@group(0) @binding(1) var weights: array; +@group(0) @binding(2) var params: Params; +@group(0) @binding(3) var bottleneck_out: texture_storage_2d; + +// Avg-pool 2x2 from enc1_tex at quarter-res coord qcoord. +// Returns zeros for OOB quarter-res coords (zero-padding for the 1x1 conv). +fn load_enc1_avg(qcoord: vec2i, half_dims: vec2i) -> array { + let quart_dims = half_dims / 2; + if (qcoord.x < 0 || qcoord.y < 0 || qcoord.x >= quart_dims.x || qcoord.y >= quart_dims.y) { + return array(0., 0., 0., 0., 0., 0., 0., 0.); + } + let base = qcoord * 2; + var s: array; + for (var dy: i32 = 0; dy < 2; dy++) { + for (var dx: i32 = 0; dx < 2; dx++) { + let hc = clamp(base + vec2i(dx, dy), vec2i(0), half_dims - vec2i(1)); + let f = unpack_8ch(enc1_tex, hc); + for (var i: u32 = 0u; i < BN_IN; i++) { s[i] += f[i]; } + } + } + for (var i: u32 = 0u; i < BN_IN; i++) { s[i] *= 0.25; } + return s; +} + +@compute @workgroup_size(8, 8) +fn bottleneck_main(@builtin(global_invocation_id) id: vec3u) { + let half_dims = vec2i(textureDimensions(enc1_tex)); + let quart_dims = half_dims / 2; + let coord = vec2i(id.xy); + if (coord.x >= quart_dims.x || coord.y >= quart_dims.y) { return; } + + let wo = params.weight_offset; + let feat = load_enc1_avg(coord, half_dims); + var out: array; + + for (var o: u32 = 0u; o < BN_OUT; o++) { + var sum = get_w(wo, BN_OUT * BN_IN + o); // bias (1x1 kernel: no spatial idx) + for (var i: u32 = 0u; i < BN_IN; i++) { + sum += get_w(wo, o * BN_IN + i) * feat[i]; + } + out[o] = max(0.0, sum); + } + + textureStore(bottleneck_out, coord, vec4u( + pack2x16float(vec2f(out[0], out[1])), + pack2x16float(vec2f(out[2], out[3])), + pack2x16float(vec2f(out[4], out[5])), + pack2x16float(vec2f(out[6], out[7])) + )); +} -- cgit v1.2.3