1 files changed, 73 insertions, 0 deletions
diff --git a/cnn_v3/shaders/cnn_v3_bottleneck.wgsl b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl
new file mode 100644
index 0000000..909fd41
--- /dev/null
+++ b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl
@@ -0,0 +1,73 @@
+// CNN v3 — Bottleneck
+// AvgPool2x2(enc1) + Conv(8->8, 1x1) + ReLU  (no FiLM)
+//
+// 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]
+
+#include "cnn_v3/common"
+
+const BN_IN:  u32 = 8u;
+const BN_OUT: u32 = 8u;
+
+struct Params {
+    weight_offset: u32,
+    _pad0: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var enc1_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 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).
+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) {
+        return array<f32, 8>(0., 0., 0., 0., 0., 0., 0., 0.);
+    }
+    let base = qcoord * 2;
+    var s: array<f32, BN_IN>;
+    for (var dy: i32 = 0; dy < 2; dy++) {
+        for (var dx: i32 = 0; dx < 2; dx++) {
+            let hc = clamp(base + vec2i(dx, dy), vec2i(0), half_dims - vec2i(1));
+            let f  = unpack_8ch(enc1_tex, hc);
+            for (var i: u32 = 0u; i < BN_IN; i++) { s[i] += f[i]; }
+        }
+    }
+    for (var i: u32 = 0u; i < BN_IN; i++) { s[i] *= 0.25; }
+    return s;
+}
+
+@compute @workgroup_size(8, 8)
+fn bottleneck_main(@builtin(global_invocation_id) id: vec3u) {
+    let half_dims  = vec2i(textureDimensions(enc1_tex));
+    let quart_dims = half_dims / 2;
+    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);
+    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];
+        }
+        out[o] = max(0.0, sum);
+    }
+
+    textureStore(bottleneck_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]))
+    ));
+}