Generate agents with compute shader

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

@@ -0,0 +1,79 @@
+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, Agent } from '../agent';
+import shader from './agent-generation.wgsl';
+import agentSchema from './agent-schema.wgsl';
+
+export class AgentGenerationPipeline {
+  private static readonly WORKGROUP_SIZE = 64;
+
+  private readonly bindGroupLayout: GPUBindGroupLayout;
+  private readonly pipeline: GPUComputePipeline;
+
+  private bindGroup?: GPUBindGroup;
+
+  public constructor(
+    private readonly device: GPUDevice,
+    private readonly commonState: CommonState
+  ) {
+    this.bindGroupLayout = device.createBindGroupLayout({
+      entries: [
+        {
+          binding: 1,
+          visibility: GPUShaderStage.COMPUTE,
+          buffer: {
+            type: 'storage',
+          },
+        },
+      ],
+    });
+
+    this.pipeline = device.createComputePipeline({
+      layout: device.createPipelineLayout({
+        bindGroupLayouts: [commonState.bindGroupLayout, this.bindGroupLayout],
+      }),
+      compute: {
+        module: smartCompile(device, CommonState.shaderCode, random, agentSchema, shader),
+        entryPoint: 'main',
+      },
+    });
+  }
+
+  public generateAgents(agentCount: number): GPUBuffer {
+    if (agentCount <= 0 || agentCount != Math.floor(agentCount)) {
+      throw new Error('Agent count must be a positive integer');
+    }
+
+    const agentsBuffer = this.device.createBuffer({
+      size: agentCount * AGENT_SIZE_IN_BYTES,
+      usage: GPUBufferUsage.STORAGE,
+    });
+
+    this.bindGroup = this.device.createBindGroup({
+      layout: this.bindGroupLayout,
+      entries: [
+        {
+          binding: 1,
+          resource: {
+            buffer: agentsBuffer,
+          },
+        },
+      ],
+    });
+
+    const commandEncoder = this.device.createCommandEncoder();
+
+    const passEncoder = commandEncoder.beginComputePass();
+    passEncoder.setPipeline(this.pipeline);
+    this.commonState.execute(passEncoder);
+    passEncoder.setBindGroup(1, this.bindGroup);
+    passEncoder.dispatchWorkgroups(
+      Math.ceil(agentCount / AgentGenerationPipeline.WORKGROUP_SIZE)
+    );
+    passEncoder.end();
+
+    this.device.queue.submit([commandEncoder.finish()]);
+    return agentsBuffer;
+  }
+}

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

@@ -0,0 +1,26 @@
+@compute @workgroup_size(64)
+fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
+  let id = global_id.x;
+
+  if (id >= arrayLength(&agents)) {
+    return;
+  }
+
+  let position = vec2(
+    hash(id) * state.size.x,
+    hash(id * id) * state.size.y,
+  );
+
+  let center = state.size / 2.0;
+
+  let direction = position - center;
+  let angle = atan2(direction.y, direction.x);
+
+  agents[id] = Agent(
+    position,
+    angle,
+    0,
+    1000000,
+    0
+  );
+}

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

@@ -0,0 +1,9 @@
+struct Agent {
+  position: vec2<f32>,
+  angle: f32,
+  species: f32,
+  timeToLive: f32,
+  timeToLive2: f32,
+}
+
+@group(1) @binding(1) var<storage, read_write> agents: array<Agent>;

src/pipelines/agents/agent-pipeline.ts

@@ -1,6 +1,8 @@
+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, Agent } from './agent';
+import agentSchme from './agent-generation/agent-schema.wgsl';
 import { AgentSettings } from './agent-settings';
 import shader from './agent.wgsl';
 
