1 files changed, 88 insertions, 0 deletions
diff --git a/cnn_v3/shaders/cnn_v3_enc1.wgsl b/cnn_v3/shaders/cnn_v3_enc1.wgsl
new file mode 100644
index 0000000..23e485d
--- /dev/null
+++ b/cnn_v3/shaders/cnn_v3_enc1.wgsl
@@ -0,0 +1,88 @@
+// CNN v3 — Encoder level 1
+// AvgPool2x2(enc0) + Conv(4->8, 3x3, zero-pad) + FiLM + ReLU
+//
+// Input:  enc0_tex (rgba16float, 4ch)   full-res
+// Output: enc1_out (rgba32uint, 8xf16)  half-res  (dispatch at half-res dims)
+//
+// Weight layout (f16, OIHW + bias):
+//   [0 .. 4*8*9)  conv: w[out][in][ky][kx]
+//   [288 .. +8)   bias: b[out]
+
+#include "cnn_v3/common"
+
+const ENC1_IN:  u32 = 4u;
+const ENC1_OUT: u32 = 8u;
+
+struct Params {
+    weight_offset: u32,
+    _pad: vec3u,
+    gamma_lo: vec4f,   // FiLM gamma ch 0-3
+    gamma_hi: vec4f,   // FiLM gamma ch 4-7
+    beta_lo:  vec4f,   // FiLM beta  ch 0-3
+    beta_hi:  vec4f,   // FiLM beta  ch 4-7
+}
+
+@group(0) @binding(0) var enc0_tex: texture_2d<f32>;
+@group(0) @binding(1) var<storage, read> weights: array<u32>;
+@group(0) @binding(2) var<uniform> params: Params;
+@group(0) @binding(3) var enc1_out: texture_storage_2d<rgba32uint, write>;
+
+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];
+}
+
+// Avg-pool 2x2 from enc0_tex at half-res coord hcoord.
+// Returns zeros for OOB half-res coords (zero-padding for the conv).
+fn load_enc0_avg(hcoord: vec2i, full_dims: vec2i) -> array<f32, 4> {
+    let half_dims = full_dims / 2;
+    if (hcoord.x < 0 || hcoord.y < 0 || hcoord.x >= half_dims.x || hcoord.y >= half_dims.y) {
+        return array<f32, 4>(0., 0., 0., 0.);
+    }
+    let base = hcoord * 2;
+    var s = vec4f(0.);
+    for (var dy: i32 = 0; dy < 2; dy++) {
+        for (var dx: i32 = 0; dx < 2; dx++) {
+            let fc = clamp(base + vec2i(dx, dy), vec2i(0), full_dims - vec2i(1));
+            s += textureLoad(enc0_tex, fc, 0);
+        }
+    }
+    let avg = s * 0.25;
+    return array<f32, 4>(avg.x, avg.y, avg.z, avg.w);
+}
+
+@compute @workgroup_size(8, 8)
+fn enc1_main(@builtin(global_invocation_id) id: vec3u) {
+    let full_dims = vec2i(textureDimensions(enc0_tex));
+    let half_dims = full_dims / 2;
+    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, ENC1_OUT>;
+
+    for (var o: u32 = 0u; o < ENC1_OUT; o++) {
+        var sum = get_w(wo, ENC1_OUT * ENC1_IN * 9u + o);  // bias
+        for (var ky: i32 = -1; ky <= 1; ky++) {
+            for (var kx: i32 = -1; kx <= 1; kx++) {
+                let feat = load_enc0_avg(coord + vec2i(kx, ky), full_dims);
+                let ki   = u32(ky + 1) * 3u + u32(kx + 1);
+                for (var i: u32 = 0u; i < ENC1_IN; i++) {
+                    sum += get_w(wo, o * ENC1_IN * 9u + i * 9u + ki) * feat[i];
+                }
+            }
+        }
+        out[o] = max(0.0, film_gamma(o) * sum + film_beta(o));
+    }
+
+    textureStore(enc1_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]))
+    ));
+}