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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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);
}
|