diff options
Diffstat (limited to 'cnn_v3/shaders/cnn_v3_bottleneck.wgsl')
| -rw-r--r-- | cnn_v3/shaders/cnn_v3_bottleneck.wgsl | 32 |
1 files changed, 19 insertions, 13 deletions
diff --git a/cnn_v3/shaders/cnn_v3_bottleneck.wgsl b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl index e24586f..e30682b 100644 --- a/cnn_v3/shaders/cnn_v3_bottleneck.wgsl +++ b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl @@ -1,17 +1,18 @@ // CNN v3 — Bottleneck -// AvgPool2x2(enc1) + Conv(8->8, 1x1) + ReLU (no FiLM) +// AvgPool2x2(enc1) + Conv(8->8, 3x3, dilation=2) + ReLU (no FiLM) // -// Input: enc1_tex (rgba32uint, 8xf16) half-res -// Output: bottleneck_out (rgba32uint, 8xf16) quarter-res (dispatch at quarter-res dims) +// Input: enc1_tex (rgba32uint, 8xf16) half-res +// Output: bottleneck_out (rgba32uint, 8xf16) quarter-res (dispatch at quarter-res dims) // // Weight layout (f16, OIHW + bias): -// [0 .. 8*8*1) conv: w[out][in] (1x1 kernel) -// [64 .. +8) bias: b[out] +// [0 .. 8*8*9) conv: w[out][in][ky*3+kx] (3x3 kernel, OIHW) +// [576 .. +8) bias: b[out] #include "cnn_v3/common" -const BN_IN: u32 = 8u; -const BN_OUT: u32 = 8u; +const BN_IN: u32 = 8u; +const BN_OUT: u32 = 8u; +const BN_DILATION: i32 = 2; struct Params { weight_offset: u32, @@ -24,7 +25,7 @@ struct Params { @group(0) @binding(3) var bottleneck_out: texture_storage_2d<rgba32uint, write>; // Avg-pool 2x2 from enc1_tex at quarter-res coord qcoord. -// Returns zeros for OOB quarter-res coords (zero-padding for the 1x1 conv). +// Returns zeros for OOB quarter-res coords (zero-padding for the 3x3 conv). fn load_enc1_avg(qcoord: vec2i, half_dims: vec2i) -> array<f32, 8> { let quart_dims = half_dims / 2; if (qcoord.x < 0 || qcoord.y < 0 || qcoord.x >= quart_dims.x || qcoord.y >= quart_dims.y) { @@ -50,14 +51,19 @@ fn bottleneck_main(@builtin(global_invocation_id) id: vec3u) { let coord = vec2i(id.xy); if (coord.x >= quart_dims.x || coord.y >= quart_dims.y) { return; } - let wo = params.weight_offset; - let feat = load_enc1_avg(coord, half_dims); + let wo = params.weight_offset; var out: array<f32, BN_OUT>; for (var o: u32 = 0u; o < BN_OUT; o++) { - var sum = get_w(wo, BN_OUT * BN_IN + o); // bias (1x1 kernel: no spatial idx) - for (var i: u32 = 0u; i < BN_IN; i++) { - sum += get_w(wo, o * BN_IN + i) * feat[i]; + var sum = get_w(wo, BN_OUT * BN_IN * 9u + o); // bias (at end of 3x3 conv weights) + for (var ky: i32 = -1; ky <= 1; ky++) { + for (var kx: i32 = -1; kx <= 1; kx++) { + let feat = load_enc1_avg(coord + vec2i(kx, ky) * BN_DILATION, half_dims); + let ki = u32(ky + 1) * 3u + u32(kx + 1); + for (var i: u32 = 0u; i < BN_IN; i++) { + sum += get_w(wo, o * BN_IN * 9u + i * 9u + ki) * feat[i]; + } + } } out[o] = max(0.0, sum); } |