80 lines
1.8 KiB
WebGPU Shading Language
80 lines
1.8 KiB
WebGPU Shading Language
struct Settings {
|
|
agentCount: u32,
|
|
padding0: u32,
|
|
padding1: u32,
|
|
padding2: u32,
|
|
};
|
|
|
|
struct Counters {
|
|
aliveAgentCount: atomic<u32>,
|
|
};
|
|
|
|
@group(1) @binding(0) var<uniform> settings: Settings;
|
|
@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>
|
|
) {
|
|
let id = get_id(global_id);
|
|
|
|
if local_id.x == 0u {
|
|
atomicStore(&workgroupAliveCount, 0u);
|
|
}
|
|
|
|
workgroupBarrier();
|
|
|
|
var localCompactedIndex = 0u;
|
|
if id < settings.agentCount {
|
|
let agent = agents[id];
|
|
if agent.colorIndex >= 0.0 {
|
|
localCompactedIndex = atomicAdd(&workgroupAliveCount, 1u);
|
|
}
|
|
}
|
|
|
|
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>
|
|
) {
|
|
let id = get_id(global_id);
|
|
|
|
if local_id.x == 0u {
|
|
workgroupCopyCount = atomicLoad(&counters.aliveAgentCount);
|
|
}
|
|
|
|
workgroupBarrier();
|
|
|
|
if id >= workgroupCopyCount {
|
|
return;
|
|
}
|
|
|
|
agents[id] = compactedAgents[id];
|
|
}
|