summaryrefslogtreecommitdiff
path: root/cnn_v3/shaders/cnn_v3_enc1.wgsl
diff options
context:
space:
mode:
authorskal <pascal.massimino@gmail.com>2026-03-26 07:03:01 +0100
committerskal <pascal.massimino@gmail.com>2026-03-26 07:03:01 +0100
commit8f14bdd66cb002b2f89265b2a578ad93249089c9 (patch)
tree2ccdb3939b673ebc3a5df429160631240239cee2 /cnn_v3/shaders/cnn_v3_enc1.wgsl
parent4ca498277b033ae10134045dae9c8c249a8d2b2b (diff)
feat(cnn_v3): upgrade architecture to enc_channels=[8,16]
Double encoder capacity: enc0 4→8ch, enc1 8→16ch, bottleneck 16→16ch, dec1 32→8ch, dec0 16→4ch. Total weights 2476→7828 f16 (~15.3 KB). FiLM MLP output 40→72 params (L1: 16×40→16×72). 16-ch textures split into _lo/_hi rgba32uint pairs (enc1, bottleneck). enc0 and dec1 textures changed from rgba16float to rgba32uint (8ch). GBUF_RGBA32UINT node gains CopySrc for parity test readback. - WGSL shaders: all 5 passes rewritten for new channel counts - C++ CNNv3Effect: new weight offsets/sizes, 8ch uniform structs - Web tool (shaders.js + tester.js): matching texture formats and bindings - Parity test: readback_rgba32uint_8ch helper, updated vector counts - Training scripts: default enc_channels=[8,16], updated docstrings - Docs + architecture PNG regenerated handoff(Gemini): CNN v3 [8,16] upgrade complete. All code, tests, web tool, training scripts, and docs updated. Next: run training pass.
Diffstat (limited to 'cnn_v3/shaders/cnn_v3_enc1.wgsl')
-rw-r--r--cnn_v3/shaders/cnn_v3_enc1.wgsl75
1 files changed, 45 insertions, 30 deletions
diff --git a/cnn_v3/shaders/cnn_v3_enc1.wgsl b/cnn_v3/shaders/cnn_v3_enc1.wgsl
index 23e485d..eb41279 100644
--- a/cnn_v3/shaders/cnn_v3_enc1.wgsl
+++ b/cnn_v3/shaders/cnn_v3_enc1.wgsl
@@ -1,58 +1,67 @@
// CNN v3 — Encoder level 1
-// AvgPool2x2(enc0) + Conv(4->8, 3x3, zero-pad) + FiLM + ReLU
+// AvgPool2x2(enc0) + Conv(8->16, 3x3, zero-pad) + FiLM + ReLU
//
-// Input: enc0_tex (rgba16float, 4ch) full-res
-// Output: enc1_out (rgba32uint, 8xf16) half-res (dispatch at half-res dims)
+// Input: enc0_tex (rgba32uint, 8xf16) full-res
+// Output: enc1_out_lo (rgba32uint, 8xf16) half-res ch 0-7
+// enc1_out_hi (rgba32uint, 8xf16) half-res ch 8-15
//
// Weight layout (f16, OIHW + bias):
-// [0 .. 4*8*9) conv: w[out][in][ky][kx]
-// [288 .. +8) bias: b[out]
+// [0 .. 8*16*9) conv: w[out][in][ky][kx]
+// [1152 .. +16) bias: b[out]
#include "cnn_v3/common"
-const ENC1_IN: u32 = 4u;
-const ENC1_OUT: u32 = 8u;
+const ENC1_IN: u32 = 8u;
+const ENC1_OUT: u32 = 16u;
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
+ gamma_0: vec4f,
+ gamma_1: vec4f,
+ gamma_2: vec4f,
+ gamma_3: vec4f,
+ beta_0: vec4f,
+ beta_1: vec4f,
+ beta_2: vec4f,
+ beta_3: vec4f,
}
-@group(0) @binding(0) var enc0_tex: texture_2d<f32>;
+@group(0) @binding(0) var enc0_tex: texture_2d<u32>;
@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>;
+@group(0) @binding(3) var enc1_out_lo: texture_storage_2d<rgba32uint, write>;
+@group(0) @binding(4) var enc1_out_hi: 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];
+ if (o < 4u) { return params.gamma_0[o]; }
+ if (o < 8u) { return params.gamma_1[o - 4u]; }
+ if (o < 12u) { return params.gamma_2[o - 8u]; }
+ return params.gamma_3[o - 12u];
}
fn film_beta(o: u32) -> f32 {
- if (o < 4u) { return params.beta_lo[o]; }
- return params.beta_hi[o - 4u];
+ if (o < 4u) { return params.beta_0[o]; }
+ if (o < 8u) { return params.beta_1[o - 4u]; }
+ if (o < 12u) { return params.beta_2[o - 8u]; }
+ return params.beta_3[o - 12u];
}
-// 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> {
+fn load_enc0_avg(hcoord: vec2i, full_dims: vec2i) -> array<f32, 8> {
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.);
+ return array<f32, 8>(0., 0., 0., 0., 0., 0., 0., 0.);
}
let base = hcoord * 2;
- var s = vec4f(0.);
+ var s: array<f32, 8>;
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 f = unpack_8ch(enc0_tex, fc);
+ for (var i: u32 = 0u; i < 8u; i++) { s[i] += f[i]; }
}
}
- let avg = s * 0.25;
- return array<f32, 4>(avg.x, avg.y, avg.z, avg.w);
+ for (var i: u32 = 0u; i < 8u; i++) { s[i] *= 0.25; }
+ return s;
}
@compute @workgroup_size(8, 8)
@@ -79,10 +88,16 @@ fn enc1_main(@builtin(global_invocation_id) id: vec3u) {
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]))
+ textureStore(enc1_out_lo, 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]))
+ ));
+ textureStore(enc1_out_hi, coord, vec4u(
+ pack2x16float(vec2f(out[8], out[9])),
+ pack2x16float(vec2f(out[10], out[11])),
+ pack2x16float(vec2f(out[12], out[13])),
+ pack2x16float(vec2f(out[14], out[15]))
));
}