diff options
Diffstat (limited to 'cnn_v3/shaders/cnn_v3_dec1.wgsl')
| -rw-r--r-- | cnn_v3/shaders/cnn_v3_dec1.wgsl | 72 |
1 files changed, 72 insertions, 0 deletions
diff --git a/cnn_v3/shaders/cnn_v3_dec1.wgsl b/cnn_v3/shaders/cnn_v3_dec1.wgsl new file mode 100644 index 0000000..28ae3dc --- /dev/null +++ b/cnn_v3/shaders/cnn_v3_dec1.wgsl @@ -0,0 +1,72 @@ +// CNN v3 — Decoder level 1 +// NearestUp2x(bottleneck) + cat(enc1_skip) -> Conv(16->4, 3x3, zero-pad) + 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) +// +// 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] + +#include "cnn_v3/common" + +const DEC1_IN: u32 = 16u; +const DEC1_OUT: u32 = 4u; + +struct Params { + weight_offset: u32, + _pad: vec3u, + gamma: vec4f, + beta: 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>; + +// 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> { + 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.); + } + 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] + ); +} + +@compute @workgroup_size(8, 8) +fn dec1_main(@builtin(global_invocation_id) id: vec3u) { + let half_dims = vec2i(textureDimensions(enc1_tex)); + 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, DEC1_OUT>; + + for (var o: u32 = 0u; o < DEC1_OUT; o++) { + var sum = get_w(wo, DEC1_OUT * DEC1_IN * 9u + o); // bias + for (var ky: i32 = -1; ky <= 1; ky++) { + for (var kx: i32 = -1; kx <= 1; kx++) { + let feat = load_dec1_concat(coord + vec2i(kx, ky), half_dims); + let ki = u32(ky + 1) * 3u + u32(kx + 1); + for (var i: u32 = 0u; i < DEC1_IN; i++) { + sum += get_w(wo, o * DEC1_IN * 9u + i * 9u + ki) * feat[i]; + } + } + } + out[o] = max(0.0, params.gamma[o] * sum + params.beta[o]); + } + + textureStore(dec1_out, coord, vec4f(out[0], out[1], out[2], out[3])); +} |
