diff options
Diffstat (limited to 'cnn_v3/shaders/cnn_v3_dec0.wgsl')
| -rw-r--r-- | cnn_v3/shaders/cnn_v3_dec0.wgsl | 43 |
1 files changed, 21 insertions, 22 deletions
diff --git a/cnn_v3/shaders/cnn_v3_dec0.wgsl b/cnn_v3/shaders/cnn_v3_dec0.wgsl index a2a70ac..617b5a2 100644 --- a/cnn_v3/shaders/cnn_v3_dec0.wgsl +++ b/cnn_v3/shaders/cnn_v3_dec0.wgsl @@ -1,19 +1,17 @@ // CNN v3 — Decoder level 0 + output -// NearestUp2x(dec1) + cat(enc0_skip) -> Conv(8->4, 3x3, zero-pad) + FiLM + ReLU + Sigmoid +// NearestUp2x(dec1) + cat(enc0_skip) -> Conv(16->4, 3x3) + 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) +// Inputs: dec1_tex (rgba32uint, 8xf16) half-res +// enc0_tex (rgba32uint, 8xf16) full-res (skip connection) +// Output: output_tex (rgba16float, 4ch) full-res // // 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 after FiLM+ReLU, not after raw conv (matches train_cnn_v3.py). +// [0 .. 16*4*9) conv: w[out][in][ky][kx] (in=16: 8 dec1 + 8 enc0 skip) +// [576 .. +4) bias: b[out] #include "cnn_v3/common" -const DEC0_IN: u32 = 8u; +const DEC0_IN: u32 = 16u; const DEC0_OUT: u32 = 4u; struct Params { @@ -23,25 +21,27 @@ struct Params { 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(0) var dec1_tex: texture_2d<u32>; +@group(0) @binding(1) var enc0_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 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> { +// Load 16ch: ch 0-7 from dec1 nearest-up, ch 8-15 from enc0 skip. +fn load_dec0_concat(coord: vec2i, full_dims: vec2i) -> array<f32, 16> { + var r: array<f32, 16>; 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.); + return r; } 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); + let hc = clamp(coord / 2, vec2i(0), half_dims - vec2i(1)); + let d = unpack_8ch(dec1_tex, hc); + let e = unpack_8ch(enc0_tex, coord); + for (var i: u32 = 0u; i < 8u; i++) { + r[i] = d[i]; + r[i + 8u] = e[i]; + } + return r; } @compute @workgroup_size(8, 8) @@ -64,7 +64,6 @@ fn dec0_main(@builtin(global_invocation_id) id: vec3u) { } } } - // 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)); } |