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            var ndc_position = clip_position.xyz / clip_position.w;
61#ifdef DEPTH_CLAMP_ORTHO
62            ndc_position.z = 1.0 / clip_position.z;
63#endif
64            let viewport_position_xy = ndc_to_uv(ndc_position.xy) * view.viewport.zw;
65
66            // Write vertex to workgroup shared memory
67            viewport_vertices[vertex_id] = vec3(viewport_position_xy, ndc_position.z);
68        }
69    }
70    workgroupBarrier();
71
72    // Load 1 triangle's worth of vertex data per thread
73    let triangle_id = local_invocation_index;
74    if triangle_id >= get_meshlet_triangle_count(&meshlet) { return; }
75    let index_ids = meshlet.start_index_id + (triangle_id * 3u) + vec3(0u, 1u, 2u);
76    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]));
77    let vertex_0 = viewport_vertices[vertex_ids[2]];
78    let vertex_1 = viewport_vertices[vertex_ids[1]];
79    let vertex_2 = viewport_vertices[vertex_ids[0]];
80    let packed_ids = (cluster_id << 7u) | triangle_id;
81
82    // Backface culling
83    let triangle_double_area = edge_function(vertex_0.xy, vertex_1.xy, vertex_2.xy);
84    if triangle_double_area <= 0.0 { return; }
85
86    // Setup triangle gradients
87    let w_x = vec3(vertex_1.y - vertex_2.y, vertex_2.y - vertex_0.y, vertex_0.y - vertex_1.y);
88    let w_y = vec3(vertex_2.x - vertex_1.x, vertex_0.x - vertex_2.x, vertex_1.x - vertex_0.x);
89    let vertices_z = vec3(vertex_0.z, vertex_1.z, vertex_2.z) / triangle_double_area;
90    let z_x = dot(vertices_z, w_x);
91    let z_y = dot(vertices_z, w_y);
92
93    // Compute triangle bounding box
94    var min_x = floor(min3(vertex_0.x, vertex_1.x, vertex_2.x));
95    var min_y = floor(min3(vertex_0.y, vertex_1.y, vertex_2.y));
96    var max_x = ceil(max3(vertex_0.x, vertex_1.x, vertex_2.x));
97    var max_y = ceil(max3(vertex_0.y, vertex_1.y, vertex_2.y));
98    min_x = max(min_x, 0.0);
99    min_y = max(min_y, 0.0);
100    max_x = min(max_x, view.viewport.z - 1.0);
101    max_y = min(max_y, view.viewport.w - 1.0);
102
103    // Setup initial triangle equations
104    let starting_pixel = vec2(min_x, min_y) + 0.5;
105    var w_row = vec3(
106        edge_function(vertex_1.xy, vertex_2.xy, starting_pixel),
107        edge_function(vertex_2.xy, vertex_0.xy, starting_pixel),
108        edge_function(vertex_0.xy, vertex_1.xy, starting_pixel),
109    );
110    var z_row = dot(vertices_z, w_row);
111
112    // Rasterize triangle
113    if subgroupAny(max_x - min_x > 4.0) {
114        // Scanline setup
115        let edge_012 = -w_x;
116        let open_edge = edge_012 < vec3(0.0);
117        let inverse_edge_012 = select(1.0 / edge_012, vec3(1e8), edge_012 == vec3(0.0));
118        let max_x_diff = vec3(max_x - min_x);
119        for (var y = min_y; y <= max_y; y += 1.0) {
120            // Calculate start and end X interval for pixels in this row within the triangle
121            let cross_x = w_row * inverse_edge_012;
122            let min_x2 = select(vec3(0.0), cross_x, open_edge);
123            let max_x2 = select(cross_x, max_x_diff, open_edge);
124            var x0 = ceil(max3(min_x2[0], min_x2[1], min_x2[2]));
125            var x1 = min3(max_x2[0], max_x2[1], max_x2[2]);
126
127            var w = w_row + w_x * x0;
128            var z = z_row + z_x * x0;
129            x0 += min_x;
130            x1 += min_x;
131
132            // Iterate scanline X interval
133            for (var x = x0; x <= x1; x += 1.0) {
134                // Check if point at pixel is within triangle (TODO: this shouldn't be needed, but there's bugs without it)
135                if min3(w[0], w[1], w[2]) >= 0.0 {
136                    write_visibility_buffer_pixel(x, y, z, packed_ids);
137                }
138
139                // Increment triangle equations along the X-axis
140                w += w_x;
141                z += z_x;
142            }
143
144            // Increment triangle equations along the Y-axis
145            w_row += w_y;
146            z_row += z_y;
147        }
148    } else {
149        // Iterate over every pixel in the triangle's bounding box
150        for (var y = min_y; y <= max_y; y += 1.0) {
151            var w = w_row;
152            var z = z_row;
153
154            for (var x = min_x; x <= max_x; x += 1.0) {
155                // Check if point at pixel is within triangle
156                if min3(w[0], w[1], w[2]) >= 0.0 {
157                    write_visibility_buffer_pixel(x, y, z, packed_ids);
158                }
159
160                // Increment triangle equations along the X-axis
161                w += w_x;
162                z += z_x;
163            }
164
165            // Increment triangle equations along the Y-axis
166            w_row += w_y;
167            z_row += z_y;
168        }
169    }
170}
171
172fn write_visibility_buffer_pixel(x: f32, y: f32, z: f32, packed_ids: u32) {
173    let frag_coord_1d = u32(y * view.viewport.z + x);
174
175#ifdef MESHLET_VISIBILITY_BUFFER_RASTER_PASS_OUTPUT
176    let depth = bitcast<u32>(z);
177    let visibility = (u64(depth) << 32u) | u64(packed_ids);
178    atomicMax(&meshlet_visibility_buffer[frag_coord_1d], visibility);
179#else ifdef DEPTH_CLAMP_ORTHO
180    let depth = bitcast<u32>(1.0 / z);
181    atomicMax(&meshlet_visibility_buffer[frag_coord_1d], depth);
182#else
183    let depth = bitcast<u32>(z);
184    atomicMax(&meshlet_visibility_buffer[frag_coord_1d], depth);
185#endif
186}
187
188fn edge_function(a: vec2<f32>, b: vec2<f32>, c: vec2<f32>) -> f32 {
189    return (b.x - a.x) * (c.y - a.y) - (b.y - a.y) * (c.x - a.x);
190}
191
192fn min3(a: f32, b: f32, c: f32) -> f32 {
193    return min(a, min(b, c));
194}
195
196fn max3(a: f32, b: f32, c: f32) -> f32 {
197    return max(a, max(b, c));
198}