Optimise

src/pipelines/agents/agent-dispatch.ts

@@ -1,8 +1,21 @@
-export const AGENT_WORKGROUP_SIZE = 64;
+// Use the device's max workgroup size so we get full SIMD/wave occupancy on
+// hardware that supports more than the WebGPU minimum of 256.
+export const getAgentWorkgroupSize = (device: GPUDevice): number =>
+  device.limits.maxComputeInvocationsPerWorkgroup;
+
+export const substituteAgentWorkgroupSize = (
+  device: GPUDevice,
+  shaderCode: string
+): string =>
+  shaderCode.replaceAll(
+    '__AGENT_WORKGROUP_SIZE__',
+    String(getAgentWorkgroupSize(device))
+  );
 
 export const dispatchAgentWorkgroups = (
   passEncoder: GPUComputePassEncoder,
+  workgroupSize: number,
   agentCount: number
 ): void => {
-  passEncoder.dispatchWorkgroups(Math.ceil(agentCount / AGENT_WORKGROUP_SIZE), 1);
+  passEncoder.dispatchWorkgroups(Math.ceil(agentCount / workgroupSize), 1);
 };

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

@@ -15,40 +15,53 @@ const clearCompactedTailStride = 4u;
 @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> scanData: array<u32, agentWorkgroupSize>;
 var<workgroup> clearAliveAgentCount: u32;
 
