1#import bevy_pbr::meshlet_bindings::{
2    meshlet_cluster_meshlet_ids,
3    meshlet_bounding_spheres,
4    meshlet_simplification_errors,
5    meshlet_cluster_instance_ids,
6    meshlet_instance_uniforms,
7    meshlet_second_pass_candidates,
8    depth_pyramid,
9    view,
10    previous_view,
11    should_cull_instance,
12    cluster_is_second_pass_candidate,
13    meshlet_software_raster_indirect_args,
14    meshlet_hardware_raster_indirect_args,
15    meshlet_raster_clusters,
16    constants,
17    MeshletBoundingSphere,
18}
19#import bevy_render::maths::affine3_to_square
20
21/// Culls individual clusters (1 per thread) in two passes (two pass occlusion culling), and outputs a bitmask of which clusters survived.
22/// 1. The first pass tests instance visibility, frustum culling, LOD selection, and finally occlusion culling using last frame's depth pyramid.
23/// 2. The second pass performs occlusion culling (using the depth buffer generated from the first pass) on all clusters that passed
24///    the instance, frustum, and LOD tests in the first pass, but were not visible last frame according to the occlusion culling.
25
26@compute
27@workgroup_size(128, 1, 1) // 128 threads per workgroup, 1 cluster per thread
28fn cull_clusters(
29    @builtin(workgroup_id) workgroup_id: vec3<u32>,
30    @builtin(num_workgroups) num_workgroups: vec3<u32>,
31    @builtin(local_invocation_index) local_invocation_index: u32,
32) {
33    // Calculate the cluster ID for this thread
34    let cluster_id = local_invocation_index + 128u * dot(workgroup_id, vec3(num_workgroups.x * num_workgroups.x, num_workgroups.x, 1u));
35    if cluster_id >= constants.scene_cluster_count { return; }
36
37#ifdef MESHLET_SECOND_CULLING_PASS
38    if !cluster_is_second_pass_candidate(cluster_id) { return; }
39#endif
40
41    // Check for instance culling
42    let instance_id = meshlet_cluster_instance_ids[cluster_id];
43#ifdef MESHLET_FIRST_CULLING_PASS
44    if should_cull_instance(instance_id) { return; }
45#endif
46
47    // Calculate world-space culling bounding sphere for the cluster
48    let instance_uniform = meshlet_instance_uniforms[instance_id];
49    let meshlet_id = meshlet_cluster_meshlet_ids[cluster_id];
50    let world_from_local = affine3_to_square(instance_uniform.world_from_local);
51    let world_scale = max(length(world_from_local[0]), max(length(world_from_local[1]), length(world_from_local[2])));
52    let bounding_spheres = meshlet_bounding_spheres[meshlet_id];
53    let culling_bounding_sphere_center = world_from_local * vec4(bounding_spheres.culling_sphere.center, 1.0);
54    let culling_bounding_sphere_radius = world_scale * bounding_spheres.culling_sphere.radius;
55
56#ifdef MESHLET_FIRST_CULLING_PASS
57    // Frustum culling
58    // TODO: Faster method from https://vkguide.dev/docs/gpudriven/compute_culling/#frustum-culling-function
59    for (var i = 0u; i < 6u; i++) {
60        if dot(view.frustum[i], culling_bounding_sphere_center) + culling_bounding_sphere_radius <= 0.0 {
61            return;
62        }
63    }
64
65    // Check LOD cut (cluster group error imperceptible, and parent group error not imperceptible)
66    let simplification_errors = unpack2x16float(meshlet_simplification_errors[meshlet_id]);
67    let lod_is_ok = lod_error_is_imperceptible(bounding_spheres.lod_group_sphere, simplification_errors.x, world_from_local, world_scale);
68    let parent_lod_is_ok = lod_error_is_imperceptible(bounding_spheres.lod_parent_group_sphere, simplification_errors.y, world_from_local, world_scale);
69    if !lod_is_ok || parent_lod_is_ok { return; }
70#endif
71
72    // Project the culling bounding sphere to view-space for occlusion culling
73#ifdef MESHLET_FIRST_CULLING_PASS
74    let previous_world_from_local = affine3_to_square(instance_uniform.previous_world_from_local);
75    let previous_world_from_local_scale = max(length(previous_world_from_local[0]), max(length(previous_world_from_local[1]), length(previous_world_from_local[2])));
76    let occlusion_culling_bounding_sphere_center = previous_world_from_local * vec4(bounding_spheres.culling_sphere.center, 1.0);
77    let occlusion_culling_bounding_sphere_radius = previous_world_from_local_scale * bounding_spheres.culling_sphere.radius;
78    let occlusion_culling_bounding_sphere_center_view_space = (previous_view.view_from_world * vec4(occlusion_culling_bounding_sphere_center.xyz, 1.0)).xyz;
79#else
80    let occlusion_culling_bounding_sphere_center = culling_bounding_sphere_center;
81    let occlusion_culling_bounding_sphere_radius = culling_bounding_sphere_radius;
82    let occlusion_culling_bounding_sphere_center_view_space = (view.view_from_world * vec4(occlusion_culling_bounding_sphere_center.xyz, 1.0)).xyz;
83#endif
84
85    var aabb = project_view_space_sphere_to_screen_space_aabb(occlusion_culling_bounding_sphere_center_view_space, occlusion_culling_bounding_sphere_radius);
86    let depth_pyramid_size_mip_0 = vec2<f32>(textureDimensions(depth_pyramid, 0));
87    var aabb_width_pixels = (aabb.z - aabb.x) * depth_pyramid_size_mip_0.x;
88    var aabb_height_pixels = (aabb.w - aabb.y) * depth_pyramid_size_mip_0.y;
89    let depth_level = max(0, i32(ceil(log2(max(aabb_width_pixels, aabb_height_pixels))))); // TODO: Naga doesn't like this being a u32
90    let depth_pyramid_size = vec2<f32>(textureDimensions(depth_pyramid, depth_level));
91    let aabb_top_left = vec2<u32>(aabb.xy * depth_pyramid_size);
92
93    let depth_quad_a = textureLoad(depth_pyramid, aabb_top_left, depth_level).x;
94    let depth_quad_b = textureLoad(depth_pyramid, aabb_top_left + vec2(1u, 0u), depth_level).x;
95    let depth_quad_c = textureLoad(depth_pyramid, aabb_top_left + vec2(0u, 1u), depth_level).x;
96    let depth_quad_d = textureLoad(depth_pyramid, aabb_top_left + vec2(1u, 1u), depth_level).x;
97    let occluder_depth = min(min(depth_quad_a, depth_quad_b), min(depth_quad_c, depth_quad_d));
98
99    // Check whether or not the cluster would be occluded if drawn
100    var cluster_visible: bool;
101    if view.clip_from_view[3][3] == 1.0 {
102        // Orthographic
103        let sphere_depth = view.clip_from_view[3][2] + (occlusion_culling_bounding_sphere_center_view_space.z + occlusion_culling_bounding_sphere_radius) * view.clip_from_view[2][2];
104        cluster_visible = sphere_depth >= occluder_depth;
105    } else {
106        // Perspective
107        let sphere_depth = -view.clip_from_view[3][2] / (occlusion_culling_bounding_sphere_center_view_space.z + occlusion_culling_bounding_sphere_radius);
108        cluster_visible = sphere_depth >= occluder_depth;
109    }
110
111    // Write if the cluster should be occlusion tested in the second pass
112#ifdef MESHLET_FIRST_CULLING_PASS
113    if !cluster_visible {
114        let bit = 1u << cluster_id % 32u;
115        atomicOr(&meshlet_second_pass_candidates[cluster_id / 32u], bit);
116    }
117#endif
118
119    // Cluster would be occluded if drawn, so don't setup a draw for it
120    if !cluster_visible { return; }
121
122    // Check how big the cluster is in screen space
123#ifdef MESHLET_FIRST_CULLING_PASS
124    let culling_bounding_sphere_center_view_space = (view.view_from_world * vec4(culling_bounding_sphere_center.xyz, 1.0)).xyz;
125    aabb = project_view_space_sphere_to_screen_space_aabb(culling_bounding_sphere_center_view_space, culling_bounding_sphere_radius);
126    aabb_width_pixels = (aabb.z - aabb.x) * view.viewport.z;
127    aabb_height_pixels = (aabb.w - aabb.y) * view.viewport.w;
128#endif
129    let cluster_is_small = all(vec2(aabb_width_pixels, aabb_height_pixels) < vec2(64.0));
130
131    // Let the hardware rasterizer handle near-plane clipping
132    let not_intersects_near_plane = dot(view.frustum[4u], culling_bounding_sphere_center) > culling_bounding_sphere_radius;
133
134    var buffer_slot: u32;
135    if cluster_is_small && not_intersects_near_plane {
136        // Append this cluster to the list for software rasterization
137        buffer_slot = atomicAdd(&meshlet_software_raster_indirect_args.x, 1u);
138    } else {
139        // Append this cluster to the list for hardware rasterization
140        buffer_slot = atomicAdd(&meshlet_hardware_raster_indirect_args.instance_count, 1u);
141        buffer_slot = constants.meshlet_raster_cluster_rightmost_slot - buffer_slot;
142    }
143    meshlet_raster_clusters[buffer_slot] = cluster_id;
144}
145
146// https://github.com/zeux/meshoptimizer/blob/1e48e96c7e8059321de492865165e9ef071bffba/demo/nanite.cpp#L115
147fn lod_error_is_imperceptible(lod_sphere: MeshletBoundingSphere, simplification_error: f32, world_from_local: mat4x4<f32>, world_scale: f32) -> bool {
148    let sphere_world_space = (world_from_local * vec4(lod_sphere.center, 1.0)).xyz;
149    let radius_world_space = world_scale * lod_sphere.radius;
150    let error_world_space = world_scale * simplification_error;
151
152    var projected_error = error_world_space;
153    if view.clip_from_view[3][3] != 1.0 {
154        // Perspective
155        let distance_to_closest_point_on_sphere = distance(sphere_world_space, view.world_position) - radius_world_space;
156        let distance_to_closest_point_on_sphere_clamped_to_znear = max(distance_to_closest_point_on_sphere, view.clip_from_view[3][2]);
157        projected_error /= distance_to_closest_point_on_sphere_clamped_to_znear;
158    }
159    projected_error *= view.clip_from_view[1][1] * 0.5;
160    projected_error *= view.viewport.w;
161
162    return projected_error < 1.0;
163}
164
165// https://zeux.io/2023/01/12/approximate-projected-bounds
166fn project_view_space_sphere_to_screen_space_aabb(cp: vec3<f32>, r: f32) -> vec4<f32> {
167    let inv_width = view.clip_from_view[0][0] * 0.5;
168    let inv_height = view.clip_from_view[1][1] * 0.5;
169    if view.clip_from_view[3][3] == 1.0 {
170        // Orthographic
171        let min_x = cp.x - r;
172        let max_x = cp.x + r;
173
174        let min_y = cp.y - r;
175        let max_y = cp.y + r;
176
177        return vec4(min_x * inv_width, 1.0 - max_y * inv_height, max_x * inv_width, 1.0 - min_y * inv_height);
178    } else {
179        // Perspective
180        let c = vec3(cp.xy, -cp.z);
181        let cr = c * r;
182        let czr2 = c.z * c.z - r * r;
183
184        let vx = sqrt(c.x * c.x + czr2);
185        let min_x = (vx * c.x - cr.z) / (vx * c.z + cr.x);
186        let max_x = (vx * c.x + cr.z) / (vx * c.z - cr.x);
187
188        let vy = sqrt(c.y * c.y + czr2);
189        let min_y = (vy * c.y - cr.z) / (vy * c.z + cr.y);
190        let max_y = (vy * c.y + cr.z) / (vy * c.z - cr.y);
191
192        return vec4(min_x * inv_width, -max_y * inv_height, max_x * inv_width, -min_y * inv_height) + vec4(0.5);
193    }
194}