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
146// TODO: We should probably be using a POT HZB texture?
147fn occlusion_cull_screen_aabb(aabb: ScreenAabb, screen: vec2<f32>) -> bool {
148    let hzb_size = ceil(screen * 0.5);
149    let aabb_min = aabb.min.xy * hzb_size;
150    let aabb_max = aabb.max.xy * hzb_size;
151
152    let min_texel = vec2<u32>(max(aabb_min, vec2<f32>(0.0)));
153    let max_texel = vec2<u32>(min(aabb_max, hzb_size - 1.0));
154    let size = max_texel - min_texel;
155    let max_size = max(size.x, size.y);
156
157    // note: add 1 before max because the unsigned overflow behavior is intentional
158    // it wraps around firstLeadingBit(0) = ~0 to 0
159    // TODO: we actually sample a 4x4 block, so ideally this would be `max(..., 3u) - 3u`.
160    // However, since our HZB is not a power of two, we need to be extra-conservative to not over-cull, so we go up a mip.
161    var mip = max(firstLeadingBit(max_size) + 1u, 2u) - 2u;
162    
163    if any((max_texel >> vec2(mip)) > (min_texel >> vec2(mip)) + 3) {
164        mip += 1u;
165    }
166
167    let smin = min_texel >> vec2<u32>(mip);
168    let smax = max_texel >> vec2<u32>(mip);
169    
170    let curr_depth = sample_hzb(smin, smax, i32(mip));
171    return aabb.max.z <= curr_depth;
172}
173
174fn occlusion_cull_projection() -> mat4x4<f32> {
175#ifdef FIRST_CULLING_PASS
176    return view.clip_from_world;
177#else
178    return previous_view.clip_from_world;
179#endif
180}
181
182fn occlusion_cull_clip_from_local(instance_id: u32) -> mat4x4<f32> {
183#ifdef FIRST_CULLING_PASS
184    let prev_world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].previous_world_from_local);
185    return previous_view.clip_from_world * prev_world_from_local;
186#else
187    let world_from_local = affine3_to_square(meshlet_instance_uniforms[instance_id].world_from_local);
188    return view.clip_from_world * world_from_local;
189#endif
190}
191
192fn should_occlusion_cull_aabb(aabb: MeshletAabb, instance_id: u32) -> bool {
193    let projection = occlusion_cull_projection();
194    var near: f32;
195    if projection[3][3] == 1.0 {
196        near = projection[3][2] / projection[2][2];
197    } else {
198        near = projection[3][2];
199    }
200
201    let clip_from_local = occlusion_cull_clip_from_local(instance_id);
202    var screen_aabb = ScreenAabb(vec3<f32>(0.0), vec3<f32>(0.0));
203    if project_aabb(clip_from_local, near, aabb, &screen_aabb) {
204        return occlusion_cull_screen_aabb(screen_aabb, view.viewport.zw);
205    }
206    return false;
207}