1#define_import_path bevy_pbr::meshlet_cull_shared
2
3#import bevy_pbr::meshlet_bindings::{
4    MeshletAabb,
5    DispatchIndirectArgs,
6    InstancedOffset,
7    depth_pyramid,
8    view,
9    previous_view,
10    meshlet_instance_uniforms,
11}
12#import bevy_render::maths::affine3_to_square
13
14// https://github.com/zeux/meshoptimizer/blob/1e48e96c7e8059321de492865165e9ef071bffba/demo/nanite.cpp#L115
15fn lod_error_is_imperceptible(lod_sphere: vec4<f32>, simplification_error: f32, instance_id: u32) -> bool {
16    let world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].world_from_local);
17    let world_scale = max(length(world_from_local[0]), max(length(world_from_local[1]), length(world_from_local[2])));
18    let camera_pos = view.world_position;
19
20    let projection = view.clip_from_view;
21    if projection[3][3] == 1.0 {
22        // Orthographic
23        let world_error = simplification_error * world_scale;
24        let proj = projection[1][1];
25        let height = 2.0 / proj;
26        let norm_error = world_error / height;
27        return norm_error * view.viewport.w < 1.0;
28    } else {
29        // Perspective
30        var near = projection[3][2];
31        let world_sphere_center = (world_from_local * vec4<f32>(lod_sphere.xyz, 1.0)).xyz;
32        let world_sphere_radius = lod_sphere.w * world_scale;
33        let d_pos = world_sphere_center - camera_pos;
34        let d = sqrt(dot(d_pos, d_pos)) - world_sphere_radius;
35        let norm_error = simplification_error / max(d, near) * projection[1][1] * 0.5;
36        return norm_error * view.viewport.w < 1.0;
37    }
38}
39
40fn normalize_plane(p: vec4<f32>) -> vec4<f32> {
41    return p / length(p.xyz);
42}
43
44// https://fgiesen.wordpress.com/2012/08/31/frustum-planes-from-the-projection-matrix/
45// https://fgiesen.wordpress.com/2010/10/17/view-frustum-culling/
46fn aabb_in_frustum(aabb: MeshletAabb, instance_id: u32) -> bool {
47    let world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].world_from_local);
48    let clip_from_local = view.clip_from_world * world_from_local;
49    let row_major = transpose(clip_from_local);
50    let planes = array(
51        row_major[3] + row_major[0],
52        row_major[3] - row_major[0],
53        row_major[3] + row_major[1],
54        row_major[3] - row_major[1],
55        row_major[2],
56    );
57
58    for (var i = 0; i < 5; i++) {
59        let plane = normalize_plane(planes[i]);
60        let flipped = aabb.half_extent * sign(plane.xyz);
61        if dot(aabb.center + flipped, plane.xyz) <= -plane.w {
62            return false;
63        }
64    }
65    return true;
66}
67
68struct ScreenAabb {
69    min: vec3<f32>,
70    max: vec3<f32>,
71}
72
73fn min8(a: vec3<f32>, b: vec3<f32>, c: vec3<f32>, d: vec3<f32>, e: vec3<f32>, f: vec3<f32>, g: vec3<f32>, h: vec3<f32>) -> vec3<f32> {
74    return min(min(min(a, b), min(c, d)), min(min(e, f), min(g, h)));
75}
76
77fn max8(a: vec3<f32>, b: vec3<f32>, c: vec3<f32>, d: vec3<f32>, e: vec3<f32>, f: vec3<f32>, g: vec3<f32>, h: vec3<f32>) -> vec3<f32> {
78    return max(max(max(a, b), max(c, d)), max(max(e, f), max(g, h)));
79}
80
81fn min8_4(a: vec4<f32>, b: vec4<f32>, c: vec4<f32>, d: vec4<f32>, e: vec4<f32>, f: vec4<f32>, g: vec4<f32>, h: vec4<f32>) -> vec4<f32> {
82    return min(min(min(a, b), min(c, d)), min(min(e, f), min(g, h)));
83}
84
85// https://zeux.io/2023/01/12/approximate-projected-bounds/
86fn project_aabb(clip_from_local: mat4x4<f32>, near: f32, aabb: MeshletAabb, out: ptr<function, ScreenAabb>) -> bool {
87    let extent = aabb.half_extent * 2.0;
88    let sx = clip_from_local * vec4<f32>(extent.x, 0.0, 0.0, 0.0);
89    let sy = clip_from_local * vec4<f32>(0.0, extent.y, 0.0, 0.0);
90    let sz = clip_from_local * vec4<f32>(0.0, 0.0, extent.z, 0.0);
91
92    let p0 = clip_from_local * vec4<f32>(aabb.center - aabb.half_extent, 1.0);
93    let p1 = p0 + sz;
94    let p2 = p0 + sy;
95    let p3 = p2 + sz;
96    let p4 = p0 + sx;
97    let p5 = p4 + sz;
98    let p6 = p4 + sy;
99    let p7 = p6 + sz;
100
101    let depth = min8_4(p0, p1, p2, p3, p4, p5, p6, p7).w;
102    // do not occlusion cull if we are inside the aabb
103    if depth < near {
104        return false;
105    }
106
107    let dp0 = p0.xyz / p0.w;
108    let dp1 = p1.xyz / p1.w;
109    let dp2 = p2.xyz / p2.w;
110    let dp3 = p3.xyz / p3.w;
111    let dp4 = p4.xyz / p4.w;
112    let dp5 = p5.xyz / p5.w;
113    let dp6 = p6.xyz / p6.w;
114    let dp7 = p7.xyz / p7.w;
115    let min = min8(dp0, dp1, dp2, dp3, dp4, dp5, dp6, dp7);
116    let max = max8(dp0, dp1, dp2, dp3, dp4, dp5, dp6, dp7);
117    var vaabb = vec4<f32>(min.xy, max.xy);
118    // convert ndc to texture coordinates by rescaling and flipping Y
119    vaabb = vaabb.xwzy * vec4<f32>(0.5, -0.5, 0.5, -0.5) + 0.5;
120    (*out).min = vec3<f32>(vaabb.xy, min.z);
121    (*out).max = vec3<f32>(vaabb.zw, max.z);
122    return true;
123}
124
125fn sample_hzb(smin: vec2<u32>, smax: vec2<u32>, mip: i32) -> f32 {
126    let texel = vec4<u32>(0, 1, 2, 3);
127    let sx = min(smin.x + texel, smax.xxxx);
128    let sy = min(smin.y + texel, smax.yyyy);
129    // TODO: switch to min samplers when wgpu has them
130    // sampling 16 times a finer mip is worth the extra cost for better culling
131    let a = sample_hzb_row(sx, sy.x, mip);
132    let b = sample_hzb_row(sx, sy.y, mip);
133    let c = sample_hzb_row(sx, sy.z, mip);
134    let d = sample_hzb_row(sx, sy.w, mip);
135    return min(min(a, b), min(c, d));
136}
137
138fn sample_hzb_row(sx: vec4<u32>, sy: u32, mip: i32) -> f32 {
139    let a = textureLoad(depth_pyramid, vec2(sx.x, sy), mip).x;
140    let b = textureLoad(depth_pyramid, vec2(sx.y, sy), mip).x;
141    let c = textureLoad(depth_pyramid, vec2(sx.z, sy), mip).x;
142    let d = textureLoad(depth_pyramid, vec2(sx.w, sy), mip).x;
143    return min(min(a, b), min(c, d));
144}
145
146fn occlusion_cull_screen_aabb(aabb: ScreenAabb, screen: vec2<f32>) -> bool {
147    let hzb_size = vec2<f32>(textureDimensions(depth_pyramid).xy);
148    let aabb_min = aabb.min.xy * hzb_size;
149    let aabb_max = aabb.max.xy * hzb_size;
150
151    let min_texel = vec2<u32>(max(aabb_min, vec2<f32>(0.0)));
152    let max_texel = vec2<u32>(min(aabb_max, hzb_size - 1.0));
153    let size = max_texel - min_texel;
154    let max_size = max(size.x, size.y);
155
156    // note: add 1 before max because the unsigned overflow behavior is intentional
157    // it wraps around firstLeadingBit(0) = ~0 to 0
158    // TODO: we actually sample a 4x4 block, so ideally this would be `max(..., 3u) - 3u`.
159    var mip = max(firstLeadingBit(max_size) + 1u, 2u) - 2u;
160    
161    if any((max_texel >> vec2(mip)) > (min_texel >> vec2(mip)) + 3) {
162        mip += 1u;
163    }
164
165    let smin = min_texel >> vec2<u32>(mip);
166    let smax = max_texel >> vec2<u32>(mip);
167    
168    let curr_depth = sample_hzb(smin, smax, i32(mip));
169    return aabb.max.z <= curr_depth;
170}
171
172fn occlusion_cull_projection() -> mat4x4<f32> {
173#ifdef MESHLET_FIRST_CULLING_PASS
174    return previous_view.clip_from_world;
175#else
176    return view.clip_from_world;
177#endif
178}
179
180fn occlusion_cull_clip_from_local(instance_id: u32) -> mat4x4<f32> {
181#ifdef MESHLET_FIRST_CULLING_PASS
182    let prev_world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].previous_world_from_local);
183    return previous_view.clip_from_world * prev_world_from_local;
184#else
185    let world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].world_from_local);
186    return view.clip_from_world * world_from_local;
187#endif
188}
189
190fn should_occlusion_cull_aabb(aabb: MeshletAabb, instance_id: u32) -> bool {
191    let projection = occlusion_cull_projection();
192    var near: f32;
193    if projection[3][3] == 1.0 {
194        near = projection[3][2] / projection[2][2];
195    } else {
196        near = projection[3][2];
197    }
198
199    let clip_from_local = occlusion_cull_clip_from_local(instance_id);
200    var screen_aabb = ScreenAabb(vec3<f32>(0.0), vec3<f32>(0.0));
201    if project_aabb(clip_from_local, near, aabb, &screen_aabb) {
202        return occlusion_cull_screen_aabb(screen_aabb, view.viewport.zw);
203    }
204    return false;
205}