1#import bevy_pbr::{
2    meshlet_bindings::{
3        meshlet_cluster_meshlet_ids,
4        meshlets,
5        meshlet_cluster_instance_ids,
6        meshlet_instance_uniforms,
7        meshlet_raster_clusters,
8        meshlet_software_raster_cluster_count,
9        meshlet_visibility_buffer,
10        view,
11        get_meshlet_vertex_count,
12        get_meshlet_triangle_count,
13        get_meshlet_vertex_id,
14        get_meshlet_vertex_position,
15    },
16    mesh_functions::mesh_position_local_to_world,
17    view_transformations::ndc_to_uv,
18}
19#import bevy_render::maths::affine3_to_square
20
21/// Compute shader for rasterizing small clusters into a visibility buffer.
22
23// TODO: Fixed-point math and top-left rule
24
25var<workgroup> viewport_vertices: array<vec3f, 255>;
26
27@compute
28@workgroup_size(128, 1, 1) // 128 threads per workgroup, 1-2 vertices per thread, 1 triangle per thread, 1 cluster per workgroup
29fn rasterize_cluster(
30    @builtin(workgroup_id) workgroup_id: vec3<u32>,
31    @builtin(local_invocation_index) local_invocation_index: u32,
32#ifdef MESHLET_2D_DISPATCH
33    @builtin(num_workgroups) num_workgroups: vec3<u32>,
34#endif
35) {
36    var workgroup_id_1d = workgroup_id.x;
37
38#ifdef MESHLET_2D_DISPATCH
39    workgroup_id_1d += workgroup_id.y * num_workgroups.x;
40    if workgroup_id_1d >= meshlet_software_raster_cluster_count { return; }
41#endif
42
43    let cluster_id = meshlet_raster_clusters[workgroup_id_1d];
44    let meshlet_id = meshlet_cluster_meshlet_ids[cluster_id];
45    var meshlet = meshlets[meshlet_id];
46
47    let instance_id = meshlet_cluster_instance_ids[cluster_id];
48    let instance_uniform = meshlet_instance_uniforms[instance_id];
49    let world_from_local = affine3_to_square(instance_uniform.world_from_local);
50
51    // Load and project 1 vertex per thread, and then again if there are more than 128 vertices in the meshlet
52    for (var i = 0u; i <= 128u; i += 128u) {
53        let vertex_id = local_invocation_index + i;
54        if vertex_id < get_meshlet_vertex_count(&meshlet) {
55            let vertex_position = get_meshlet_vertex_position(&meshlet, vertex_id);
56
57            // Project vertex to viewport space
58            let world_position = mesh_position_local_to_world(world_from_local, vec4(vertex_position, 1.0));
59            let clip_position = view.clip_from_world * vec4(world_position.xyz, 1.0);
60            let ndc_position = clip_position.xyz / clip_position.w;
61            let viewport_position_xy = ndc_to_uv(ndc_position.xy) * view.viewport.zw;
62
63            // Write vertex to workgroup shared memory
64            viewport_vertices[vertex_id] = vec3(viewport_position_xy, ndc_position.z);
65        }
66    }
67    workgroupBarrier();
68
69    // Load 1 triangle's worth of vertex data per thread
70    let triangle_id = local_invocation_index;
71    if triangle_id >= get_meshlet_triangle_count(&meshlet) { return; }
72    let index_ids = meshlet.start_index_id + (triangle_id * 3u) + vec3(0u, 1u, 2u);
73    let vertex_ids = vec3(get_meshlet_vertex_id(index_ids[0]), get_meshlet_vertex_id(index_ids[1]), get_meshlet_vertex_id(index_ids[2]));
74    let vertex_0 = viewport_vertices[vertex_ids[2]];
75    let vertex_1 = viewport_vertices[vertex_ids[1]];
76    let vertex_2 = viewport_vertices[vertex_ids[0]];
77    let packed_ids = (cluster_id << 7u) | triangle_id;
78
79    // Backface culling
80    let triangle_double_area = edge_function(vertex_0.xy, vertex_1.xy, vertex_2.xy);
81    if triangle_double_area <= 0.0 { return; }
82
83    // Setup triangle gradients
84    let w_x = vec3(vertex_1.y - vertex_2.y, vertex_2.y - vertex_0.y, vertex_0.y - vertex_1.y);
85    let w_y = vec3(vertex_2.x - vertex_1.x, vertex_0.x - vertex_2.x, vertex_1.x - vertex_0.x);
86    let vertices_z = vec3(vertex_0.z, vertex_1.z, vertex_2.z) / triangle_double_area;
87    let z_x = dot(vertices_z, w_x);
88    let z_y = dot(vertices_z, w_y);
89
90    // Compute triangle bounding box
91    var min_x = floor(min3(vertex_0.x, vertex_1.x, vertex_2.x));
92    var min_y = floor(min3(vertex_0.y, vertex_1.y, vertex_2.y));
93    var max_x = ceil(max3(vertex_0.x, vertex_1.x, vertex_2.x));
94    var max_y = ceil(max3(vertex_0.y, vertex_1.y, vertex_2.y));
95    min_x = max(min_x, 0.0);
96    min_y = max(min_y, 0.0);
97    max_x = min(max_x, view.viewport.z - 1.0);
98    max_y = min(max_y, view.viewport.w - 1.0);
99
100    // Setup initial triangle equations
101    let starting_pixel = vec2(min_x, min_y) + 0.5;
102    var w_row = vec3(
103        edge_function(vertex_1.xy, vertex_2.xy, starting_pixel),
104        edge_function(vertex_2.xy, vertex_0.xy, starting_pixel),
105        edge_function(vertex_0.xy, vertex_1.xy, starting_pixel),
106    );
107    var z_row = dot(vertices_z, w_row);
108
109    // Rasterize triangle
110    if subgroupAny(max_x - min_x > 4.0) {
111        // Scanline setup
112        let edge_012 = -w_x;
113        let open_edge = edge_012 < vec3(0.0);
114        let inverse_edge_012 = select(1.0 / edge_012, vec3(1e8), edge_012 == vec3(0.0));
115        let max_x_diff = vec3(max_x - min_x);
116        for (var y = min_y; y <= max_y; y += 1.0) {
117            // Calculate start and end X interval for pixels in this row within the triangle
118            let cross_x = w_row * inverse_edge_012;
119            let min_x2 = select(vec3(0.0), cross_x, open_edge);
120            let max_x2 = select(cross_x, max_x_diff, open_edge);
121            var x0 = ceil(max3(min_x2[0], min_x2[1], min_x2[2]));
122            var x1 = min3(max_x2[0], max_x2[1], max_x2[2]);
123
124            var w = w_row + w_x * x0;
125            var z = z_row + z_x * x0;
126            x0 += min_x;
127            x1 += min_x;
128
129            // Iterate scanline X interval
130            for (var x = x0; x <= x1; x += 1.0) {
131                // Check if point at pixel is within triangle (TODO: this shouldn't be needed, but there's bugs without it)
132                if min3(w[0], w[1], w[2]) >= 0.0 {
133                    write_visibility_buffer_pixel(x, y, z, packed_ids);
134                }
135
136                // Increment triangle equations along the X-axis
137                w += w_x;
138                z += z_x;
139            }
140
141            // Increment triangle equations along the Y-axis
142            w_row += w_y;
143            z_row += z_y;
144        }
145    } else {
146        // Iterate over every pixel in the triangle's bounding box
147        for (var y = min_y; y <= max_y; y += 1.0) {
148            var w = w_row;
149            var z = z_row;
150
151            for (var x = min_x; x <= max_x; x += 1.0) {
152                // Check if point at pixel is within triangle
153                if min3(w[0], w[1], w[2]) >= 0.0 {
154                    write_visibility_buffer_pixel(x, y, z, packed_ids);
155                }
156
157                // Increment triangle equations along the X-axis
158                w += w_x;
159                z += z_x;
160            }
161
162            // Increment triangle equations along the Y-axis
163            w_row += w_y;
164            z_row += z_y;
165        }
166    }
167}
168
169fn write_visibility_buffer_pixel(x: f32, y: f32, z: f32, packed_ids: u32) {
170    let depth = bitcast<u32>(z);
171#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
172    let visibility = (u64(depth) << 32u) | u64(packed_ids);
173#else
174    let visibility = depth;
175#endif
176    textureAtomicMax(meshlet_visibility_buffer, vec2(u32(x), u32(y)), visibility);
177}
178
179fn edge_function(a: vec2<f32>, b: vec2<f32>, c: vec2<f32>) -> f32 {
180    return (b.x - a.x) * (c.y - a.y) - (b.y - a.y) * (c.x - a.x);
181}
182
183fn min3(a: f32, b: f32, c: f32) -> f32 {
184    return min(a, min(b, c));
185}
186
187fn max3(a: f32, b: f32, c: f32) -> f32 {
188    return max(a, max(b, c));
189}