1#import bevy_pbr::meshlet_bindings::{
2    scene_instance_count,
3    meshlet_global_cluster_count,
4    meshlet_instance_meshlet_counts,
5    meshlet_instance_meshlet_slice_starts,
6    meshlet_cluster_instance_ids,
7    meshlet_cluster_meshlet_ids,
8}
9
10/// Writes out instance_id and meshlet_id to the global buffers for each cluster in the scene.
11
12var<workgroup> cluster_slice_start_workgroup: u32;
13
14@compute
15@workgroup_size(1024, 1, 1) // 1024 threads per workgroup, 1 instance per workgroup
16fn fill_cluster_buffers(
17    @builtin(workgroup_id) workgroup_id: vec3<u32>,
18    @builtin(num_workgroups) num_workgroups: vec3<u32>,
19    @builtin(local_invocation_index) local_invocation_index: u32,
20) {
21    // Calculate the instance ID for this workgroup
22    var instance_id = workgroup_id.x + (workgroup_id.y * num_workgroups.x);
23    if instance_id >= scene_instance_count { return; }
24
25    let instance_meshlet_count = meshlet_instance_meshlet_counts[instance_id];
26    let instance_meshlet_slice_start = meshlet_instance_meshlet_slice_starts[instance_id];
27
28    // Reserve cluster slots for the instance and broadcast to the workgroup
29    if local_invocation_index == 0u {
30        cluster_slice_start_workgroup = atomicAdd(&meshlet_global_cluster_count, instance_meshlet_count);
31    }
32    let cluster_slice_start = workgroupUniformLoad(&cluster_slice_start_workgroup);
33
34    // Loop enough times to write out all the meshlets for the instance given that each thread writes 1 meshlet in each iteration
35    for (var clusters_written = 0u; clusters_written < instance_meshlet_count; clusters_written += 1024u) {
36        // Calculate meshlet ID within this instance's MeshletMesh to process for this thread
37        let meshlet_id_local = clusters_written + local_invocation_index;
38        if meshlet_id_local >= instance_meshlet_count { return; }
39
40        // Find the overall cluster ID in the global cluster buffer
41        let cluster_id = cluster_slice_start + meshlet_id_local;
42
43        // Find the overall meshlet ID in the global meshlet buffer
44        let meshlet_id = instance_meshlet_slice_start + meshlet_id_local;
45
46        // Write results to buffers
47        meshlet_cluster_instance_ids[cluster_id] = instance_id;
48        meshlet_cluster_meshlet_ids[cluster_id] = meshlet_id;
49    }
50}