diff options
Diffstat (limited to 'cnn_v3/shaders/gbuf_pack.wgsl')
| -rw-r--r-- | cnn_v3/shaders/gbuf_pack.wgsl | 123 |
1 files changed, 123 insertions, 0 deletions
diff --git a/cnn_v3/shaders/gbuf_pack.wgsl b/cnn_v3/shaders/gbuf_pack.wgsl new file mode 100644 index 0000000..71d8471 --- /dev/null +++ b/cnn_v3/shaders/gbuf_pack.wgsl @@ -0,0 +1,123 @@ +// G-buffer pack compute shader for CNN v3 +// Pass 4: Pack all G-buffer channels into two rgba32uint feature textures (32 bytes/pixel) +// Output feat_tex0 holds 8×f16 geometric channels; feat_tex1 holds 12×u8 context channels. + +struct GBufRes { + resolution: vec2f, +} + +@group(0) @binding(0) var<uniform> gbuf_res: GBufRes; +@group(0) @binding(1) var gbuf_albedo: texture_2d<f32>; +@group(0) @binding(2) var gbuf_normal_mat: texture_2d<f32>; +@group(0) @binding(3) var gbuf_depth: texture_depth_2d; +@group(0) @binding(4) var gbuf_shadow: texture_2d<f32>; +@group(0) @binding(5) var gbuf_transp: texture_2d<f32>; +@group(0) @binding(6) var prev_cnn: texture_2d<f32>; +@group(0) @binding(7) var feat_tex0: texture_storage_2d<rgba32uint, write>; +@group(0) @binding(8) var feat_tex1: texture_storage_2d<rgba32uint, write>; +@group(0) @binding(9) var bilinear_sampler: sampler; + +// Sample depth texture at integer coordinate, clamp to borders. +fn load_depth(coord: vec2i) -> f32 { + let dims = vec2i(textureDimensions(gbuf_depth)); + let c = clamp(coord, vec2i(0), dims - vec2i(1)); + return textureLoad(gbuf_depth, c, 0); +} + +// Box-filter albedo: average of 2×2 texels starting at top-left corner `tl`. +fn box2(tl: vec2i) -> vec3f { + let a = textureLoad(gbuf_albedo, tl + vec2i(0, 0), 0).rgb; + let b = textureLoad(gbuf_albedo, tl + vec2i(1, 0), 0).rgb; + let c = textureLoad(gbuf_albedo, tl + vec2i(0, 1), 0).rgb; + let d = textureLoad(gbuf_albedo, tl + vec2i(1, 1), 0).rgb; + return (a + b + c + d) * 0.25; +} + +// Box-filter albedo: average of 4×4 texels starting at top-left corner `tl`. +fn box4(tl: vec2i) -> vec3f { + var acc = vec3f(0.0); + for (var dy: i32 = 0; dy < 4; dy++) { + for (var dx: i32 = 0; dx < 4; dx++) { + acc += textureLoad(gbuf_albedo, tl + vec2i(dx, dy), 0).rgb; + } + } + return acc * (1.0 / 16.0); +} + +// Decode oct-normal from [0,1] storage → [-1,1] encoded xy → reconstruct z. +fn decode_oct_normal(rg: vec2f) -> vec3f { + let f = rg * 2.0 - vec2f(1.0); + var n = vec3f(f.x, f.y, 1.0 - abs(f.x) - abs(f.y)); + let t = max(-n.z, 0.0); + n.x += select(t, -t, n.x >= 0.0); + n.y += select(t, -t, n.y >= 0.0); + return normalize(n); +} + +@compute @workgroup_size(8, 8) +fn pack_features(@builtin(global_invocation_id) id: vec3u) { + let coord = vec2i(id.xy); + let dims = vec2i(textureDimensions(gbuf_albedo)); + if (coord.x >= dims.x || coord.y >= dims.y) { return; } + + let uv = (vec2f(coord) + vec2f(0.5)) / gbuf_res.resolution; + + // --- Geometric channels (high precision, f16 packed) --- + let albedo = textureLoad(gbuf_albedo, coord, 0).rgb; + let nm = textureLoad(gbuf_normal_mat, coord, 0); + let depth_raw = load_depth(coord); + + // Finite-difference depth gradient (central difference, clamped coords) + let dzdx = (load_depth(coord + vec2i(1, 0)) - load_depth(coord - vec2i(1, 0))) * 0.5; + let dzdy = (load_depth(coord + vec2i(0, 1)) - load_depth(coord - vec2i(0, 1))) * 0.5; + + // Normal: stored as oct-encoded [0,1] in RG; extract just the encoded xy for feat_tex0 + let normal_enc = nm.rg; // already in [0,1] — decode to get the xy for CNN input + let n3 = decode_oct_normal(normal_enc); + // Store oct-encoded in [-1,1] remapped back to what CNN expects (the [-1,1] oct xy) + let oct_xy = normal_enc * 2.0 - vec2f(1.0); // remap [0,1] → [-1,1] + + // Texture 0: 4 u32, each = pack2x16float of two f16 values + // [0] albedo.r | albedo.g + // [1] albedo.b | normal.x (oct, [-1,1]) + // [2] normal.y | depth + // [3] dzdx | dzdy + let t0 = vec4u( + pack2x16float(albedo.rg), + pack2x16float(vec2f(albedo.b, oct_xy.x)), + pack2x16float(vec2f(oct_xy.y, depth_raw)), + pack2x16float(vec2f(dzdx, dzdy)) + ); + textureStore(feat_tex0, coord, t0); + + // --- Context channels (low precision, u8 packed) --- + let mat_id_u8 = nm.b; // mat_id already in [0,1] + let shadow = textureLoad(gbuf_shadow, coord, 0).r; + let transp = textureLoad(gbuf_transp, coord, 0).r; + let prev = textureSampleLevel(prev_cnn, bilinear_sampler, uv, 0.0).rgb; + + // MIP 1: 2×2 box filter (half resolution context) + // Use top-left aligned 2×2 block at half-res position + let tl1 = coord * 2; // this pixel's 2×2 region in full-res (mip1 is at half-res) + // Actually we want to sample the neighborhood around this pixel for downsampled context. + // mip1: sample a 2×2 box centered on the pixel in full-res coordinates + let tl1c = max(coord - vec2i(0), vec2i(0)); + let mip1 = box2(tl1c); + + // mip2: sample a 4×4 box + let tl2c = max(coord - vec2i(1), vec2i(0)); + let mip2 = box4(tl2c); + + // Texture 1: 4 u32, each = pack4x8unorm of four u8 values + // [0] mat_id | prev.r | prev.g | prev.b + // [1] mip1.r | mip1.g | mip1.b | mip2.r + // [2] mip2.g | mip2.b | shadow | transp + // [3] spare (0) + let t1 = vec4u( + pack4x8unorm(vec4f(mat_id_u8, prev.r, prev.g, prev.b)), + pack4x8unorm(vec4f(mip1.r, mip1.g, mip1.b, mip2.r)), + pack4x8unorm(vec4f(mip2.g, mip2.b, shadow, transp)), + 0u + ); + textureStore(feat_tex1, coord, t1); +} |