diff options
| author | skal <pascal.massimino@gmail.com> | 2026-03-21 08:38:29 +0100 |
|---|---|---|
| committer | skal <pascal.massimino@gmail.com> | 2026-03-21 08:38:29 +0100 |
| commit | a4ff60233fce134e8f779ef001872dfd9a8f9923 (patch) | |
| tree | 3a5466273ecb42269b4d6443c893c61b84ee7d93 /cnn_v3/shaders/cnn_v3_enc1.wgsl | |
| parent | 4d055080d2ab4b674d5f0fd611ea051e87454a31 (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_enc1.wgsl')
| -rw-r--r-- | cnn_v3/shaders/cnn_v3_enc1.wgsl | 88 |
1 files changed, 88 insertions, 0 deletions
diff --git a/cnn_v3/shaders/cnn_v3_enc1.wgsl b/cnn_v3/shaders/cnn_v3_enc1.wgsl new file mode 100644 index 0000000..23e485d --- /dev/null +++ b/cnn_v3/shaders/cnn_v3_enc1.wgsl @@ -0,0 +1,88 @@ +// CNN v3 — Encoder level 1 +// AvgPool2x2(enc0) + Conv(4->8, 3x3, zero-pad) + FiLM + ReLU +// +// Input: enc0_tex (rgba16float, 4ch) full-res +// Output: enc1_out (rgba32uint, 8xf16) half-res (dispatch at half-res dims) +// +// Weight layout (f16, OIHW + bias): +// [0 .. 4*8*9) conv: w[out][in][ky][kx] +// [288 .. +8) bias: b[out] + +#include "cnn_v3/common" + +const ENC1_IN: u32 = 4u; +const ENC1_OUT: u32 = 8u; + +struct Params { + weight_offset: u32, + _pad: vec3u, + gamma_lo: vec4f, // FiLM gamma ch 0-3 + gamma_hi: vec4f, // FiLM gamma ch 4-7 + beta_lo: vec4f, // FiLM beta ch 0-3 + beta_hi: vec4f, // FiLM beta ch 4-7 +} + +@group(0) @binding(0) var enc0_tex: texture_2d<f32>; +@group(0) @binding(1) var<storage, read> weights: array<u32>; +@group(0) @binding(2) var<uniform> params: Params; +@group(0) @binding(3) var enc1_out: texture_storage_2d<rgba32uint, write>; + +fn film_gamma(o: u32) -> f32 { + if (o < 4u) { return params.gamma_lo[o]; } + return params.gamma_hi[o - 4u]; +} +fn film_beta(o: u32) -> f32 { + if (o < 4u) { return params.beta_lo[o]; } + return params.beta_hi[o - 4u]; +} + +// Avg-pool 2x2 from enc0_tex at half-res coord hcoord. +// Returns zeros for OOB half-res coords (zero-padding for the conv). +fn load_enc0_avg(hcoord: vec2i, full_dims: vec2i) -> array<f32, 4> { + let half_dims = full_dims / 2; + if (hcoord.x < 0 || hcoord.y < 0 || hcoord.x >= half_dims.x || hcoord.y >= half_dims.y) { + return array<f32, 4>(0., 0., 0., 0.); + } + let base = hcoord * 2; + var s = vec4f(0.); + for (var dy: i32 = 0; dy < 2; dy++) { + for (var dx: i32 = 0; dx < 2; dx++) { + let fc = clamp(base + vec2i(dx, dy), vec2i(0), full_dims - vec2i(1)); + s += textureLoad(enc0_tex, fc, 0); + } + } + let avg = s * 0.25; + return array<f32, 4>(avg.x, avg.y, avg.z, avg.w); +} + +@compute @workgroup_size(8, 8) +fn enc1_main(@builtin(global_invocation_id) id: vec3u) { + let full_dims = vec2i(textureDimensions(enc0_tex)); + let half_dims = full_dims / 2; + let coord = vec2i(id.xy); + if (coord.x >= half_dims.x || coord.y >= half_dims.y) { return; } + + let wo = params.weight_offset; + var out: array<f32, ENC1_OUT>; + + for (var o: u32 = 0u; o < ENC1_OUT; o++) { + var sum = get_w(wo, ENC1_OUT * ENC1_IN * 9u + o); // bias + for (var ky: i32 = -1; ky <= 1; ky++) { + for (var kx: i32 = -1; kx <= 1; kx++) { + let feat = load_enc0_avg(coord + vec2i(kx, ky), full_dims); + let ki = u32(ky + 1) * 3u + u32(kx + 1); + for (var i: u32 = 0u; i < ENC1_IN; i++) { + sum += get_w(wo, o * ENC1_IN * 9u + i * 9u + ki) * feat[i]; + } + } + } + out[o] = max(0.0, film_gamma(o) * sum + film_beta(o)); + } + + textureStore(enc1_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])) + )); +} |
