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_enc0.wgsl | 75 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 75 insertions(+) create mode 100644 cnn_v3/shaders/cnn_v3_enc0.wgsl (limited to 'cnn_v3/shaders/cnn_v3_enc0.wgsl') diff --git a/cnn_v3/shaders/cnn_v3_enc0.wgsl b/cnn_v3/shaders/cnn_v3_enc0.wgsl new file mode 100644 index 0000000..f52a167 --- /dev/null +++ b/cnn_v3/shaders/cnn_v3_enc0.wgsl @@ -0,0 +1,75 @@ +// CNN v3 — Encoder level 0 +// Conv(20->4, 3x3, zero-pad) + FiLM + ReLU +// +// Input: feat_tex0 (rgba32uint, 8xf16), feat_tex1 (rgba32uint, 12xu8) full-res +// Output: enc0_out (rgba16float, 4ch) full-res +// +// Weight layout (f16, OIHW + bias): +// [0 .. 20*4*9) conv: w[out][in][ky][kx] +// [720 .. +4) bias: b[out] + +#include "cnn_v3/common" + +const ENC0_IN: u32 = 20u; +const ENC0_OUT: u32 = 4u; + +struct Params { + weight_offset: u32, + _pad: vec3u, + gamma: vec4f, + beta: vec4f, +} + +@group(0) @binding(0) var feat_tex0: texture_2d; +@group(0) @binding(1) var feat_tex1: texture_2d; +@group(0) @binding(2) var weights: array; +@group(0) @binding(3) var params: Params; +@group(0) @binding(4) var enc0_out: texture_storage_2d; + +// Unpack all 20 feature channels at coord. Returns zeros for OOB (zero-padding). +fn load_feat(coord: vec2i, dims: vec2i) -> array { + if (coord.x < 0 || coord.y < 0 || coord.x >= dims.x || coord.y >= dims.y) { + return array(0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.); + } + let t0 = textureLoad(feat_tex0, coord, 0); + let t1 = textureLoad(feat_tex1, coord, 0); + let a = unpack2x16float(t0.x); + let b = unpack2x16float(t0.y); + let c = unpack2x16float(t0.z); + let d = unpack2x16float(t0.w); + let e = unpack4x8unorm(t1.x); + let f = unpack4x8unorm(t1.y); + let g = unpack4x8unorm(t1.z); + return array( + a.x, a.y, b.x, b.y, c.x, c.y, d.x, d.y, + e.x, e.y, e.z, e.w, + f.x, f.y, f.z, f.w, + g.x, g.y, g.z, g.w + ); +} + +@compute @workgroup_size(8, 8) +fn enc0_main(@builtin(global_invocation_id) id: vec3u) { + let coord = vec2i(id.xy); + let dims = vec2i(textureDimensions(feat_tex0)); + if (coord.x >= dims.x || coord.y >= dims.y) { return; } + + let wo = params.weight_offset; + var out: array; + + for (var o: u32 = 0u; o < ENC0_OUT; o++) { + var sum = get_w(wo, ENC0_OUT * ENC0_IN * 9u + o); // bias + for (var ky: i32 = -1; ky <= 1; ky++) { + for (var kx: i32 = -1; kx <= 1; kx++) { + let feat = load_feat(coord + vec2i(kx, ky), dims); + let ki = u32(ky + 1) * 3u + u32(kx + 1); + for (var i: u32 = 0u; i < ENC0_IN; i++) { + sum += get_w(wo, o * ENC0_IN * 9u + i * 9u + ki) * feat[i]; + } + } + } + out[o] = max(0.0, params.gamma[o] * sum + params.beta[o]); + } + + textureStore(enc0_out, coord, vec4f(out[0], out[1], out[2], out[3])); +} -- cgit v1.2.3