WIP

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

@@ -1,27 +1,42 @@
+import { vec2 } from 'gl-matrix';
+
 import { getWorkgroupCounts } from '../../../utils/graphics/get-workgroup-counts';
 import { smartCompile } from '../../../utils/graphics/smart-compile';
 import { CommonState } from '../../common-state/common-state';
 import { AGENT_SIZE_IN_BYTES } from './agent';
+import compactionShader from './agent-compaction.wgsl?raw';
 import countingShader from './agent-counting.wgsl?raw';
 import firstGenerationShader from './agent-first-generation.wgsl?raw';
+import resizeShader from './agent-resize.wgsl?raw';
 import agentSchema from './agent-schema.wgsl?raw';
 import { GenerationCounts } from './generation-counts';
 
 export class AgentGenerationPipeline {
   private static readonly WORKGROUP_SIZE = 64;
-  private static readonly UNIFORM_COUNT = 1;
+  private static readonly UNIFORM_COUNT = 4;
   private static readonly COUNTER_COUNT = 3;
 
   private readonly bindGroupLayout: GPUBindGroupLayout;
+  private readonly compactionBindGroupLayout: GPUBindGroupLayout;
   private readonly uniforms: GPUBuffer;
   private readonly bindGroup: GPUBindGroup;
+  private readonly compactionBindGroup: GPUBindGroup;
 
   private readonly firstGenerationPipeline: GPUComputePipeline;
   private readonly countingPipeline: GPUComputePipeline;
+  private readonly resizePipeline: GPUComputePipeline;
+  private readonly compactionPipeline: GPUComputePipeline;
 
   public readonly agentsBuffer: GPUBuffer;
+  private readonly compactedAgentsBuffer: GPUBuffer;
   public readonly countersBuffer: GPUBuffer;
   public readonly countersStagingBuffer: GPUBuffer;
+  private readonly counterClearValues = new Uint32Array(
+    AgentGenerationPipeline.COUNTER_COUNT
+  );
+  private readonly agentCountUniformValues = new Uint32Array(
+    AgentGenerationPipeline.UNIFORM_COUNT
+  );
 
   public constructor(
     private readonly device: GPUDevice,
@@ -54,9 +69,47 @@ export class AgentGenerationPipeline {
       ],
     });
 
+    this.compactionBindGroupLayout = device.createBindGroupLayout({
+      entries: [
+        {
+          binding: 0,
+          visibility: GPUShaderStage.COMPUTE,
+          buffer: {
+            type: 'uniform',
+          },
+        },
+        {
+          binding: 1,
+          visibility: GPUShaderStage.COMPUTE,
+          buffer: {
+            type: 'storage',
+          },
+        },
+        {
+          binding: 2,
+          visibility: GPUShaderStage.COMPUTE,
+          buffer: {
+            type: 'storage',
+          },
+        },
+        {
+          binding: 3,
+          visibility: GPUShaderStage.COMPUTE,
+          buffer: {
+            type: 'storage',
+          },
+        },
+      ],
+    });
+
     this.agentsBuffer = this.device.createBuffer({
       size: this.maxAgentCount * AGENT_SIZE_IN_BYTES,
-      usage: GPUBufferUsage.STORAGE,
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
+    });
+
+    this.compactedAgentsBuffer = this.device.createBuffer({
+      size: this.maxAgentCount * AGENT_SIZE_IN_BYTES,
+      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
     });
 
     this.countersBuffer = this.device.createBuffer({
@@ -98,6 +151,36 @@ export class AgentGenerationPipeline {
       ],
     });
 
