Allow more agents & dynamic setting of agent count

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

@@ -1,15 +1,25 @@
+struct Settings {
+  agentCount: u32 // might be smaller than the length of the agents array
+};
+
+@group(1) @binding(0) var<uniform> settings: Settings;
+
 struct Counters {
-  evenGenerationAlive: atomic<i32>,
-  oddGenerationAlive: atomic<i32>,
+  evenGenerationAlive: atomic<u32>,
+  oddGenerationAlive: atomic<u32>,
 };
 
 @group(1) @binding(2) var<storage, read_write> counters: Counters;
 
-@compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  let id = global_id.x;
 
-  if id >= arrayLength(&agents) {
+@compute @workgroup_size(64)
+fn main(
+  @builtin(global_invocation_id) global_id: vec3<u32>,
+  @builtin(num_workgroups) workgroup_count: vec3<u32>
+) {
+  let id = global_id.x + global_id.y * (workgroup_count.x * 64) + global_id.z * (workgroup_count.x * workgroup_count.y * 64);
+
+  if id >= settings.agentCount {
     return;
   }
 
@@ -18,4 +28,7 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
   } else {
     atomicAdd(&counters.oddGenerationAlive, 1);
   }
+
+  // atomicStore(&counters.evenGenerationAlive, settings.agentCount);
+  // atomicStore(&counters.oddGenerationAlive,  workgroup_count.y);
 }

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

@@ -1,6 +1,9 @@
 @compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  let id = global_id.x;
