1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
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]))
));
}
|
