,

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];
 }

src/pipelines/agents/agent-generation/agent-generation-pipeline.test.ts

@@ -0,0 +1,235 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import type { CommonState } from '../../common-state/common-state';
+import { AGENT_SIZE_IN_BYTES } from './agent';
+import { AgentGenerationPipeline } from './agent-generation-pipeline';
+
+const installGpuConstants = () => {
+  Object.defineProperties(globalThis, {
+    GPUBufferUsage: {
+      configurable: true,
+      value: {
+        MAP_READ: 1,
+        COPY_DST: 2,
+        COPY_SRC: 4,
+        STORAGE: 8,
+        UNIFORM: 16,
+      },
+    },
+    GPUMapMode: {
+      configurable: true,
+      value: {
+        READ: 1,
+      },
+    },
+    GPUShaderStage: {
+      configurable: true,
+      value: {
+        COMPUTE: 1,
+      },
+    },
+  });
+};
+
+type CopyCall = {
+  source: GPUBuffer;
+  sourceOffset: number;
+  destination: GPUBuffer;
+  destinationOffset: number;
+  size: number;
+};
+
+type DispatchCall = {
+  entryPoint: string;
+  workgroups: [number, number, number];
+};
+
+type FakePipeline = {
+  entryPoint: string;
+};
+
+class FakeBuffer {
+  private readonly mappedRange: ArrayBuffer;
+
+  public readonly destroy = vi.fn();
+  public readonly mapAsync = vi.fn(async (_mode: number) => undefined);
+  public readonly getMappedRange = vi.fn(() => this.mappedRange);
+  public readonly unmap = vi.fn();
+
+  public constructor(
+    public readonly label: string,
+    size: number,
+    mappedValue = 0
+  ) {
+    this.mappedRange = new ArrayBuffer(Math.max(size, Uint32Array.BYTES_PER_ELEMENT));
+    new Uint32Array(this.mappedRange)[0] = mappedValue;
+  }
+}
+
+class FakeComputePass {
+  private pipeline: FakePipeline | null = null;
+
+  public readonly setPipeline = vi.fn((pipeline: GPUComputePipeline) => {
+    this.pipeline = pipeline as unknown as FakePipeline;
+  });
+  public readonly setBindGroup = vi.fn(
+    (_index: number, _bindGroup: GPUBindGroup) => undefined
+  );
+  public readonly dispatchWorkgroups = vi.fn((x: number, y = 1, z = 1) => {
+    this.device.dispatchCalls.push({
+      entryPoint: this.pipeline?.entryPoint ?? 'unset',
+      workgroups: [x, y, z],
+    });
+  });
+  public readonly end = vi.fn();
+
+  public constructor(private readonly device: FakeDevice) {}
+}
+
+class FakeCommandEncoder {
+  public readonly beginComputePass = vi.fn(() => new FakeComputePass(this.device));
+  public readonly copyBufferToBuffer = vi.fn(
+    (
+      source: GPUBuffer,
+      sourceOffset: number,
+      destination: GPUBuffer,
+      destinationOffset: number,
+      size: number
+    ) => {
+      this.device.copyCalls.push({
+        source,
+        sourceOffset,
+        destination,
+        destinationOffset,
+        size,
+      });
+    }
+  );
+  public readonly finish = vi.fn(() => ({}) as GPUCommandBuffer);
+
+  public constructor(private readonly device: FakeDevice) {}
+}
+
+class FakeQueue {
+  public readonly writeBuffer = vi.fn(
+    (_buffer: GPUBuffer, _offset: number, _data: BufferSource) => undefined
+  );
+  public readonly submit = vi.fn(
+    (_commandBuffers: Iterable<GPUCommandBuffer>) => undefined
+  );
+}
+
+class FakeShaderModule {
+  public readonly getCompilationInfo = vi.fn(async () => ({
+    messages: [],
+  }));
+}
+
+class FakeDevice {
+  public readonly copyCalls: Array<CopyCall> = [];
+  public readonly dispatchCalls: Array<DispatchCall> = [];
+  public readonly createdComputeEntryPoints: Array<string> = [];
+  public readonly limits = {
+    maxBufferSize: 1024 * 1024 * 1024,
+    maxComputeWorkgroupsPerDimension: 65_535,
+  };
+  public readonly queue = new FakeQueue();
+
+  private bufferIndex = 0;
+
+  public readonly createBindGroupLayout = vi.fn(
+    (_descriptor: GPUBindGroupLayoutDescriptor) => ({}) as GPUBindGroupLayout
+  );
+  public readonly createBuffer = vi.fn((descriptor: GPUBufferDescriptor) => {
+    const label =
+      ['agents', 'compactedAgents', 'counters', 'countersStaging', 'uniforms'][
+        this.bufferIndex
+      ] ?? `buffer${this.bufferIndex}`;
+    this.bufferIndex += 1;
+
+    const isMappedReadBuffer = (Number(descriptor.usage) & GPUBufferUsage.MAP_READ) !== 0;
+
+    return new FakeBuffer(
+      label,
+      Number(descriptor.size),
+      isMappedReadBuffer ? this.compactedCount : 0
+    ) as unknown as GPUBuffer;
+  });
+  public readonly createBindGroup = vi.fn(
+    (_descriptor: GPUBindGroupDescriptor) => ({}) as GPUBindGroup
+  );
+  public readonly createPipelineLayout = vi.fn(
+    (_descriptor: GPUPipelineLayoutDescriptor) => ({}) as GPUPipelineLayout
+  );
+  public readonly createShaderModule = vi.fn(
+    (_descriptor: GPUShaderModuleDescriptor) =>
+      new FakeShaderModule() as unknown as GPUShaderModule
+  );
+  public readonly createComputePipeline = vi.fn(
+    (descriptor: GPUComputePipelineDescriptor) => {
+      const pipeline = {
+        entryPoint: descriptor.compute.entryPoint ?? 'main',
+      };
+      this.createdComputeEntryPoints.push(pipeline.entryPoint);
+      return pipeline as unknown as GPUComputePipeline;
+    }
+  );
+  public readonly createCommandEncoder = vi.fn(() => new FakeCommandEncoder(this));
+
+  public constructor(private readonly compactedCount: number) {}
+}
+
+const createPipeline = (compactedCount: number) => {
+  installGpuConstants();
+
+  const device = new FakeDevice(compactedCount);
+  const commonState = {
+    bindGroupLayout: {} as GPUBindGroupLayout,
+    execute: vi.fn(),
+  } as unknown as CommonState;
+
+  return {
+    device,
+    pipeline: new AgentGenerationPipeline(
+      device as unknown as GPUDevice,
+      commonState,
+      1024
+    ),
+  };
+};
+
+describe('AgentGenerationPipeline compaction', () => {
+  it('copies compacted agents back with compute instead of a full agent buffer copy', async () => {
+    const agentCount = 10;
+    const { device, pipeline } = createPipeline(3);
+
+    await expect(pipeline.compactAgents(agentCount)).resolves.toBe(3);
+
+    expect(device.createdComputeEntryPoints).toContain('copyCompactedAgents');
+    expect(device.dispatchCalls.map((call) => call.entryPoint)).toEqual([
+      'main',
+      'copyCompactedAgents',
+    ]);
+    expect(device.copyCalls.map((call) => call.size)).toEqual([
+      Uint32Array.BYTES_PER_ELEMENT,
+    ]);
+    expect(
+      device.copyCalls.some((call) => call.size === agentCount * AGENT_SIZE_IN_BYTES)
+    ).toBe(false);
+    expect(device.queue.submit).toHaveBeenCalledTimes(1);
+
+    pipeline.destroy();
+  });
+
+  it('does not encode work for empty compaction requests', async () => {
+    const { device, pipeline } = createPipeline(0);
+
+    await expect(pipeline.compactAgents(0)).resolves.toBe(0);
+
+    expect(device.dispatchCalls).toEqual([]);
+    expect(device.copyCalls).toEqual([]);
+    expect(device.queue.submit).not.toHaveBeenCalled();
+
+    pipeline.destroy();
+  });
+});

