1#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
2@group(0) @binding(0) var<storage, read> mip_0: array<u64>; // Per pixel
3#else
4@group(0) @binding(0) var<storage, read> mip_0: array<u32>; // Per pixel
5#endif
6@group(0) @binding(1) var mip_1: texture_storage_2d<r32float, write>;
7@group(0) @binding(2) var mip_2: texture_storage_2d<r32float, write>;
8@group(0) @binding(3) var mip_3: texture_storage_2d<r32float, write>;
9@group(0) @binding(4) var mip_4: texture_storage_2d<r32float, write>;
10@group(0) @binding(5) var mip_5: texture_storage_2d<r32float, write>;
11@group(0) @binding(6) var mip_6: texture_storage_2d<r32float, read_write>;
12@group(0) @binding(7) var mip_7: texture_storage_2d<r32float, write>;
13@group(0) @binding(8) var mip_8: texture_storage_2d<r32float, write>;
14@group(0) @binding(9) var mip_9: texture_storage_2d<r32float, write>;
15@group(0) @binding(10) var mip_10: texture_storage_2d<r32float, write>;
16@group(0) @binding(11) var mip_11: texture_storage_2d<r32float, write>;
17@group(0) @binding(12) var mip_12: texture_storage_2d<r32float, write>;
18@group(0) @binding(13) var samplr: sampler;
19struct Constants { max_mip_level: u32, view_width: u32 }
20var<push_constant> constants: Constants;
21
22/// Generates a hierarchical depth buffer.
23/// Based on FidelityFX SPD v2.1 https://github.com/GPUOpen-LibrariesAndSDKs/FidelityFX-SDK/blob/d7531ae47d8b36a5d4025663e731a47a38be882f/sdk/include/FidelityFX/gpu/spd/ffx_spd.h#L528
24
25// TODO:
26// * Subgroup support
27// * True single pass downsampling
28
29var<workgroup> intermediate_memory: array<array<f32, 16>, 16>;
30
31@compute
32@workgroup_size(256, 1, 1)
33fn downsample_depth_first(
34    @builtin(num_workgroups) num_workgroups: vec3u,
35    @builtin(workgroup_id) workgroup_id: vec3u,
36    @builtin(local_invocation_index) local_invocation_index: u32,
37) {
38    let sub_xy = remap_for_wave_reduction(local_invocation_index % 64u);
39    let x = sub_xy.x + 8u * ((local_invocation_index >> 6u) % 2u);
40    let y = sub_xy.y + 8u * (local_invocation_index >> 7u);
41
42    downsample_mips_0_and_1(x, y, workgroup_id.xy, local_invocation_index);
43
44    downsample_mips_2_to_5(x, y, workgroup_id.xy, local_invocation_index);
45}
46
47@compute
48@workgroup_size(256, 1, 1)
49fn downsample_depth_second(@builtin(local_invocation_index) local_invocation_index: u32) {
50    let sub_xy = remap_for_wave_reduction(local_invocation_index % 64u);
51    let x = sub_xy.x + 8u * ((local_invocation_index >> 6u) % 2u);
52    let y = sub_xy.y + 8u * (local_invocation_index >> 7u);
53
54    downsample_mips_6_and_7(x, y);
55
56    downsample_mips_8_to_11(x, y, local_invocation_index);
57}
58
59fn downsample_mips_0_and_1(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32) {
60    var v: vec4f;
61
62    var tex = vec2(workgroup_id * 64u) + vec2(x * 2u, y * 2u);
63    var pix = vec2(workgroup_id * 32u) + vec2(x, y);
64    v[0] = reduce_load_mip_0(tex);
65    textureStore(mip_1, pix, vec4(v[0]));
66
67    tex = vec2(workgroup_id * 64u) + vec2(x * 2u + 32u, y * 2u);
68    pix = vec2(workgroup_id * 32u) + vec2(x + 16u, y);
69    v[1] = reduce_load_mip_0(tex);
70    textureStore(mip_1, pix, vec4(v[1]));
71
72    tex = vec2(workgroup_id * 64u) + vec2(x * 2u, y * 2u + 32u);
73    pix = vec2(workgroup_id * 32u) + vec2(x, y + 16u);
74    v[2] = reduce_load_mip_0(tex);
75    textureStore(mip_1, pix, vec4(v[2]));
76
77    tex = vec2(workgroup_id * 64u) + vec2(x * 2u + 32u, y * 2u + 32u);
78    pix = vec2(workgroup_id * 32u) + vec2(x + 16u, y + 16u);
79    v[3] = reduce_load_mip_0(tex);
80    textureStore(mip_1, pix, vec4(v[3]));
81
82    if constants.max_mip_level <= 1u { return; }
83
84    for (var i = 0u; i < 4u; i++) {
85        intermediate_memory[x][y] = v[i];
86        workgroupBarrier();
87        if local_invocation_index < 64u {
88            v[i] = reduce_4(vec4(
89                intermediate_memory[x * 2u + 0u][y * 2u + 0u],
90                intermediate_memory[x * 2u + 1u][y * 2u + 0u],
91                intermediate_memory[x * 2u + 0u][y * 2u + 1u],
92                intermediate_memory[x * 2u + 1u][y * 2u + 1u],
93            ));
94            pix = (workgroup_id * 16u) + vec2(
95                x + (i % 2u) * 8u,
96                y + (i / 2u) * 8u,
97            );
98            textureStore(mip_2, pix, vec4(v[i]));
99        }
100        workgroupBarrier();
101    }
102
103    if local_invocation_index < 64u {
104        intermediate_memory[x + 0u][y + 0u] = v[0];
105        intermediate_memory[x + 8u][y + 0u] = v[1];
106        intermediate_memory[x + 0u][y + 8u] = v[2];
107        intermediate_memory[x + 8u][y + 8u] = v[3];
108    }
109}
110
111fn downsample_mips_2_to_5(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32) {
112    if constants.max_mip_level <= 2u { return; }
113    workgroupBarrier();
114    downsample_mip_2(x, y, workgroup_id, local_invocation_index);
115
116    if constants.max_mip_level <= 3u { return; }
117    workgroupBarrier();
118    downsample_mip_3(x, y, workgroup_id, local_invocation_index);
119
120    if constants.max_mip_level <= 4u { return; }
121    workgroupBarrier();
122    downsample_mip_4(x, y, workgroup_id, local_invocation_index);
123
124    if constants.max_mip_level <= 5u { return; }
125    workgroupBarrier();
126    downsample_mip_5(workgroup_id, local_invocation_index);
127}
128
129fn downsample_mip_2(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32) {
130    if local_invocation_index < 64u {
131        let v = reduce_4(vec4(
132            intermediate_memory[x * 2u + 0u][y * 2u + 0u],
133            intermediate_memory[x * 2u + 1u][y * 2u + 0u],
134            intermediate_memory[x * 2u + 0u][y * 2u + 1u],
135            intermediate_memory[x * 2u + 1u][y * 2u + 1u],
136        ));
137        textureStore(mip_3, (workgroup_id * 8u) + vec2(x, y), vec4(v));
138        intermediate_memory[x * 2u + y % 2u][y * 2u] = v;
139    }
140}
141
142fn downsample_mip_3(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32) {
143    if local_invocation_index < 16u {
144        let v = reduce_4(vec4(
145            intermediate_memory[x * 4u + 0u + 0u][y * 4u + 0u],
146            intermediate_memory[x * 4u + 2u + 0u][y * 4u + 0u],
147            intermediate_memory[x * 4u + 0u + 1u][y * 4u + 2u],
148            intermediate_memory[x * 4u + 2u + 1u][y * 4u + 2u],
149        ));
150        textureStore(mip_4, (workgroup_id * 4u) + vec2(x, y), vec4(v));
151        intermediate_memory[x * 4u + y][y * 4u] = v;
152    }
153}
154
155fn downsample_mip_4(x: u32, y: u32, workgroup_id: vec2u, local_invocation_index: u32) {
156    if local_invocation_index < 4u {
157        let v = reduce_4(vec4(
158            intermediate_memory[x * 8u + 0u + 0u + y * 2u][y * 8u + 0u],
159            intermediate_memory[x * 8u + 4u + 0u + y * 2u][y * 8u + 0u],
160            intermediate_memory[x * 8u + 0u + 1u + y * 2u][y * 8u + 4u],
161            intermediate_memory[x * 8u + 4u + 1u + y * 2u][y * 8u + 4u],
162        ));
163        textureStore(mip_5, (workgroup_id * 2u) + vec2(x, y), vec4(v));
164        intermediate_memory[x + y * 2u][0u] = v;
165    }
166}
167
168fn downsample_mip_5(workgroup_id: vec2u, local_invocation_index: u32) {
169    if local_invocation_index < 1u {
170        let v = reduce_4(vec4(
171            intermediate_memory[0u][0u],
172            intermediate_memory[1u][0u],
173            intermediate_memory[2u][0u],
174            intermediate_memory[3u][0u],
175        ));
176        textureStore(mip_6, workgroup_id, vec4(v));
177    }
178}
179
180fn downsample_mips_6_and_7(x: u32, y: u32) {
181    var v: vec4f;
182
183    var tex = vec2(x * 4u + 0u, y * 4u + 0u);
184    var pix = vec2(x * 2u + 0u, y * 2u + 0u);
185    v[0] = reduce_load_mip_6(tex);
186    textureStore(mip_7, pix, vec4(v[0]));
187
188    tex = vec2(x * 4u + 2u, y * 4u + 0u);
189    pix = vec2(x * 2u + 1u, y * 2u + 0u);
190    v[1] = reduce_load_mip_6(tex);
191    textureStore(mip_7, pix, vec4(v[1]));
192
193    tex = vec2(x * 4u + 0u, y * 4u + 2u);
194    pix = vec2(x * 2u + 0u, y * 2u + 1u);
195    v[2] = reduce_load_mip_6(tex);
196    textureStore(mip_7, pix, vec4(v[2]));
197
198    tex = vec2(x * 4u + 2u, y * 4u + 2u);
199    pix = vec2(x * 2u + 1u, y * 2u + 1u);
200    v[3] = reduce_load_mip_6(tex);
201    textureStore(mip_7, pix, vec4(v[3]));
202
203    if constants.max_mip_level <= 7u { return; }
204
205    let vr = reduce_4(v);
206    textureStore(mip_8, vec2(x, y), vec4(vr));
207    intermediate_memory[x][y] = vr;
208}
209
210fn downsample_mips_8_to_11(x: u32, y: u32, local_invocation_index: u32) {
211    if constants.max_mip_level <= 8u { return; }
212    workgroupBarrier();
213    downsample_mip_8(x, y, local_invocation_index);
214
215    if constants.max_mip_level <= 9u { return; }
216    workgroupBarrier();
217    downsample_mip_9(x, y, local_invocation_index);
218
219    if constants.max_mip_level <= 10u { return; }
220    workgroupBarrier();
221    downsample_mip_10(x, y, local_invocation_index);
222
223    if constants.max_mip_level <= 11u { return; }
224    workgroupBarrier();
225    downsample_mip_11(local_invocation_index);
226}
227
228fn downsample_mip_8(x: u32, y: u32, local_invocation_index: u32) {
229    if local_invocation_index < 64u {
230        let v = reduce_4(vec4(
231            intermediate_memory[x * 2u + 0u][y * 2u + 0u],
232            intermediate_memory[x * 2u + 1u][y * 2u + 0u],
233            intermediate_memory[x * 2u + 0u][y * 2u + 1u],
234            intermediate_memory[x * 2u + 1u][y * 2u + 1u],
235        ));
236        textureStore(mip_9, vec2(x, y), vec4(v));
237        intermediate_memory[x * 2u + y % 2u][y * 2u] = v;
238    }
239}
240
241fn downsample_mip_9(x: u32, y: u32, local_invocation_index: u32) {
242    if local_invocation_index < 16u {
243        let v = reduce_4(vec4(
244            intermediate_memory[x * 4u + 0u + 0u][y * 4u + 0u],
245            intermediate_memory[x * 4u + 2u + 0u][y * 4u + 0u],
246            intermediate_memory[x * 4u + 0u + 1u][y * 4u + 2u],
247            intermediate_memory[x * 4u + 2u + 1u][y * 4u + 2u],
248        ));
249        textureStore(mip_10, vec2(x, y), vec4(v));
250        intermediate_memory[x * 4u + y][y * 4u] = v;
251    }
252}
253
254fn downsample_mip_10(x: u32, y: u32, local_invocation_index: u32) {
255    if local_invocation_index < 4u {
256        let v = reduce_4(vec4(
257            intermediate_memory[x * 8u + 0u + 0u + y * 2u][y * 8u + 0u],
258            intermediate_memory[x * 8u + 4u + 0u + y * 2u][y * 8u + 0u],
259            intermediate_memory[x * 8u + 0u + 1u + y * 2u][y * 8u + 4u],
260            intermediate_memory[x * 8u + 4u + 1u + y * 2u][y * 8u + 4u],
261        ));
262        textureStore(mip_11, vec2(x, y), vec4(v));
263        intermediate_memory[x + y * 2u][0u] = v;
264    }
265}
266
267fn downsample_mip_11(local_invocation_index: u32) {
268    if local_invocation_index < 1u {
269        let v = reduce_4(vec4(
270            intermediate_memory[0u][0u],
271            intermediate_memory[1u][0u],
272            intermediate_memory[2u][0u],
273            intermediate_memory[3u][0u],
274        ));
275        textureStore(mip_12, vec2(0u, 0u), vec4(v));
276    }
277}
278
279fn remap_for_wave_reduction(a: u32) -> vec2u {
280    return vec2(
281        insertBits(extractBits(a, 2u, 3u), a, 0u, 1u),
282        insertBits(extractBits(a, 3u, 3u), extractBits(a, 1u, 2u), 0u, 2u),
283    );
284}
285
286fn reduce_load_mip_0(tex: vec2u) -> f32 {
287    let a = load_mip_0(tex.x, tex.y);
288    let b = load_mip_0(tex.x + 1u, tex.y);
289    let c = load_mip_0(tex.x, tex.y + 1u);
290    let d = load_mip_0(tex.x + 1u, tex.y + 1u);
291    return reduce_4(vec4(a, b, c, d));
292}
293
294fn reduce_load_mip_6(tex: vec2u) -> f32 {
295    return reduce_4(vec4(
296        textureLoad(mip_6, tex + vec2(0u, 0u)).r,
297        textureLoad(mip_6, tex + vec2(0u, 1u)).r,
298        textureLoad(mip_6, tex + vec2(1u, 0u)).r,
299        textureLoad(mip_6, tex + vec2(1u, 1u)).r,
300    ));
301}
302
303fn load_mip_0(x: u32, y: u32) -> f32 {
304    let i = y * constants.view_width + x;
305#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
306    return bitcast<f32>(u32(mip_0[i] >> 32u));
307#else
308    return bitcast<f32>(mip_0[i]);
309#endif
310}
311
312fn reduce_4(v: vec4f) -> f32 {
313    return min(min(v.x, v.y), min(v.z, v.w));
314}