1#import bevy_pbr::meshlet_bindings::{
2    InstancedOffset,
3    get_aabb,
4    get_aabb_error,
5    get_aabb_child_offset,
6    constants,
7    meshlet_bvh_nodes,
8    meshlet_bvh_cull_count_read,
9    meshlet_bvh_cull_count_write,
10    meshlet_bvh_cull_dispatch,
11    meshlet_bvh_cull_queue,
12    meshlet_meshlet_cull_count_early,
13    meshlet_meshlet_cull_count_late,
14    meshlet_meshlet_cull_dispatch_early,
15    meshlet_meshlet_cull_dispatch_late,
16    meshlet_meshlet_cull_queue,
17    meshlet_second_pass_bvh_count,
18    meshlet_second_pass_bvh_dispatch,
19    meshlet_second_pass_bvh_queue,
20}
21#import bevy_pbr::meshlet_cull_shared::{
22    lod_error_is_imperceptible,
23    aabb_in_frustum,
24    should_occlusion_cull_aabb,
25}
26
27@compute
28@workgroup_size(128, 1, 1) // 8 threads per node, 16 nodes per workgroup
29fn cull_bvh(@builtin(global_invocation_id) global_invocation_id: vec3<u32>) {
30    // Calculate the queue ID for this thread
31    let dispatch_id = global_invocation_id.x;
32    var node = dispatch_id >> 3u;
33    let subnode = dispatch_id & 7u;
34    if node >= meshlet_bvh_cull_count_read { return; }
35
36    node = select(node, constants.rightmost_slot - node, constants.read_from_front == 0u);
37    let instanced_offset = meshlet_bvh_cull_queue[node];
38    let instance_id = instanced_offset.instance_id;
39    let bvh_node = &meshlet_bvh_nodes[instanced_offset.offset];
40
41    var aabb_error_offset = (*bvh_node).aabbs[subnode];
42    let aabb = get_aabb(&aabb_error_offset);
43    let parent_error = get_aabb_error(&aabb_error_offset);
44    let lod_sphere = (*bvh_node).lod_bounds[subnode];
45
46    let parent_is_imperceptible = lod_error_is_imperceptible(lod_sphere, parent_error, instance_id);
47    // Error and frustum cull, in both passes
48    if parent_is_imperceptible || !aabb_in_frustum(aabb, instance_id) { return; }
49
50    let child_offset = get_aabb_child_offset(&aabb_error_offset);
51    let index = subnode >> 2u;
52    let bit_offset = subnode & 3u;
53    let packed_child_count = (*bvh_node).child_counts[index];
54    let child_count = extractBits(packed_child_count, bit_offset * 8u, 8u);
55    var value = InstancedOffset(instance_id, child_offset);
56
57    // If we pass, try occlusion culling
58    // If this node was occluded, push it's children to the second pass to check against this frame's HZB
59    if should_occlusion_cull_aabb(aabb, instance_id) {
60#ifdef MESHLET_FIRST_CULLING_PASS
61        if child_count == 255u {
62            let id = atomicAdd(&meshlet_second_pass_bvh_count, 1u);
63            meshlet_second_pass_bvh_queue[id] = value;
64            if ((id & 15u) == 0u) {
65                atomicAdd(&meshlet_second_pass_bvh_dispatch.x, 1u);
66            }
67        } else {
68            let base = atomicAdd(&meshlet_meshlet_cull_count_late, child_count);
69            let start = constants.rightmost_slot - base;
70            for (var i = start; i < start - child_count; i--) {
71                meshlet_meshlet_cull_queue[i] = value;
72                value.offset += 1u;
73            }
74            let req = (base + child_count + 127u) >> 7u;
75            atomicMax(&meshlet_meshlet_cull_dispatch_late.x, req);
76        }
77#endif
78        return;
79    }
80
81    // If we pass, push the children to the next BVH cull
82    if child_count == 255u {
83        let id = atomicAdd(&meshlet_bvh_cull_count_write, 1u);
84        let index = select(constants.rightmost_slot - id, id, constants.read_from_front == 0u);
85        meshlet_bvh_cull_queue[index] = value;
86        if ((id & 15u) == 0u) {
87            atomicAdd(&meshlet_bvh_cull_dispatch.x, 1u);
88        }
89    } else {
90#ifdef MESHLET_FIRST_CULLING_PASS
91        let base = atomicAdd(&meshlet_meshlet_cull_count_early, child_count);
92        let end = base + child_count;
93        for (var i = base; i < end; i++) {
94            meshlet_meshlet_cull_queue[i] = value;
95            value.offset += 1u;
96        }
97        let req = (end + 127u) >> 7u;
98        atomicMax(&meshlet_meshlet_cull_dispatch_early.x, req);
99#else
100        let base = atomicAdd(&meshlet_meshlet_cull_count_late, child_count);
101        let start = constants.rightmost_slot - base;
102        for (var i = start; i < start - child_count; i--) {
103            meshlet_meshlet_cull_queue[i] = value;
104            value.offset += 1u;
105        }
106        let req = (base + child_count + 127u) >> 7u;
107        atomicMax(&meshlet_meshlet_cull_dispatch_late.x, req);
108#endif
109    }
110}