@@ -11,7 +13,6 @@ export class AgentPipeline {
   private readonly bindGroupLayout: GPUBindGroupLayout;
   private readonly pipeline: GPUComputePipeline;
   private readonly uniforms: GPUBuffer;
-  private readonly agentsBuffer: GPUBuffer;
 
   private bindGroup?: GPUBindGroup;
   private previousTrailMapIn?: GPUTextureView;
@@ -19,13 +20,9 @@ export class AgentPipeline {
 
   public constructor(
     private readonly device: GPUDevice,
-    agents: Array<Agent>,
-    private readonly commonState: CommonState
+    private readonly commonState: CommonState,
+    private readonly agentsBuffer: GPUBuffer
   ) {
-    if (agents.length === 0) {
-      throw new Error('No agents provided');
-    }
-
     this.bindGroupLayout = device.createBindGroupLayout(AgentPipeline.bindGroupLayout);
 
     this.pipeline = device.createComputePipeline({
@@ -33,7 +30,7 @@ export class AgentPipeline {
         bindGroupLayouts: [commonState.bindGroupLayout, this.bindGroupLayout],
       }),
       compute: {
-        module: smartCompile(device, CommonState.shaderCode, shader),
+        module: smartCompile(device, CommonState.shaderCode, random, agentSchme, shader),
         entryPoint: 'main',
       },
     });
@@ -42,23 +39,6 @@ export class AgentPipeline {
       size: AgentPipeline.UNIFORM_COUNT * Float32Array.BYTES_PER_ELEMENT,
       usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
     });
-
-    this.agentsBuffer = device.createBuffer({
-      size: agents.length * AGENT_SIZE_IN_BYTES,
-      usage: GPUBufferUsage.STORAGE,
-      mappedAtCreation: true,
-    });
-
-    new Float32Array(this.agentsBuffer.getMappedRange()).set(
-      agents.flatMap((agent) => [
-        ...agent.position,
-        agent.angle,
-        0, // padding
-        agent.species,
-        agent.timeToLive,
-      ])
-    );
-    this.agentsBuffer.unmap();
   }
 
   public setParameters({
@@ -81,25 +61,6 @@ export class AgentPipeline {
     );
   }
 
-  public executeRenderPass(
-    commandEncoder: GPUCommandEncoder,
-    trailMapIn: GPUTextureView,
-    trailMapOut: GPUTextureView
-  ) {
-    this.ensureBindGroupExists(trailMapIn, trailMapOut);
-
-    const passEncoder = commandEncoder.beginComputePass();
-    passEncoder.setPipeline(this.pipeline);
-    this.commonState.execute(passEncoder);
-    passEncoder.setBindGroup(1, this.bindGroup);
-    passEncoder.dispatchWorkgroups(
-      Math.ceil(
-        this.agentsBuffer.size / AGENT_SIZE_IN_BYTES / AgentPipeline.WORKGROUP_SIZE
-      )
-    );
-    passEncoder.end();
-  }
-
   public execute(
     commandEncoder: GPUCommandEncoder,
     trailMapIn: GPUTextureView,

src/pipelines/agents/agent.wgsl

@@ -1,10 +1,3 @@
-struct Agent {
-  position: vec2<f32>,
-  angle: f32,
-  species: f32,
-  timeToLive: f32
-}
-
 struct Settings {
   brushTrailWeight: f32,
   moveRate: f32,
@@ -14,7 +7,6 @@ struct Settings {
 };
 
 @group(1) @binding(0) var<uniform> settings: Settings;
-@group(1) @binding(1) var<storage, read_write> agents: array<Agent>;
 @group(1) @binding(2) var trailMapIn: texture_2d<f32>;
 @group(1) @binding(3) var trailMapOut: texture_storage_2d<rgba16float, write>;
 
@@ -28,31 +20,31 @@ fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
 
   var agent = agents[id];
 
-  if (agent.timeToLive <= 0.) {
-    // agent.position = vec2(
-    //   random_with_seed(agent.position, f32(id) + state.time),
-    //   random_with_seed(agent.position, f32(id) + state.time + 12),
-    // );
-    // agent.angle = random_with_seed(vec2(agent.angle), f32(id) + state.time);
-    // agent.species = 1;
-    // agent.timeToLive = 1000;
-    // agents[id] = agent;
-    // return;
-  }
+  // if (agent.timeToLive <= 0.) {
+  //   agent.position = vec2(
+  //     random_with_seed(agent.position, f32(id) + state.time),
+  //     random_with_seed(agent.position, f32(id) + state.time + 12),
+  //   );
+  //   agent.angle = random_with_seed(vec2(agent.angle), f32(id) + state.time);
+  //   agent.species = 1;
+  //   agent.timeToLive = 1000;
+  //   agents[id] = agent;
+  //   return;
+  // }
 
-  let random = random_with_seed(agent.position, f32(id) + state.time);
+  let random = hash(id + u32(state.time * 16732.0));
   let trailCurrent = textureLoad(trailMapIn, vec2<i32>(agent.position), 0);
 
-  var weight: f32;
-  if(agent.species == 0) {
-    weight = trailCurrent.r - trailCurrent.g;
-  } else {
-    weight = trailCurrent.g - trailCurrent.r;
-  }
-  if (weight < 0) {
-    agent.timeToLive = 0;
-    return;
-  }
+  // var weight: f32;
+  // if(agent.species == 0) {
+  //   weight = trailCurrent.r - trailCurrent.g;
+  // } else {
+  //   weight = trailCurrent.g - trailCurrent.r;
+  // }
+  // if (weight < 0) {
+  //   agent.timeToLive = 0;
+  //   return;
+  // }
 
   
   let trailForward = sense(agent.position, agent.angle, settings.sensorOffset, 0);
@@ -111,7 +103,3 @@ fn sense(agentPosition: vec2<f32>, agentAngle: f32, sensorOffset: f32, sensorOff
   let sensorPosition = vec2<i32>(agentPosition + sensorDirection * sensorOffset);
   return textureLoad(trailMapIn, sensorPosition, 0); 
 }
-
-fn random_with_seed(uv: vec2<f32>, seed: f32) -> f32 {
-  return fract(sin(dot(uv, vec2(12.9898 + seed, 78.233 + seed)))* 43758.5453123 + seed);
-}