,

src/pipelines/agents/agent-generation/agent-compaction.wgsl

@@ -15,22 +15,70 @@ struct Counters {
 @group(1) @binding(2) var<storage, read_write> counters: Counters;
 @group(1) @binding(3) var<storage, read_write> compactedAgents: array<Agent>;
 
+var<workgroup> workgroupAliveCount: atomic<u32>;
+var<workgroup> workgroupCompactedOffset: u32;
+var<workgroup> workgroupCopyCount: u32;
+
 @compute @workgroup_size(64)
 fn main(
   @builtin(global_invocation_id) global_id: vec3<u32>,
+  @builtin(local_invocation_id) local_id: vec3<u32>,
   @builtin(num_workgroups) workgroup_count: vec3<u32>
 ) {
   let id = get_id(global_id, workgroup_count);
 
-  if id >= settings.agentCount {
-    return;
+  if local_id.x == 0u {
+    atomicStore(&workgroupAliveCount, 0u);
   }
 
-  let agent = agents[id];
-  if agent.colorIndex < 0.0 {
-    return;
+  workgroupBarrier();
+
+  var localCompactedIndex = 0u;
+  if id < settings.agentCount {
+    let agent = agents[id];
+    if agent.colorIndex >= 0.0 {
+      localCompactedIndex = atomicAdd(&workgroupAliveCount, 1u);
+    }
   }
 
-  let compactedIndex = atomicAdd(&counters.aliveAgentCount, 1);
-  compactedAgents[compactedIndex] = agent;
+  workgroupBarrier();
+
+  if local_id.x == 0u {
+    let groupAliveCount = atomicLoad(&workgroupAliveCount);
+    if groupAliveCount > 0u {
+      workgroupCompactedOffset = atomicAdd(&counters.aliveAgentCount, groupAliveCount);
+    } else {
+      workgroupCompactedOffset = 0u;
+    }
+  }
+
+  workgroupBarrier();
+
+  if id < settings.agentCount {
+    let agent = agents[id];
+    if agent.colorIndex >= 0.0 {
+      compactedAgents[workgroupCompactedOffset + localCompactedIndex] = agent;
+    }
+  }
+}
+
+@compute @workgroup_size(64)
+fn copyCompactedAgents(
+  @builtin(global_invocation_id) global_id: vec3<u32>,
+  @builtin(local_invocation_id) local_id: vec3<u32>,
+  @builtin(num_workgroups) workgroup_count: vec3<u32>
+) {
+  let id = get_id(global_id, workgroup_count);
+
+  if local_id.x == 0u {
+    workgroupCopyCount = atomicLoad(&counters.aliveAgentCount);
+  }
+
+  workgroupBarrier();
+
+  if id >= workgroupCopyCount {
+    return;
+  }
+
+  agents[id] = compactedAgents[id];
 }