+    this.compactionBindGroup = this.device.createBindGroup({
+      layout: this.compactionBindGroupLayout,
+      entries: [
+        {
+          binding: 0,
+          resource: {
+            buffer: this.uniforms,
+          },
+        },
+        {
+          binding: 1,
+          resource: {
+            buffer: this.agentsBuffer,
+          },
+        },
+        {
+          binding: 2,
+          resource: {
+            buffer: this.countersBuffer,
+          },
+        },
+        {
+          binding: 3,
+          resource: {
+            buffer: this.compactedAgentsBuffer,
+          },
+        },
+      ],
+    });
+
     this.firstGenerationPipeline = device.createComputePipeline({
       layout: device.createPipelineLayout({
         bindGroupLayouts: [commonState.bindGroupLayout, this.bindGroupLayout],
@@ -122,16 +205,79 @@ export class AgentGenerationPipeline {
         entryPoint: 'main',
       },
     });
+
+    this.resizePipeline = device.createComputePipeline({
+      layout: device.createPipelineLayout({
+        bindGroupLayouts: [commonState.bindGroupLayout, this.bindGroupLayout],
+      }),
+      compute: {
+        module: smartCompile(device, CommonState.shaderCode, agentSchema, resizeShader),
+        entryPoint: 'main',
+      },
+    });
+
+    this.compactionPipeline = device.createComputePipeline({
+      layout: device.createPipelineLayout({
+        bindGroupLayouts: [commonState.bindGroupLayout, this.compactionBindGroupLayout],
+      }),
+      compute: {
+        module: smartCompile(
+          device,
+          CommonState.shaderCode,
+          agentSchema,
+          compactionShader
+        ),
+        entryPoint: 'main',
+      },
+    });
   }
 
   public get maxAgentCount(): number {
     return Math.min(
-      this.maxAgentCountUpperLimit,
+      Number.isFinite(this.maxAgentCountUpperLimit)
+        ? this.maxAgentCountUpperLimit
+        : Number.POSITIVE_INFINITY,
       Math.floor(this.device.limits.maxBufferSize / AGENT_SIZE_IN_BYTES) - 1,
       this.device.limits.maxComputeWorkgroupsPerDimension ** 3
     );
   }
 
+  public writeAgents(agentOffset: number, data: Float32Array): void {
+    this.device.queue.writeBuffer(
+      this.agentsBuffer,
+      agentOffset * AGENT_SIZE_IN_BYTES,
+      data
+    );
+  }
+
+  public resizeAgents(agentCount: number, scale: vec2): void {
+    if (agentCount <= 0 || vec2.equals(scale, vec2.fromValues(1, 1))) {
+      return;
+    }
+
+    this.device.queue.writeBuffer(
+      this.uniforms,
+      0,
+      new Float32Array([scale[0], scale[1], agentCount, 0])
+    );
+
+    const commandEncoder = this.device.createCommandEncoder();
+    const passEncoder = commandEncoder.beginComputePass();
+    this.commonState.execute(passEncoder);
+    passEncoder.setPipeline(this.resizePipeline);
+    passEncoder.setBindGroup(1, this.bindGroup);
+    passEncoder.dispatchWorkgroups(
+      ...getWorkgroupCounts(
+        this.device,
+        agentCount,
+        AgentGenerationPipeline.WORKGROUP_SIZE
+      )
+    );
+    passEncoder.end();
+
+    this.device.queue.submit([commandEncoder.finish()]);
+  }
+
   public spawnFirstGeneration(): void {
     const commandEncoder = this.device.createCommandEncoder();
 
@@ -152,8 +298,11 @@ export class AgentGenerationPipeline {
   }
 
   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]));
+    this.counterClearValues.fill(0);
+    this.agentCountUniformValues.fill(0);
+    this.agentCountUniformValues[0] = agentCount;
+    this.device.queue.writeBuffer(this.countersBuffer, 0, this.counterClearValues);
+    this.device.queue.writeBuffer(this.uniforms, 0, this.agentCountUniformValues);
 
     const commandEncoder = this.device.createCommandEncoder();
 
@@ -190,10 +339,62 @@ export class AgentGenerationPipeline {
     };
   }
 
+  public async compactAgents(agentCount: number): Promise<number> {
+    if (agentCount <= 0) {
+      return 0;
+    }
+
+    this.counterClearValues.fill(0);
+    this.agentCountUniformValues.fill(0);
+    this.agentCountUniformValues[0] = agentCount;
+    this.device.queue.writeBuffer(this.countersBuffer, 0, this.counterClearValues);
+    this.device.queue.writeBuffer(this.uniforms, 0, this.agentCountUniformValues);
+
+    const commandEncoder = this.device.createCommandEncoder();
+    const passEncoder = commandEncoder.beginComputePass();
+    passEncoder.setPipeline(this.compactionPipeline);
+    this.commonState.execute(passEncoder);
+    passEncoder.setBindGroup(1, this.compactionBindGroup);
+    passEncoder.dispatchWorkgroups(
+      ...getWorkgroupCounts(
+        this.device,
+        agentCount,
+        AgentGenerationPipeline.WORKGROUP_SIZE
+      )
+    );
+    passEncoder.end();
+
+    commandEncoder.copyBufferToBuffer(
+      this.compactedAgentsBuffer,
+      0,
+      this.agentsBuffer,
+      0,
+      agentCount * AGENT_SIZE_IN_BYTES
+    );
+    commandEncoder.copyBufferToBuffer(
+      this.countersBuffer,
+      0,
+      this.countersStagingBuffer,
+      0,
+      Uint32Array.BYTES_PER_ELEMENT
+    );
+
+    this.device.queue.submit([commandEncoder.finish()]);
+
+    await this.countersStagingBuffer.mapAsync(GPUMapMode.READ);
+    const compactedCount = new Uint32Array(
+      this.countersStagingBuffer.getMappedRange().slice(0, Uint32Array.BYTES_PER_ELEMENT)
+    )[0];
+    this.countersStagingBuffer.unmap();
+
+    return compactedCount;
+  }
+
   public destroy() {
     this.uniforms.destroy();
     this.countersBuffer.destroy();
     this.countersStagingBuffer.destroy();
+    this.compactedAgentsBuffer.destroy();
     this.agentsBuffer.destroy();
   }
 }