summaryrefslogtreecommitdiff
path: root/cnn_v3/shaders/cnn_v3_dec1.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_dec1.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_dec1.wgsl')
-rw-r--r--cnn_v3/shaders/cnn_v3_dec1.wgsl85
1 files changed, 54 insertions, 31 deletions
diff --git a/cnn_v3/shaders/cnn_v3_dec1.wgsl b/cnn_v3/shaders/cnn_v3_dec1.wgsl
index 28ae3dc..fadea3b 100644
--- a/cnn_v3/shaders/cnn_v3_dec1.wgsl
+++ b/cnn_v3/shaders/cnn_v3_dec1.wgsl
@@ -1,53 +1,71 @@
// CNN v3 — Decoder level 1
-// NearestUp2x(bottleneck) + cat(enc1_skip) -> Conv(16->4, 3x3, zero-pad) + FiLM + ReLU
+// NearestUp2x(bottleneck) + cat(enc1_skip) -> Conv(32->8, 3x3) + FiLM + ReLU
//
-// Inputs: bottleneck_tex (rgba32uint, 8xf16) quarter-res
-// enc1_tex (rgba32uint, 8xf16) half-res (skip connection)
-// Output: dec1_out (rgba16float, 4ch) half-res (dispatch at half-res dims)
+// Inputs: bn_tex_lo (rgba32uint, 8xf16) quarter-res ch 0-7
+// bn_tex_hi (rgba32uint, 8xf16) quarter-res ch 8-15
+// enc1_tex_lo (rgba32uint, 8xf16) half-res skip ch 0-7
+// enc1_tex_hi (rgba32uint, 8xf16) half-res skip ch 8-15
+// Output: dec1_out (rgba32uint, 8xf16) half-res
//
// Weight layout (f16, OIHW + bias):
-// [0 .. 16*4*9) conv: w[out][in][ky][kx] (in=16: 8 bottleneck + 8 enc1 skip)
-// [576 .. +4) bias: b[out]
+// [0 .. 32*8*9) conv: w[out][in][ky][kx] (in=32: 16 bn + 16 enc1 skip)
+// [2304 .. +8) bias: b[out]
#include "cnn_v3/common"
-const DEC1_IN: u32 = 16u;
-const DEC1_OUT: u32 = 4u;
+const DEC1_IN: u32 = 32u;
+const DEC1_OUT: u32 = 8u;
struct Params {
weight_offset: u32,
_pad: vec3u,
- gamma: vec4f,
- beta: vec4f,
+ gamma_lo: vec4f,
+ gamma_hi: vec4f,
+ beta_lo: vec4f,
+ beta_hi: vec4f,
}
-@group(0) @binding(0) var bottleneck_tex: texture_2d<u32>;
-@group(0) @binding(1) var enc1_tex: texture_2d<u32>;
-@group(0) @binding(2) var<storage, read> weights: array<u32>;
-@group(0) @binding(3) var<uniform> params: Params;
-@group(0) @binding(4) var dec1_out: texture_storage_2d<rgba16float, write>;
+@group(0) @binding(0) var bn_tex_lo: texture_2d<u32>;
+@group(0) @binding(1) var bn_tex_hi: texture_2d<u32>;
+@group(0) @binding(2) var enc1_tex_lo: texture_2d<u32>;
+@group(0) @binding(3) var enc1_tex_hi: texture_2d<u32>;
+@group(0) @binding(4) var<storage, read> weights: array<u32>;
+@group(0) @binding(5) var<uniform> params: Params;
+@group(0) @binding(6) var dec1_out: texture_storage_2d<rgba32uint, write>;
-// Load 16 concatenated channels at half-res coord hcoord:
-// ch 0-7: bottleneck nearest-up (bottleneck_tex[hcoord/2])
-// ch 8-15: enc1 skip (enc1_tex[hcoord])
-// Returns zeros for OOB hcoord (zero-padding for the conv).
-fn load_dec1_concat(hcoord: vec2i, half_dims: vec2i) -> array<f32, 16> {
+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];
+}
+
+// Load 32ch: [bn_nearest_up(16ch), enc1_skip(16ch)]
+fn load_dec1_concat(hcoord: vec2i, half_dims: vec2i) -> array<f32, 32> {
+ var r: array<f32, 32>;
if (hcoord.x < 0 || hcoord.y < 0 || hcoord.x >= half_dims.x || hcoord.y >= half_dims.y) {
- return array<f32, 16>(0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.);
+ return r;
}
let quart_dims = half_dims / 2;
- let qc = clamp(hcoord / 2, vec2i(0), quart_dims - vec2i(1));
- let b = unpack_8ch(bottleneck_tex, qc);
- let s = unpack_8ch(enc1_tex, hcoord);
- return array<f32, 16>(
- b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7],
- s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7]
- );
+ let qc = clamp(hcoord / 2, vec2i(0), quart_dims - vec2i(1));
+ let blo = unpack_8ch(bn_tex_lo, qc);
+ let bhi = unpack_8ch(bn_tex_hi, qc);
+ let slo = unpack_8ch(enc1_tex_lo, hcoord);
+ let shi = unpack_8ch(enc1_tex_hi, hcoord);
+ for (var i: u32 = 0u; i < 8u; i++) {
+ r[i] = blo[i];
+ r[i + 8u] = bhi[i];
+ r[i + 16u] = slo[i];
+ r[i + 24u] = shi[i];
+ }
+ return r;
}
@compute @workgroup_size(8, 8)
fn dec1_main(@builtin(global_invocation_id) id: vec3u) {
- let half_dims = vec2i(textureDimensions(enc1_tex));
+ let half_dims = vec2i(textureDimensions(enc1_tex_lo));
let coord = vec2i(id.xy);
if (coord.x >= half_dims.x || coord.y >= half_dims.y) { return; }
@@ -65,8 +83,13 @@ fn dec1_main(@builtin(global_invocation_id) id: vec3u) {
}
}
}
- out[o] = max(0.0, params.gamma[o] * sum + params.beta[o]);
+ out[o] = max(0.0, film_gamma(o) * sum + film_beta(o));
}
- textureStore(dec1_out, coord, vec4f(out[0], out[1], out[2], out[3]));
+ textureStore(dec1_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]))
+ ));
}