diff options
Diffstat (limited to 'cnn_v3/shaders/cnn_v3_enc1.wgsl')
| -rw-r--r-- | cnn_v3/shaders/cnn_v3_enc1.wgsl | 75 |
1 files changed, 45 insertions, 30 deletions
diff --git a/cnn_v3/shaders/cnn_v3_enc1.wgsl b/cnn_v3/shaders/cnn_v3_enc1.wgsl index 23e485d..eb41279 100644 --- a/cnn_v3/shaders/cnn_v3_enc1.wgsl +++ b/cnn_v3/shaders/cnn_v3_enc1.wgsl @@ -1,58 +1,67 @@ // CNN v3 — Encoder level 1 -// AvgPool2x2(enc0) + Conv(4->8, 3x3, zero-pad) + FiLM + ReLU +// AvgPool2x2(enc0) + Conv(8->16, 3x3, zero-pad) + FiLM + ReLU // -// Input: enc0_tex (rgba16float, 4ch) full-res -// Output: enc1_out (rgba32uint, 8xf16) half-res (dispatch at half-res dims) +// Input: enc0_tex (rgba32uint, 8xf16) full-res +// Output: enc1_out_lo (rgba32uint, 8xf16) half-res ch 0-7 +// enc1_out_hi (rgba32uint, 8xf16) half-res ch 8-15 // // Weight layout (f16, OIHW + bias): -// [0 .. 4*8*9) conv: w[out][in][ky][kx] -// [288 .. +8) bias: b[out] +// [0 .. 8*16*9) conv: w[out][in][ky][kx] +// [1152 .. +16) bias: b[out] #include "cnn_v3/common" -const ENC1_IN: u32 = 4u; -const ENC1_OUT: u32 = 8u; +const ENC1_IN: u32 = 8u; +const ENC1_OUT: u32 = 16u; 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 + gamma_0: vec4f, + gamma_1: vec4f, + gamma_2: vec4f, + gamma_3: vec4f, + beta_0: vec4f, + beta_1: vec4f, + beta_2: vec4f, + beta_3: vec4f, } -@group(0) @binding(0) var enc0_tex: texture_2d<f32>; +@group(0) @binding(0) var enc0_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 enc1_out: texture_storage_2d<rgba32uint, write>; +@group(0) @binding(3) var enc1_out_lo: texture_storage_2d<rgba32uint, write>; +@group(0) @binding(4) var enc1_out_hi: 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]; + if (o < 4u) { return params.gamma_0[o]; } + if (o < 8u) { return params.gamma_1[o - 4u]; } + if (o < 12u) { return params.gamma_2[o - 8u]; } + return params.gamma_3[o - 12u]; } fn film_beta(o: u32) -> f32 { - if (o < 4u) { return params.beta_lo[o]; } - return params.beta_hi[o - 4u]; + if (o < 4u) { return params.beta_0[o]; } + if (o < 8u) { return params.beta_1[o - 4u]; } + if (o < 12u) { return params.beta_2[o - 8u]; } + return params.beta_3[o - 12u]; } -// 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> { +fn load_enc0_avg(hcoord: vec2i, full_dims: vec2i) -> array<f32, 8> { 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.); + return array<f32, 8>(0., 0., 0., 0., 0., 0., 0., 0.); } let base = hcoord * 2; - var s = vec4f(0.); + var s: array<f32, 8>; 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 f = unpack_8ch(enc0_tex, fc); + for (var i: u32 = 0u; i < 8u; i++) { s[i] += f[i]; } } } - let avg = s * 0.25; - return array<f32, 4>(avg.x, avg.y, avg.z, avg.w); + for (var i: u32 = 0u; i < 8u; i++) { s[i] *= 0.25; } + return s; } @compute @workgroup_size(8, 8) @@ -79,10 +88,16 @@ fn enc1_main(@builtin(global_invocation_id) id: vec3u) { 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])) + textureStore(enc1_out_lo, 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])) + )); + textureStore(enc1_out_hi, coord, vec4u( + pack2x16float(vec2f(out[8], out[9])), + pack2x16float(vec2f(out[10], out[11])), + pack2x16float(vec2f(out[12], out[13])), + pack2x16float(vec2f(out[14], out[15])) )); } |