summaryrefslogtreecommitdiff
path: root/cnn_v3/shaders
diff options
context:
space:
mode:
Diffstat (limited to 'cnn_v3/shaders')
-rw-r--r--cnn_v3/shaders/cnn_v3_bottleneck.wgsl54
-rw-r--r--cnn_v3/shaders/cnn_v3_dec0.wgsl43
-rw-r--r--cnn_v3/shaders/cnn_v3_dec1.wgsl85
-rw-r--r--cnn_v3/shaders/cnn_v3_enc0.wgsl37
-rw-r--r--cnn_v3/shaders/cnn_v3_enc1.wgsl75
5 files changed, 179 insertions, 115 deletions
diff --git a/cnn_v3/shaders/cnn_v3_bottleneck.wgsl b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl
index e30682b..09819cc 100644
--- a/cnn_v3/shaders/cnn_v3_bottleneck.wgsl
+++ b/cnn_v3/shaders/cnn_v3_bottleneck.wgsl
@@ -1,43 +1,49 @@
// CNN v3 — Bottleneck
-// AvgPool2x2(enc1) + Conv(8->8, 3x3, dilation=2) + ReLU (no FiLM)
+// AvgPool2x2(enc1) + Conv(16->16, 3x3, dilation=2) + ReLU (no FiLM)
//
-// Input: enc1_tex (rgba32uint, 8xf16) half-res
-// Output: bottleneck_out (rgba32uint, 8xf16) quarter-res (dispatch at quarter-res dims)
+// Input: enc1_tex_lo (rgba32uint, 8xf16) half-res ch 0-7
+// enc1_tex_hi (rgba32uint, 8xf16) half-res ch 8-15
+// Output: bn_out_lo (rgba32uint, 8xf16) quarter-res
+// bn_out_hi (rgba32uint, 8xf16) quarter-res
//
// Weight layout (f16, OIHW + bias):
-// [0 .. 8*8*9) conv: w[out][in][ky*3+kx] (3x3 kernel, OIHW)
-// [576 .. +8) bias: b[out]
+// [0 .. 16*16*9) conv: w[out][in][ky*3+kx]
+// [2304 .. +16) bias: b[out]
#include "cnn_v3/common"
-const BN_IN: u32 = 8u;
-const BN_OUT: u32 = 8u;
+const BN_IN: u32 = 16u;
+const BN_OUT: u32 = 16u;
const BN_DILATION: i32 = 2;
struct Params {
weight_offset: u32,
- _pad0: u32, _pad1: u32, _pad2: u32, // 3 explicit pads: array<u32,3> invalid in uniform
+ _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>;
+@group(0) @binding(0) var enc1_tex_lo: texture_2d<u32>;
+@group(0) @binding(1) var enc1_tex_hi: 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 bn_out_lo: texture_storage_2d<rgba32uint, write>;
+@group(0) @binding(5) var bn_out_hi: 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 3x3 conv).
-fn load_enc1_avg(qcoord: vec2i, half_dims: vec2i) -> array<f32, 8> {
+fn load_enc1_avg(qcoord: vec2i, half_dims: vec2i) -> array<f32, 16> {
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.);
+ return array<f32, 16>(0.,0.,0.,0.,0.,0.,0.,0.,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]; }
+ let lo = unpack_8ch(enc1_tex_lo, hc);
+ let hi = unpack_8ch(enc1_tex_hi, hc);
+ for (var i: u32 = 0u; i < 8u; i++) {
+ s[i] += lo[i];
+ s[i + 8u] += hi[i];
+ }
}
}
for (var i: u32 = 0u; i < BN_IN; i++) { s[i] *= 0.25; }
@@ -46,7 +52,7 @@ fn load_enc1_avg(qcoord: vec2i, half_dims: vec2i) -> array<f32, 8> {
@compute @workgroup_size(8, 8)
fn bottleneck_main(@builtin(global_invocation_id) id: vec3u) {
- let half_dims = vec2i(textureDimensions(enc1_tex));
+ let half_dims = vec2i(textureDimensions(enc1_tex_lo));
let quart_dims = half_dims / 2;
let coord = vec2i(id.xy);
if (coord.x >= quart_dims.x || coord.y >= quart_dims.y) { return; }
@@ -55,7 +61,7 @@ fn bottleneck_main(@builtin(global_invocation_id) id: vec3u) {
var out: array<f32, BN_OUT>;
for (var o: u32 = 0u; o < BN_OUT; o++) {
- var sum = get_w(wo, BN_OUT * BN_IN * 9u + o); // bias (at end of 3x3 conv weights)
+ var sum = get_w(wo, BN_OUT * BN_IN * 9u + o); // bias
for (var ky: i32 = -1; ky <= 1; ky++) {
for (var kx: i32 = -1; kx <= 1; kx++) {
let feat = load_enc1_avg(coord + vec2i(kx, ky) * BN_DILATION, half_dims);
@@ -68,10 +74,16 @@ fn bottleneck_main(@builtin(global_invocation_id) id: vec3u) {
out[o] = max(0.0, sum);
}
- textureStore(bottleneck_out, coord, vec4u(
+ textureStore(bn_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(bn_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]))
+ ));
}
diff --git a/cnn_v3/shaders/cnn_v3_dec0.wgsl b/cnn_v3/shaders/cnn_v3_dec0.wgsl
index a2a70ac..617b5a2 100644
--- a/cnn_v3/shaders/cnn_v3_dec0.wgsl
+++ b/cnn_v3/shaders/cnn_v3_dec0.wgsl
@@ -1,19 +1,17 @@
// CNN v3 — Decoder level 0 + output
-// NearestUp2x(dec1) + cat(enc0_skip) -> Conv(8->4, 3x3, zero-pad) + FiLM + ReLU + Sigmoid
+// NearestUp2x(dec1) + cat(enc0_skip) -> Conv(16->4, 3x3) + FiLM + ReLU + Sigmoid
//
-// Inputs: dec1_tex (rgba16float, 4ch) half-res
-// enc0_tex (rgba16float, 4ch) full-res (skip connection)
-// Output: output_tex (rgba16float, 4ch) full-res (dispatch at full-res dims)
+// 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 .. 8*4*9) conv: w[out][in][ky][kx] (in=8: 4 dec1 + 4 enc0 skip)
-// [288 .. +4) bias: b[out]
-//
-// Parity note: sigmoid applied after FiLM+ReLU, not after raw conv (matches train_cnn_v3.py).
+// [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 = 8u;
+const DEC0_IN: u32 = 16u;
const DEC0_OUT: u32 = 4u;
struct Params {
@@ -23,25 +21,27 @@ struct Params {
beta: vec4f,
}
-@group(0) @binding(0) var dec1_tex: texture_2d<f32>;
-@group(0) @binding(1) var enc0_tex: texture_2d<f32>;
+@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 8 concatenated channels at full-res coord:
-// ch 0-3: dec1 nearest-up (dec1_tex[coord/2])
-// ch 4-7: enc0 skip (enc0_tex[coord])
-// Returns zeros for OOB coord (zero-padding for the conv).
-fn load_dec0_concat(coord: vec2i, full_dims: vec2i) -> array<f32, 8> {
+// 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 array<f32, 8>(0., 0., 0., 0., 0., 0., 0., 0.);
+ return r;
}
let half_dims = vec2i(textureDimensions(dec1_tex));
- let hc = clamp(coord / 2, vec2i(0), half_dims - vec2i(1));
- let d = textureLoad(dec1_tex, hc, 0);
- let e = textureLoad(enc0_tex, coord, 0);
- return array<f32, 8>(d.x, d.y, d.z, d.w, e.x, e.y, e.z, e.w);
+ 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)
@@ -64,7 +64,6 @@ fn dec0_main(@builtin(global_invocation_id) id: vec3u) {
}
}
}
- // FiLM + ReLU + Sigmoid (matches training forward())
let v = max(0.0, params.gamma[o] * sum + params.beta[o]);
out[o] = 1.0 / (1.0 + exp(-v));
}
diff --git a/cnn_v3/shaders/cnn_v3_dec1.wgsl b/cnn_v3/shaders/cnn_v3_dec1.wgsl
index 28ae3dc..fadea3b 100644
--- a/cnn_v3/shaders/cnn_v3_dec1.wgsl
+++ b/cnn_v3/shaders/cnn_v3_dec1.wgsl
@@ -1,53 +1,71 @@
// CNN v3 — Decoder level 1
-// NearestUp2x(bottleneck) + cat(enc1_skip) -> Conv(16->4, 3x3, zero-pad) + FiLM + ReLU
+// NearestUp2x(bottleneck) + cat(enc1_skip) -> Conv(32->8, 3x3) + FiLM + ReLU
//
-// Inputs: bottleneck_tex (rgba32uint, 8xf16) quarter-res
-// enc1_tex (rgba32uint, 8xf16) half-res (skip connection)
-// Output: dec1_out (rgba16float, 4ch) half-res (dispatch at half-res dims)
+// Inputs: bn_tex_lo (rgba32uint, 8xf16) quarter-res ch 0-7
+// bn_tex_hi (rgba32uint, 8xf16) quarter-res ch 8-15
+// enc1_tex_lo (rgba32uint, 8xf16) half-res skip ch 0-7
+// enc1_tex_hi (rgba32uint, 8xf16) half-res skip ch 8-15
+// Output: dec1_out (rgba32uint, 8xf16) half-res
//
// Weight layout (f16, OIHW + bias):
-// [0 .. 16*4*9) conv: w[out][in][ky][kx] (in=16: 8 bottleneck + 8 enc1 skip)
-// [576 .. +4) bias: b[out]
+// [0 .. 32*8*9) conv: w[out][in][ky][kx] (in=32: 16 bn + 16 enc1 skip)
+// [2304 .. +8) bias: b[out]
#include "cnn_v3/common"
-const DEC1_IN: u32 = 16u;
-const DEC1_OUT: u32 = 4u;
+const DEC1_IN: u32 = 32u;
+const DEC1_OUT: u32 = 8u;
struct Params {
weight_offset: u32,
_pad: vec3u,
- gamma: vec4f,
- beta: vec4f,
+ gamma_lo: vec4f,
+ gamma_hi: vec4f,
+ beta_lo: vec4f,
+ beta_hi: vec4f,
}
-@group(0) @binding(0) var bottleneck_tex: texture_2d<u32>;
-@group(0) @binding(1) var enc1_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 dec1_out: texture_storage_2d<rgba16float, write>;
+@group(0) @binding(0) var bn_tex_lo: texture_2d<u32>;
+@group(0) @binding(1) var bn_tex_hi: texture_2d<u32>;
+@group(0) @binding(2) var enc1_tex_lo: texture_2d<u32>;
+@group(0) @binding(3) var enc1_tex_hi: texture_2d<u32>;
+@group(0) @binding(4) var<storage, read> weights: array<u32>;
+@group(0) @binding(5) var<uniform> params: Params;
+@group(0) @binding(6) var dec1_out: texture_storage_2d<rgba32uint, write>;
-// Load 16 concatenated channels at half-res coord hcoord:
-// ch 0-7: bottleneck nearest-up (bottleneck_tex[hcoord/2])
-// ch 8-15: enc1 skip (enc1_tex[hcoord])
-// Returns zeros for OOB hcoord (zero-padding for the conv).
-fn load_dec1_concat(hcoord: vec2i, half_dims: vec2i) -> array<f32, 16> {
+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];
+}
+
+// Load 32ch: [bn_nearest_up(16ch), enc1_skip(16ch)]
+fn load_dec1_concat(hcoord: vec2i, half_dims: vec2i) -> array<f32, 32> {
+ var r: array<f32, 32>;
if (hcoord.x < 0 || hcoord.y < 0 || hcoord.x >= half_dims.x || hcoord.y >= half_dims.y) {
- return array<f32, 16>(0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.);
+ return r;
}
let quart_dims = half_dims / 2;
- let qc = clamp(hcoord / 2, vec2i(0), quart_dims - vec2i(1));
- let b = unpack_8ch(bottleneck_tex, qc);
- let s = unpack_8ch(enc1_tex, hcoord);
- return array<f32, 16>(
- b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7],
- s[0], s[1], s[2], s[3], s[4], s[5], s[6], s[7]
- );
+ let qc = clamp(hcoord / 2, vec2i(0), quart_dims - vec2i(1));
+ let blo = unpack_8ch(bn_tex_lo, qc);
+ let bhi = unpack_8ch(bn_tex_hi, qc);
+ let slo = unpack_8ch(enc1_tex_lo, hcoord);
+ let shi = unpack_8ch(enc1_tex_hi, hcoord);
+ for (var i: u32 = 0u; i < 8u; i++) {
+ r[i] = blo[i];
+ r[i + 8u] = bhi[i];
+ r[i + 16u] = slo[i];
+ r[i + 24u] = shi[i];
+ }
+ return r;
}
@compute @workgroup_size(8, 8)
fn dec1_main(@builtin(global_invocation_id) id: vec3u) {
- let half_dims = vec2i(textureDimensions(enc1_tex));
+ let half_dims = vec2i(textureDimensions(enc1_tex_lo));
let coord = vec2i(id.xy);
if (coord.x >= half_dims.x || coord.y >= half_dims.y) { return; }
@@ -65,8 +83,13 @@ fn dec1_main(@builtin(global_invocation_id) id: vec3u) {
}
}
}
- out[o] = max(0.0, params.gamma[o] * sum + params.beta[o]);
+ out[o] = max(0.0, film_gamma(o) * sum + film_beta(o));
}
- textureStore(dec1_out, coord, vec4f(out[0], out[1], out[2], out[3]));
+ textureStore(dec1_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]))
+ ));
}
diff --git a/cnn_v3/shaders/cnn_v3_enc0.wgsl b/cnn_v3/shaders/cnn_v3_enc0.wgsl
index e171ca7..84d40fd 100644
--- a/cnn_v3/shaders/cnn_v3_enc0.wgsl
+++ b/cnn_v3/shaders/cnn_v3_enc0.wgsl
@@ -1,32 +1,42 @@
// CNN v3 — Encoder level 0
-// Conv(20->4, 3x3, zero-pad) + FiLM + ReLU
+// Conv(20->8, 3x3, zero-pad) + FiLM + ReLU
//
// Input: feat_tex0 (rgba32uint, 8xf16), feat_tex1 (rgba32uint, 12ch u8norm) full-res
-// Output: enc0_out (rgba16float, 4ch) full-res
+// Output: enc0_out (rgba32uint, 8xf16) full-res
//
// Weight layout (f16, OIHW + bias):
-// [0 .. 20*4*9) conv: w[out][in][ky][kx]
-// [720 .. +4) bias: b[out]
+// [0 .. 20*8*9) conv: w[out][in][ky][kx]
+// [1440 .. +8) bias: b[out]
#include "cnn_v3/common"
const ENC0_IN: u32 = 20u;
-const ENC0_OUT: u32 = 4u;
+const ENC0_OUT: u32 = 8u;
struct Params {
weight_offset: u32,
_pad: vec3u,
- gamma: vec4f,
- beta: vec4f,
+ gamma_lo: vec4f,
+ gamma_hi: vec4f,
+ beta_lo: vec4f,
+ beta_hi: vec4f,
}
@group(0) @binding(0) var feat_tex0: texture_2d<u32>;
@group(0) @binding(1) var feat_tex1: 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 enc0_out: texture_storage_2d<rgba16float, write>;
+@group(0) @binding(4) var enc0_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];
+}
-// Unpack all 20 feature channels at coord. Returns zeros for OOB (zero-padding).
fn load_feat(coord: vec2i, dims: vec2i) -> array<f32, 20> {
if (coord.x < 0 || coord.y < 0 || coord.x >= dims.x || coord.y >= dims.y) {
return array<f32, 20>(0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.,0.);
@@ -68,8 +78,13 @@ fn enc0_main(@builtin(global_invocation_id) id: vec3u) {
}
}
}
- out[o] = max(0.0, params.gamma[o] * sum + params.beta[o]);
+ out[o] = max(0.0, film_gamma(o) * sum + film_beta(o));
}
- textureStore(enc0_out, coord, vec4f(out[0], out[1], out[2], out[3]));
+ textureStore(enc0_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]))
+ ));
}
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]))
));
}
generated by cgit v1.2.3 (git 2.51.0) at 2026-03-26 20:06:06 +0000