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
|
// CNN v3 — Decoder level 0 + output
// NearestUp2x(dec1) + cat(enc0_skip) -> Conv(16->4, 3x3) + FiLM + ReLU + Sigmoid
//
// Inputs: dec1_tex (rgba32uint, 8xf16) half-res
// enc0_tex (rgba32uint, 8xf16) full-res (skip connection)
// Output: output_tex (rgba16float, 4ch) full-res
//
// Weight layout (f16, OIHW + bias):
// [0 .. 16*4*9) conv: w[out][in][ky][kx] (in=16: 8 dec1 + 8 enc0 skip)
// [576 .. +4) bias: b[out]
#include "cnn_v3/common"
const DEC0_IN: u32 = 16u;
const DEC0_OUT: u32 = 4u;
struct Params {
weight_offset: u32,
_pad: vec3u,
gamma: vec4f,
beta: vec4f,
}
@group(0) @binding(0) var dec1_tex: texture_2d<u32>;
@group(0) @binding(1) var enc0_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 output_tex: texture_storage_2d<rgba16float, write>;
// Load 16ch: ch 0-7 from dec1 nearest-up, ch 8-15 from enc0 skip.
fn load_dec0_concat(coord: vec2i, full_dims: vec2i) -> array<f32, 16> {
var r: array<f32, 16>;
if (coord.x < 0 || coord.y < 0 || coord.x >= full_dims.x || coord.y >= full_dims.y) {
return r;
}
let half_dims = vec2i(textureDimensions(dec1_tex));
let hc = clamp(coord / 2, vec2i(0), half_dims - vec2i(1));
let d = unpack_8ch(dec1_tex, hc);
let e = unpack_8ch(enc0_tex, coord);
for (var i: u32 = 0u; i < 8u; i++) {
r[i] = d[i];
r[i + 8u] = e[i];
}
return r;
}
@compute @workgroup_size(8, 8)
fn dec0_main(@builtin(global_invocation_id) id: vec3u) {
let full_dims = vec2i(textureDimensions(enc0_tex));
let coord = vec2i(id.xy);
if (coord.x >= full_dims.x || coord.y >= full_dims.y) { return; }
let wo = params.weight_offset;
var out: array<f32, DEC0_OUT>;
for (var o: u32 = 0u; o < DEC0_OUT; o++) {
var sum = get_w(wo, DEC0_OUT * DEC0_IN * 9u + o); // bias
for (var ky: i32 = -1; ky <= 1; ky++) {
for (var kx: i32 = -1; kx <= 1; kx++) {
let feat = load_dec0_concat(coord + vec2i(kx, ky), full_dims);
let ki = u32(ky + 1) * 3u + u32(kx + 1);
for (var i: u32 = 0u; i < DEC0_IN; i++) {
sum += get_w(wo, o * DEC0_IN * 9u + i * 9u + ki) * feat[i];
}
}
}
let v = max(0.0, params.gamma[o] * sum + params.beta[o]);
out[o] = 1.0 / (1.0 + exp(-v));
}
textureStore(output_tex, coord, vec4f(out[0], out[1], out[2], out[3]));
}
|