-@compute @workgroup_size(64)
+@compute @workgroup_size(agentWorkgroupSize)
 fn main(
   @builtin(global_invocation_id) global_id: vec3<u32>,
   @builtin(local_invocation_id) local_id: vec3<u32>
 ) {
   let id = get_id(global_id);
+  let lid = local_id.x;
 
-  if local_id.x == 0u {
-    atomicStore(&workgroupAliveCount, 0u);
-  }
-
-  workgroupBarrier();
-
-  var localCompactedIndex = 0u;
-  var agent: Agent;
   var isAlive = false;
+  var agent: Agent;
   if id < settings.agentCount {
     isAlive = agents[id].colorIndex >= 0.0;
     if isAlive {
       agent = agents[id];
-      localCompactedIndex = atomicAdd(&workgroupAliveCount, 1u);
     }
   }
 
+  // Hillis-Steele inclusive prefix sum across the workgroup. Replaces a
+  // per-thread atomicAdd to a workgroup counter, eliminating serialization
+  // on dense workgroups.
+  scanData[lid] = select(0u, 1u, isAlive);
   workgroupBarrier();
 
-  if local_id.x == 0u {
-    let groupAliveCount = atomicLoad(&workgroupAliveCount);
-    if groupAliveCount > 0u {
-      workgroupCompactedOffset = atomicAdd(&counters.aliveAgentCount, groupAliveCount);
+  var offset: u32 = 1u;
+  while offset < agentWorkgroupSize {
+    let own = scanData[lid];
+    var contribution: u32 = 0u;
+    if lid >= offset {
+      contribution = scanData[lid - offset];
+    }
+    workgroupBarrier();
+    scanData[lid] = own + contribution;
+    workgroupBarrier();
+    offset = offset * 2u;
+  }
+
+  let inclusivePrefix = scanData[lid];
+  let workgroupAliveTotal = scanData[agentWorkgroupSize - 1u];
+  let exclusivePrefix = inclusivePrefix - select(0u, 1u, isAlive);
+
+  if lid == 0u {
+    if workgroupAliveTotal > 0u {
+      workgroupCompactedOffset = atomicAdd(&counters.aliveAgentCount, workgroupAliveTotal);
     } else {
       workgroupCompactedOffset = 0u;
     }
@@ -57,11 +70,11 @@ fn main(
   workgroupBarrier();
 
   if isAlive {
-    compactedAgents[workgroupCompactedOffset + localCompactedIndex] = agent;
+    compactedAgents[workgroupCompactedOffset + exclusivePrefix] = agent;
   }
 }
 
-@compute @workgroup_size(64)
+@compute @workgroup_size(agentWorkgroupSize)
 fn clearCompactedTail(
   @builtin(global_invocation_id) global_id: vec3<u32>,
   @builtin(local_invocation_id) local_id: vec3<u32>

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

@@ -2,14 +2,16 @@ import { vec2 } from 'gl-matrix';
 
 import { createBindGroupCache } from '../../../utils/graphics/bind-group-cache';
 import { smartCompile } from '../../../utils/graphics/smart-compile';
-import { dispatchAgentWorkgroups } from '../agent-dispatch';
+import {
+  dispatchAgentWorkgroups,
+  getAgentWorkgroupSize,
+  substituteAgentWorkgroupSize,
+} from '../agent-dispatch';
 import { AGENT_SIZE_IN_BYTES, getMaxSupportedAgentCount } from '../agent-limits';
 import compactionShader from './agent-compaction.wgsl?raw';
 import resizeShader from './agent-resize.wgsl?raw';
 import agentSchema from './agent-schema.wgsl?raw';
 
-export { AGENT_FLOAT_COUNT } from '../agent-limits';
-
 export class AgentGenerationPipeline {
   private static readonly UNIFORM_COUNT = 4;
   private static readonly COUNTER_COUNT = 1;
@@ -34,6 +36,7 @@ export class AgentGenerationPipeline {
   private readonly resizePipeline: GPUComputePipeline;
   private readonly compactionPipeline: GPUComputePipeline;
   private readonly clearCompactedTailPipeline: GPUComputePipeline;
+  private readonly workgroupSize: number;
 
   private activeAgentsBuffer: GPUBuffer;
   private inactiveAgentsBuffer: GPUBuffer;
@@ -90,7 +93,7 @@ export class AgentGenerationPipeline {
     });
 
     this.activeAgentsBuffer = this.createAgentsBuffer();
-    this.inactiveAgentsBuffer = this.createInactivePlaceholderBuffer();
+    this.inactiveAgentsBuffer = this.createAgentsBuffer();
 
     this.countersBuffer = this.device.createBuffer({
       size: AgentGenerationPipeline.COUNTER_COUNT * Uint32Array.BYTES_PER_ELEMENT,
@@ -107,17 +110,20 @@ export class AgentGenerationPipeline {
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
     });
 
+    this.workgroupSize = getAgentWorkgroupSize(device);
+    const sizedSchema = substituteAgentWorkgroupSize(device, agentSchema);
+
     this.resizePipeline = device.createComputePipeline({
       layout: device.createPipelineLayout({
         bindGroupLayouts: [emptyBindGroupLayout, this.bindGroupLayout],
       }),
       compute: {
-        module: smartCompile(device, agentSchema, resizeShader),
+        module: smartCompile(device, sizedSchema, resizeShader),
         entryPoint: 'main',
       },
     });
 
-    const compactionModule = smartCompile(device, agentSchema, compactionShader);
+    const compactionModule = smartCompile(device, sizedSchema, compactionShader);
 
     this.compactionPipeline = device.createComputePipeline({
       layout: device.createPipelineLayout({
@@ -151,16 +157,6 @@ export class AgentGenerationPipeline {
     });
   }
 
-  // The inactive slot only needs a real allocation during compaction. The rest of
-  // the time we keep a one-agent placeholder so the bind group at binding 3 stays
-  // valid for resize without holding a second N-agent buffer in GPU memory.
-  private createInactivePlaceholderBuffer(): GPUBuffer {
-    return this.device.createBuffer({
-      size: AGENT_SIZE_IN_BYTES,
-      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
-    });
-  }
-
   public get maxAgentCount(): number {
     return this.allocatedMaxAgentCount;
   }
@@ -187,9 +183,11 @@ export class AgentGenerationPipeline {
       )
     );
     const previousActiveAgentsBuffer = this.activeAgentsBuffer;
+    const previousInactiveAgentsBuffer = this.inactiveAgentsBuffer;
     const previousMaxAgentCount = this.allocatedMaxAgentCount;
     this.allocatedMaxAgentCount = nextMaxAgentCount;
     this.activeAgentsBuffer = this.createAgentsBuffer();
+    this.inactiveAgentsBuffer = this.createAgentsBuffer();
 
     const copyAgentCount = Math.min(
       Math.max(0, Math.floor(activeAgentCount)),
@@ -209,10 +207,9 @@ export class AgentGenerationPipeline {
     }
 
     // GPUBuffer.destroy() defers actual freeing until pending submissions
-    // finish, so calling it synchronously after submit is safe and avoids the
-    // transient 4-buffers-live spike that pushes iOS Safari past its per-tab
-    // memory ceiling.
+    // finish, so calling it synchronously after submit is safe.
     previousActiveAgentsBuffer.destroy();
+    previousInactiveAgentsBuffer.destroy();
     return this.allocatedMaxAgentCount;
   }
 
@@ -251,7 +248,7 @@ export class AgentGenerationPipeline {
     const passEncoder = commandEncoder.beginComputePass();
     passEncoder.setPipeline(this.resizePipeline);
     passEncoder.setBindGroup(1, this.getBindGroup());
-    dispatchAgentWorkgroups(passEncoder, agentCount);
+    dispatchAgentWorkgroups(passEncoder, this.workgroupSize, agentCount);
     passEncoder.end();
 
     this.device.queue.submit([commandEncoder.finish()]);
@@ -262,12 +259,6 @@ export class AgentGenerationPipeline {
       return 0;
     }
 
-    // Stash the placeholder, swap in a real N-agent destination buffer just
-    // for this compaction so the rest of the time we only carry one full
-    // agent buffer in memory.
-    const placeholder = this.inactiveAgentsBuffer;
-    this.inactiveAgentsBuffer = this.createAgentsBuffer();
-
     this.agentCountUniformValues[0] = agentCount;
     this.device.queue.writeBuffer(this.uniforms, 0, this.agentCountUniformValues);
 
@@ -276,10 +267,11 @@ export class AgentGenerationPipeline {
     const passEncoder = commandEncoder.beginComputePass();
     passEncoder.setPipeline(this.compactionPipeline);
     passEncoder.setBindGroup(1, this.getBindGroup());
-    dispatchAgentWorkgroups(passEncoder, agentCount);
+    dispatchAgentWorkgroups(passEncoder, this.workgroupSize, agentCount);
     passEncoder.setPipeline(this.clearCompactedTailPipeline);
     dispatchAgentWorkgroups(
       passEncoder,
+      this.workgroupSize,
       Math.ceil(agentCount / AgentGenerationPipeline.CLEAR_COMPACTED_TAIL_STRIDE)
     );
     passEncoder.end();
@@ -295,13 +287,6 @@ export class AgentGenerationPipeline {
     this.device.queue.submit([commandEncoder.finish()]);
     this.swapAgentBuffers();
 
-    // After swap, inactive is the previous active (full size). Destroy it and
-    // restore the placeholder; the destroy is deferred by WebGPU until the
-    // submitted compaction work has finished.
-    const previousActiveAgentsBuffer = this.inactiveAgentsBuffer;
-    this.inactiveAgentsBuffer = placeholder;
-    previousActiveAgentsBuffer.destroy();
-
     await this.countersStagingBuffer.mapAsync(GPUMapMode.READ);
     const compactedCount = new Uint32Array(
       this.countersStagingBuffer.getMappedRange(),

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

@@ -5,7 +5,7 @@ struct ResizeSettings {
 
 @group(1) @binding(0) var<uniform> resizeSettings: ResizeSettings;
 
-@compute @workgroup_size(64)
+@compute @workgroup_size(agentWorkgroupSize)
 fn main(
   @builtin(global_invocation_id) global_id: vec3<u32>
 ) {

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

@@ -9,7 +9,7 @@ struct Agent {
 
 @group(1) @binding(1) var<storage, read_write> agents: array<Agent>;
 
-const agentWorkgroupSize = 64u;
+const agentWorkgroupSize = __AGENT_WORKGROUP_SIZE__u;
 
 fn get_id(global_id: vec3<u32>) -> u32 {
   return global_id.x;

src/pipelines/agents/agent-limits.ts

@@ -1,8 +1,46 @@
-import { AGENT_WORKGROUP_SIZE } from './agent-dispatch';
+import { getAgentWorkgroupSize } from './agent-dispatch';
 
 export const AGENT_FLOAT_COUNT = 8;
 export const AGENT_SIZE_IN_BYTES = AGENT_FLOAT_COUNT * Float32Array.BYTES_PER_ELEMENT;
 
+const AGENT_LAYOUT = {
+  positionX: 0,
+  positionY: 1,
+  angle: 2,
+  colorIndex: 3,
+  targetPositionX: 4,
+  targetPositionY: 5,
+  targetAngle: 6,
+  introDelay: 7,
+} as const;
+
+export interface AgentLayoutValues {
+  angle: number;
+  colorIndex: number;
+  introDelay: number;
+  positionX: number;
+  positionY: number;
+  targetAngle: number;
+  targetPositionX: number;
+  targetPositionY: number;
+}
+
+export const writeAgentValues = (
+  target: Float32Array,
+  agentIndex: number,
+  values: AgentLayoutValues
+): void => {
+  const base = agentIndex * AGENT_FLOAT_COUNT;
+  target[base + AGENT_LAYOUT.positionX] = values.positionX;
+  target[base + AGENT_LAYOUT.positionY] = values.positionY;
+  target[base + AGENT_LAYOUT.angle] = values.angle;
+  target[base + AGENT_LAYOUT.colorIndex] = values.colorIndex;
+  target[base + AGENT_LAYOUT.targetPositionX] = values.targetPositionX;
+  target[base + AGENT_LAYOUT.targetPositionY] = values.targetPositionY;
+  target[base + AGENT_LAYOUT.targetAngle] = values.targetAngle;
+  target[base + AGENT_LAYOUT.introDelay] = values.introDelay;
+};
+
 export const getMaxSupportedAgentCount = (
   device: GPUDevice,
   maxAgentCountUpperLimit = Number.POSITIVE_INFINITY
@@ -19,7 +57,8 @@ export const getMaxSupportedAgentCount = (
       upperLimit,
       Math.floor(device.limits.maxBufferSize / AGENT_SIZE_IN_BYTES),
       Math.floor(storageBufferBindingSize / AGENT_SIZE_IN_BYTES),
-      Math.floor(device.limits.maxComputeWorkgroupsPerDimension) * AGENT_WORKGROUP_SIZE
+      Math.floor(device.limits.maxComputeWorkgroupsPerDimension) *
+        getAgentWorkgroupSize(device)
     )
   );
 };

src/pipelines/agents/agent-pipeline.ts

@@ -1,11 +1,16 @@
 import { createBindGroupCache3 } from '../../utils/graphics/bind-group-cache';
 import {
-  createCachedFloat32BufferWrite,
-  writeFloat32BufferIfChanged,
+  createCachedBufferWrite,
+  writeBufferIfChanged,
 } from '../../utils/graphics/cached-buffer-write';
 import { smartCompile } from '../../utils/graphics/smart-compile';
 import { CommonState } from '../common-state/common-state';
-import { dispatchAgentWorkgroups } from './agent-dispatch';
+import { TRAIL_SOURCE_TEXTURE_FORMAT } from '../texture-formats';
+import {
+  dispatchAgentWorkgroups,
+  getAgentWorkgroupSize,
+  substituteAgentWorkgroupSize,
+} from './agent-dispatch';
 import agentSchema from './agent-generation/agent-schema.wgsl?raw';
 import shader from './agent.wgsl?raw';
 
@@ -44,11 +49,13 @@ const UNIFORM_COUNT = 30;
 export class AgentPipeline {
   private readonly bindGroupLayout: GPUBindGroupLayout;
   private readonly pipeline: GPUComputePipeline;
-  private readonly normalPipeline: GPUComputePipeline;
   private readonly uniforms: GPUBuffer;
+  private readonly workgroupSize: number;
   private readonly uniformValues = new Float32Array(UNIFORM_COUNT);
   private readonly uniformUintValues = new Uint32Array(this.uniformValues.buffer);
-  private readonly uniformCache = createCachedFloat32BufferWrite(UNIFORM_COUNT);
+  private readonly uniformCache = createCachedBufferWrite(
+    UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT
+  );
   private readonly bindGroupCache = createBindGroupCache3<
     GPUBuffer,
     GPUTextureView,
@@ -66,7 +73,6 @@ export class AgentPipeline {
   );
 
   private agentCount = 0;
-  private useIntroPipeline = true;
 
   public constructor(
     private readonly device: GPUDevice,
@@ -93,15 +99,16 @@ export class AgentPipeline {
         {
           binding: 3,
           visibility: GPUShaderStage.COMPUTE,
-          storageTexture: { format: 'rgba16float' },
+          storageTexture: { format: TRAIL_SOURCE_TEXTURE_FORMAT },
         },
       ],
     });
 
+    this.workgroupSize = getAgentWorkgroupSize(device);
     const shaderModule = smartCompile(
       device,
       CommonState.shaderCode,
-      agentSchema,
+      substituteAgentWorkgroupSize(device, agentSchema),
       shader
     );
     const pipelineLayout = device.createPipelineLayout({
@@ -114,13 +121,6 @@ export class AgentPipeline {
         entryPoint: 'main',
       },
     });
-    this.normalPipeline = device.createComputePipeline({
-      layout: pipelineLayout,
-      compute: {
-        module: shaderModule,
-        entryPoint: 'mainNormal',
-      },
-    });
 
     this.uniforms = device.createBuffer({
       size: UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT,
@@ -167,7 +167,6 @@ export class AgentPipeline {
     introProgress?: number;
   }) {
     this.agentCount = agentCount;
-    this.useIntroPipeline = (introProgress ?? 1) < introProgressCutoff;
     this.uniformValues[0] = moveSpeed * deltaTime;
     this.uniformValues[1] = turnSpeed * deltaTime;
     const sensorAngle = (sensorOffsetAngle * Math.PI) / 180;
@@ -199,7 +198,7 @@ export class AgentPipeline {
     this.uniformValues[27] = introNearMoveMultiplier;
     this.uniformValues[28] = introStepStopDistance;
     this.uniformUintValues[29] = Math.max(0, Math.floor(time * randomTimeScale)) >>> 0;
-    writeFloat32BufferIfChanged(
+    writeBufferIfChanged(
       this.device,
       this.uniforms,
       this.uniformValues,
@@ -220,13 +219,13 @@ export class AgentPipeline {
     const passEncoder = commandEncoder.beginComputePass(
       timestampWrites ? { timestampWrites } : undefined
     );
-    passEncoder.setPipeline(this.useIntroPipeline ? this.pipeline : this.normalPipeline);
+    passEncoder.setPipeline(this.pipeline);
     this.commonState.execute(passEncoder);
     passEncoder.setBindGroup(
       1,
       this.bindGroupCache(this.getAgentsBuffer(), trailMapIn, trailMapOut)
     );
-    dispatchAgentWorkgroups(passEncoder, this.agentCount);
+    dispatchAgentWorkgroups(passEncoder, this.workgroupSize, this.agentCount);
     passEncoder.end();
   }
 

src/pipelines/agents/agent.wgsl

@@ -1,4 +1,6 @@
 const PI: f32 = 3.14159265359;
+const TAU: f32 = 6.28318530718;
+const INV_TAU: f32 = 0.15915494309;
 
 struct Settings {
   moveRate: f32,
@@ -35,9 +37,9 @@ struct Settings {
 
 @group(1) @binding(0) var<uniform> settings: Settings;
 @group(1) @binding(2) var trailMapIn: texture_2d<f32>;
-@group(1) @binding(3) var trailMapOut: texture_storage_2d<rgba16float, write>;
+@group(1) @binding(3) var trailMapOut: texture_storage_2d<rgba8unorm, write>;
 
-@compute @workgroup_size(64)
+@compute @workgroup_size(agentWorkgroupSize)
 fn main(
   @builtin(global_invocation_id) global_id: vec3<u32>
 ) {
@@ -158,79 +160,6 @@ fn main(
   agents[id].position = nextPosition;
 }
 
-@compute @workgroup_size(64)
-fn mainNormal(
-  @builtin(global_invocation_id) global_id: vec3<u32>
-) {
-  let id = get_id(global_id);
-
-  if id >= settings.agentCount {
-    return;
-  }
-
-  let colorIndex = agents[id].colorIndex;
-  if colorIndex < 0.0 || colorIndex >= 2.5 {
-    return;
-  }
-
-  var position = agents[id].position;
-  var angle = agents[id].angle;
-  let channelMask = get_channel_mask(colorIndex);
-  let reactionMask = get_reaction_mask(colorIndex);
-  let randomSeed = random_seed(id);
-  let maxPosition = state.size - vec2<f32>(1.0, 1.0);
-  let randomTurn = random_float(randomSeed);
-  let direction = vec2(cos(angle), sin(angle));
-
-  let forwardSensor = sensor_position(position, direction, settings.sensorOffset, maxPosition);
-  let leftSensor = sensor_position(
-    position,
-    rotate_direction(direction, settings.sensorAngleSin, settings.sensorAngleCos),
-    settings.sensorOffset,
-    maxPosition
-  );
-  let rightSensor = sensor_position(
-    position,
-    rotate_direction(direction, -settings.sensorAngleSin, settings.sensorAngleCos),
-    settings.sensorOffset,
-    maxPosition
-  );
-
-  let trailForward = textureLoad(trailMapIn, forwardSensor, 0);
-  let trailLeft = textureLoad(trailMapIn, leftSensor, 0);
-  let trailRight = textureLoad(trailMapIn, rightSensor, 0);
-
-  let weightForward = dot(trailForward.rgb, reactionMask);
-  let weightLeft = dot(trailLeft.rgb, reactionMask);
-  let weightRight = dot(trailRight.rgb, reactionMask);
-
-  var rotation = (randomTurn - 0.5) * settings.turnWhenLost;
-  if weightForward >= weightLeft && weightForward >= weightRight {
-    rotation = rotation * settings.forwardRotationScale;
-  } else {
-    rotation += sign(weightLeft - weightRight) * settings.turnRate;
-  }
-
-  let nextPosition = clamp(
-    position + direction * settings.moveRate,
-    vec2<f32>(0, 0),
-    maxPosition
-  );
-  if nextPosition.x == 0 || nextPosition.x == maxPosition.x || nextPosition.y == 0 || nextPosition.y == maxPosition.y {
-    rotation = PI + random_float(randomSeed + 22695477u) - 0.5;
-  }
-
-  var trailBelow = textureLoad(trailMapIn, vec2<i32>(nextPosition), 0);
-  trailBelow = vec4<f32>(
-    trailBelow.rgb + channelMask * settings.individualTrailWeight,
-    max(trailBelow.a, 0.0)
-  );
-
-  textureStore(trailMapOut, vec2<i32>(nextPosition), trailBelow);
-  agents[id].angle = angle + rotation;
-  agents[id].position = nextPosition;
-}
-
 fn sensor_position(
   agentPosition: vec2<f32>,
   direction: vec2<f32>,
@@ -290,7 +219,8 @@ fn get_reaction_mask(colorIndex: f32) -> vec3<f32> {
 }
 
 fn angle_delta(sourceAngle: f32, targetAngle: f32) -> f32 {
-  return atan2(sin(targetAngle - sourceAngle), cos(targetAngle - sourceAngle));
+  // Wraps to (-π, π] via fract(); replaces atan2(sin(d), cos(d)).
+  return (fract((targetAngle - sourceAngle) * INV_TAU + 0.5) - 0.5) * TAU;
 }
 
 fn random_seed(id: u32) -> u32 {