summaryrefslogtreecommitdiff
path: root/cnn_v3/shaders/cnn_v3_dec0.wgsl
diff options
context:
space:
mode:
authorskal <pascal.massimino@gmail.com>2026-03-21 08:38:29 +0100
committerskal <pascal.massimino@gmail.com>2026-03-21 08:38:29 +0100
commita4ff60233fce134e8f779ef001872dfd9a8f9923 (patch)
tree3a5466273ecb42269b4d6443c893c61b84ee7d93 /cnn_v3/shaders/cnn_v3_dec0.wgsl
parent4d055080d2ab4b674d5f0fd611ea051e87454a31 (diff)
feat(cnn_v3): Phase 3 complete — WGSL U-Net inference shaders
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 <noreply@anthropic.com>
Diffstat (limited to 'cnn_v3/shaders/cnn_v3_dec0.wgsl')
-rw-r--r--cnn_v3/shaders/cnn_v3_dec0.wgsl73
1 files changed, 73 insertions, 0 deletions
diff --git a/cnn_v3/shaders/cnn_v3_dec0.wgsl b/cnn_v3/shaders/cnn_v3_dec0.wgsl
new file mode 100644
index 0000000..7a4e7c9
--- /dev/null
+++ b/cnn_v3/shaders/cnn_v3_dec0.wgsl
@@ -0,0 +1,73 @@
+// CNN v3 — Decoder level 0 + output
+// NearestUp2x(dec1) + cat(enc0_skip) -> Conv(8->4, 3x3, zero-pad) + FiLM + ReLU + Sigmoid
+//
+// Inputs: dec1_tex (rgba16float, 4ch) half-res
+// enc0_tex (rgba16float, 4ch) full-res (skip connection)
+// Output: output_tex (rgba16float, 4ch) full-res (dispatch at full-res dims)
+//
+// Weight layout (f16, OIHW + bias):
+// [0 .. 8*4*9) conv: w[out][in][ky][kx] (in=8: 4 dec1 + 4 enc0 skip)
+// [288 .. +4) bias: b[out]
+//
+// Parity note: sigmoid applied directly to dec0 output (matches train_cnn_v3.py forward()).
+
+#include "cnn_v3/common"
+
+const DEC0_IN: u32 = 8u;
+const DEC0_OUT: u32 = 4u;
+
+struct Params {
+ weight_offset: u32,
+ _pad: vec3u,
+ gamma: vec4f,
+ beta: vec4f,
+}
+
+@group(0) @binding(0) var dec1_tex: texture_2d<f32>;
+@group(0) @binding(1) var enc0_tex: texture_2d<f32>;
+@group(0) @binding(2) var<storage, read> weights: array<u32>;
+@group(0) @binding(3) var<uniform> params: Params;
+@group(0) @binding(4) var output_tex: texture_storage_2d<rgba16float, write>;
+
+// Load 8 concatenated channels at full-res coord:
+// ch 0-3: dec1 nearest-up (dec1_tex[coord/2])
+// ch 4-7: enc0 skip (enc0_tex[coord])
+// Returns zeros for OOB coord (zero-padding for the conv).
+fn load_dec0_concat(coord: vec2i, full_dims: vec2i) -> array<f32, 8> {
+ if (coord.x < 0 || coord.y < 0 || coord.x >= full_dims.x || coord.y >= full_dims.y) {
+ return array<f32, 8>(0., 0., 0., 0., 0., 0., 0., 0.);
+ }
+ let half_dims = vec2i(textureDimensions(dec1_tex));
+ let hc = clamp(coord / 2, vec2i(0), half_dims - vec2i(1));
+ let d = textureLoad(dec1_tex, hc, 0);
+ let e = textureLoad(enc0_tex, coord, 0);
+ return array<f32, 8>(d.x, d.y, d.z, d.w, e.x, e.y, e.z, e.w);
+}
+
+@compute @workgroup_size(8, 8)
+fn dec0_main(@builtin(global_invocation_id) id: vec3u) {
+ let full_dims = vec2i(textureDimensions(enc0_tex));
+ let coord = vec2i(id.xy);
+ if (coord.x >= full_dims.x || coord.y >= full_dims.y) { return; }
+
+ let wo = params.weight_offset;
+ var out: array<f32, DEC0_OUT>;
+
+ for (var o: u32 = 0u; o < DEC0_OUT; o++) {
+ var sum = get_w(wo, DEC0_OUT * DEC0_IN * 9u + o); // bias
+ for (var ky: i32 = -1; ky <= 1; ky++) {
+ for (var kx: i32 = -1; kx <= 1; kx++) {
+ let feat = load_dec0_concat(coord + vec2i(kx, ky), full_dims);
+ let ki = u32(ky + 1) * 3u + u32(kx + 1);
+ for (var i: u32 = 0u; i < DEC0_IN; i++) {
+ sum += get_w(wo, o * DEC0_IN * 9u + i * 9u + ki) * feat[i];
+ }
+ }
+ }
+ // FiLM + ReLU + Sigmoid (matches training forward())
+ let v = max(0.0, params.gamma[o] * sum + params.beta[o]);
+ out[o] = 1.0 / (1.0 + exp(-v));
+ }
+
+ textureStore(output_tex, coord, vec4f(out[0], out[1], out[2], out[3]));
+}