+fn main(
+  @builtin(global_invocation_id) global_id: vec3<u32>,
+  @builtin(num_workgroups) workgroup_count: vec3<u32>
+) {
+  let id = global_id.x + global_id.y * (workgroup_count.x * 64) + global_id.z * (workgroup_count.x * workgroup_count.y * 64);
 
   if id >= arrayLength(&agents) {
     return;

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

@@ -1,3 +1,4 @@
+import { getWorkgroupCounts } from '../../../utils/graphics/get-workgroup-counts';
 import random from '../../../utils/graphics/random.wgsl';
 import { smartCompile } from '../../../utils/graphics/smart-compile';
 import { CommonState } from '../../common-state/common-state';
@@ -9,7 +10,7 @@ import { GenerationCounts } from './generation-counts';
 
 export class AgentGenerationPipeline {
   private static readonly WORKGROUP_SIZE = 64;
-  private static readonly UNIFORM_COUNT = 4;
+  private static readonly UNIFORM_COUNT = 1;
   private static readonly COUNTER_COUNT = 3;
 
   private readonly bindGroupLayout: GPUBindGroupLayout;
@@ -60,17 +61,17 @@ export class AgentGenerationPipeline {
     });
 
     this.countersBuffer = this.device.createBuffer({
-      size: AgentGenerationPipeline.COUNTER_COUNT * Int32Array.BYTES_PER_ELEMENT,
+      size: AgentGenerationPipeline.COUNTER_COUNT * Uint32Array.BYTES_PER_ELEMENT,
       usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
     });
 
     this.countersStagingBuffer = this.device.createBuffer({
-      size: AgentGenerationPipeline.COUNTER_COUNT * Int32Array.BYTES_PER_ELEMENT,
+      size: AgentGenerationPipeline.COUNTER_COUNT * Uint32Array.BYTES_PER_ELEMENT,
       usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
     });
 
     this.uniforms = this.device.createBuffer({
-      size: AgentGenerationPipeline.UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT,
+      size: AgentGenerationPipeline.UNIFORM_COUNT * Uint32Array.BYTES_PER_ELEMENT,
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
     });
 
@@ -134,7 +135,8 @@ export class AgentGenerationPipeline {
   public get maxAgentCount(): number {
     return Math.min(
       this.maxAgentCountUpperLimit,
-      Math.floor(this.device.limits.maxBufferSize / AGENT_SIZE_IN_BYTES)
+      Math.floor(this.device.limits.maxBufferSize / AGENT_SIZE_IN_BYTES),
+      this.device.limits.maxComputeWorkgroupsPerDimension ** 3
     );
   }
 
@@ -146,15 +148,20 @@ export class AgentGenerationPipeline {
     passEncoder.setPipeline(this.firstGenerationPipeline);
     passEncoder.setBindGroup(1, this.bindGroup);
     passEncoder.dispatchWorkgroups(
-      Math.ceil(this.maxAgentCount / AgentGenerationPipeline.WORKGROUP_SIZE)
+      ...getWorkgroupCounts(
+        this.device,
+        this.maxAgentCount,
+        AgentGenerationPipeline.WORKGROUP_SIZE
+      )
     );
     passEncoder.end();
 
     this.device.queue.submit([commandEncoder.finish()]);
   }
 
-  public async countAgents(): Promise<GenerationCounts> {
-    this.device.queue.writeBuffer(this.countersBuffer, 0, new Int32Array([0, 0]));
+  public async countAgents(agentCount: number): Promise<GenerationCounts> {
+    this.device.queue.writeBuffer(this.countersBuffer, 0, new Uint32Array([0, 0]));
+    this.device.queue.writeBuffer(this.uniforms, 0, new Uint32Array([agentCount]));
 
     const commandEncoder = this.device.createCommandEncoder();
 
@@ -163,7 +170,11 @@ export class AgentGenerationPipeline {
     this.commonState.execute(passEncoder);
     passEncoder.setBindGroup(1, this.bindGroup);
     passEncoder.dispatchWorkgroups(
-      Math.ceil(this.maxAgentCount / AgentGenerationPipeline.WORKGROUP_SIZE)
+      ...getWorkgroupCounts(
+        this.device,
+        agentCount,
+        AgentGenerationPipeline.WORKGROUP_SIZE
+      )
     );
     passEncoder.end();
 
@@ -172,14 +183,14 @@ export class AgentGenerationPipeline {
       0,
       this.countersStagingBuffer,
       0,
-      AgentGenerationPipeline.COUNTER_COUNT * Int32Array.BYTES_PER_ELEMENT
+      AgentGenerationPipeline.COUNTER_COUNT * Uint32Array.BYTES_PER_ELEMENT
     );
 
     this.device.queue.submit([commandEncoder.finish()]);
 
     await this.countersStagingBuffer.mapAsync(GPUMapMode.READ);
 
-    const data = new Int32Array(this.countersStagingBuffer.getMappedRange().slice(0));
+    const data = new Uint32Array(this.countersStagingBuffer.getMappedRange().slice(0));
     this.countersStagingBuffer.unmap();
     return {
       evenGenerationCount: data[0],

src/pipelines/agents/agent-pipeline.ts

@@ -1,7 +1,7 @@
+import { getWorkgroupCounts } from '../../utils/graphics/get-workgroup-counts';
 import random from '../../utils/graphics/random.wgsl';
 import { smartCompile } from '../../utils/graphics/smart-compile';
 import { CommonState } from '../common-state/common-state';
-import { AGENT_SIZE_IN_BYTES } from './agent-generation/agent';
 import agentSchme from './agent-generation/agent-schema.wgsl';
 import { AgentSettings } from './agent-settings';
 import shader from './agent.wgsl';
@@ -10,7 +10,7 @@ import { vec2 } from 'gl-matrix';
 
 export class AgentPipeline {
   private static readonly WORKGROUP_SIZE = 64;
-  private static readonly UNIFORM_COUNT = 16;
+  private static readonly UNIFORM_COUNT = 17;
 
   private readonly bindGroupLayout: GPUBindGroupLayout;
   private readonly pipeline: GPUComputePipeline;
@@ -19,6 +19,8 @@ export class AgentPipeline {
   private previousTrailMapIn?: GPUTextureView;
   private previousTrailMapOut?: GPUTextureView;
 
+  private agentCount = 0;
+
   public constructor(
     private readonly device: GPUDevice,
     private readonly commonState: CommonState,
@@ -57,13 +59,16 @@ export class AgentPipeline {
     turnWhenLost,
     individualTrailWeight,
     deinfectionProbability,
+    agentCount,
   }: AgentSettings & {
     evenGenerationAggression: number;
     oddGenerationAggression: number;
     nextGenerationId: number;
     center: vec2;
     radius: number;
+    agentCount: number;
   }) {
+    this.agentCount = agentCount;
     this.device.queue.writeBuffer(
       this.uniforms,
       0,
@@ -82,6 +87,7 @@ export class AgentPipeline {
         turnWhenLost,
         individualTrailWeight,
         deinfectionProbability,
+        agentCount,
       ])
     );
   }
@@ -98,9 +104,7 @@ export class AgentPipeline {
     this.commonState.execute(passEncoder);
     passEncoder.setBindGroup(1, this.bindGroup);
     passEncoder.dispatchWorkgroups(
-      Math.ceil(
-        this.agentsBuffer.size / AGENT_SIZE_IN_BYTES / AgentPipeline.WORKGROUP_SIZE
-      )
+      ...getWorkgroupCounts(this.device, this.agentCount, AgentPipeline.WORKGROUP_SIZE)
     );
     passEncoder.end();
   }

src/pipelines/agents/agent.wgsl

@@ -17,7 +17,9 @@ struct Settings {
   turnWhenGoingInTheRightDirection: f32,
   turnWhenLost: f32,
   individualTrailWeight: f32,
-  deinfectionProbability: f32
+  deinfectionProbability: f32,
+
+  agentCount: f32 // might be smaller than the length of the agents array
 };
 
 
@@ -31,10 +33,13 @@ struct Settings {
 @group(1) @binding(3) var trailMapOut: texture_storage_2d<rgba16float, write>;
 
 @compute @workgroup_size(64)
-fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
-  let id = global_id.x;
+fn main(
+  @builtin(global_invocation_id) global_id: vec3<u32>,
+  @builtin(num_workgroups) workgroup_count: vec3<u32>
+) {
+  let id = global_id.x + global_id.y * (workgroup_count.x * 64) + global_id.z * (workgroup_count.x * workgroup_count.y * 64);
 
-  if (id >= arrayLength(&agents)) {
+  if id >= u32(settings.agentCount) {
     return;
   }