src/pipelines/agents/agent-generation/agent-generation-pipeline.ts

@@ -26,6 +26,7 @@ export class AgentGenerationPipeline {
   private readonly countingPipeline: GPUComputePipeline;
   private readonly resizePipeline: GPUComputePipeline;
   private readonly compactionPipeline: GPUComputePipeline;
+  private readonly compactedAgentsCopyPipeline: GPUComputePipeline;
 
   public readonly agentsBuffer: GPUBuffer;
   private readonly compactedAgentsBuffer: GPUBuffer;
@@ -109,7 +110,7 @@ export class AgentGenerationPipeline {
 
     this.compactedAgentsBuffer = this.device.createBuffer({
       size: this.maxAgentCount * AGENT_SIZE_IN_BYTES,
-      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
+      usage: GPUBufferUsage.STORAGE,
     });
 
     this.countersBuffer = this.device.createBuffer({
@@ -216,20 +217,32 @@ export class AgentGenerationPipeline {
       },
     });
 
+    const compactionModule = smartCompile(
+      device,
+      CommonState.shaderCode,
+      agentSchema,
+      compactionShader
+    );
+
     this.compactionPipeline = device.createComputePipeline({
       layout: device.createPipelineLayout({
         bindGroupLayouts: [commonState.bindGroupLayout, this.compactionBindGroupLayout],
       }),
       compute: {
-        module: smartCompile(
-          device,
-          CommonState.shaderCode,
-          agentSchema,
-          compactionShader
-        ),
+        module: compactionModule,
         entryPoint: 'main',
       },
     });
+
+    this.compactedAgentsCopyPipeline = device.createComputePipeline({
+      layout: device.createPipelineLayout({
+        bindGroupLayouts: [commonState.bindGroupLayout, this.compactionBindGroupLayout],
+      }),
+      compute: {
+        module: compactionModule,
+        entryPoint: 'copyCompactedAgents',
+      },
+    });
   }
 
   public get maxAgentCount(): number {
@@ -364,13 +377,19 @@ export class AgentGenerationPipeline {
     );
     passEncoder.end();
 
-    commandEncoder.copyBufferToBuffer(
-      this.compactedAgentsBuffer,
-      0,
-      this.agentsBuffer,
-      0,
-      agentCount * AGENT_SIZE_IN_BYTES
+    const copyPassEncoder = commandEncoder.beginComputePass();
+    copyPassEncoder.setPipeline(this.compactedAgentsCopyPipeline);
+    this.commonState.execute(copyPassEncoder);
+    copyPassEncoder.setBindGroup(1, this.compactionBindGroup);
+    copyPassEncoder.dispatchWorkgroups(
+      ...getWorkgroupCounts(
+        this.device,
+        agentCount,
+        AgentGenerationPipeline.WORKGROUP_SIZE
+      )
     );
+    copyPassEncoder.end();
+
     commandEncoder.copyBufferToBuffer(
       this.countersBuffer,
       0,

src/pipelines/agents/agent-generation/agent-schema.test.ts

@@ -69,6 +69,27 @@ describe('Agent TS/WGSL contract', () => {
     expect(agentSchema).toContain('workgroup_count.x * 64');
     expect(agentSchema).toContain('workgroup_count.x * workgroup_count.y * 64');
     expect(compactionShader).toContain('let id = get_id(global_id, workgroup_count);');
-    expect(compactionShader).toContain('if id >= settings.agentCount');
+    expect(compactionShader).toContain('if id < settings.agentCount');
+  });
+
+  it('keeps compaction copy-back bounded by the compacted count', () => {
+    expect(compactionShader).toContain('fn copyCompactedAgents');
+    expect(compactionShader).toContain(
+      'workgroupCopyCount = atomicLoad(&counters.aliveAgentCount);'
+    );
+    expect(compactionShader).toContain('if id >= workgroupCopyCount');
+    expect(compactionShader).toContain('agents[id] = compactedAgents[id];');
+  });
+
+  it('uses workgroup-local counting before allocating global compacted ranges', () => {
+    expect(compactionShader).toContain(
+      'var<workgroup> workgroupAliveCount: atomic<u32>;'
+    );
+    expect(compactionShader).toContain(
+      'localCompactedIndex = atomicAdd(&workgroupAliveCount, 1u);'
+    );
+    expect(
+      compactionShader.match(/atomicAdd\(&counters\.aliveAgentCount/g) ?? []
+    ).toHaveLength(1);
   });
 });

src/pipelines/agents/agent.wgsl

@@ -35,7 +35,7 @@ fn main(
   }
 
   var agent = agents[id];
-  if agent.colorIndex < 0.0 {
+  if agent.colorIndex < 0.0 || agent.colorIndex >= 2.5 {
     return;
   }
 
@@ -48,12 +48,7 @@ fn main(
     return;
   }
 
-  let random = textureSampleLevel(
-    noise,
-    noiseSampler,
-    fract(vec2(f32(id) * 0.7548777, state.time * 0.00017 + f32(id) * 0.5698403)),
-    0
-  );
+  let random = random_vec4(id, state.time);
 
   let forwardSensor = sensor_position(agent.position, agent.angle, settings.sensorOffset, 0);
   let leftSensor = sensor_position(agent.position, agent.angle, settings.sensorOffset, settings.sensorAngle);
@@ -154,7 +149,10 @@ fn get_channel_mask(colorIndex: f32) -> vec3<f32> {
   if colorIndex < 1.5 {
     return vec3<f32>(0, 1, 0);
   }
-  return vec3<f32>(0, 0, 1);
+  if colorIndex < 2.5 {
+    return vec3<f32>(0, 0, 1);
+  }
+  return vec3<f32>(0.0, 0.0, 0.0);
 }
 
 fn get_reaction_mask(colorIndex: f32) -> vec3<f32> {
@@ -172,13 +170,37 @@ fn get_reaction_mask(colorIndex: f32) -> vec3<f32> {
       settings.color2ToColor3
     );
   }
-  return vec3<f32>(
-    settings.color3ToColor1,
-    settings.color3ToColor2,
-    settings.color3ToColor3
-  );
+  if colorIndex < 2.5 {
+    return vec3<f32>(
+      settings.color3ToColor1,
+      settings.color3ToColor2,
+      settings.color3ToColor3
+    );
+  }
+  return vec3<f32>(0.0, 0.0, 0.0);
 }
 
 fn angle_delta(sourceAngle: f32, targetAngle: f32) -> f32 {
   return atan2(sin(targetAngle - sourceAngle), cos(targetAngle - sourceAngle));
 }
+
+fn random_vec4(id: u32, time: f32) -> vec4<f32> {
+  let timeSeed = u32(time * 0.34816);
+  let seed = id * 747796405u + timeSeed * 2891336453u;
+  return vec4<f32>(
+    random_float(seed),
+    random_float(seed + 1013904223u),
+    random_float(seed + 1664525u),
+    random_float(seed + 22695477u)
+  );
+}
+
+fn random_float(seed: u32) -> f32 {
+  return f32(hash_u32(seed) >> 8u) * (1.0 / 16777216.0);
+}
+
+fn hash_u32(seed: u32) -> u32 {
+  let value = seed * 747796405u + 2891336453u;
+  let word = ((value >> ((value >> 28u) + 4u)) ^ value) * 277803737u;
+  return (word >> 22u) ^ word;
+}