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