more cleaning up

This commit is contained in:
Andras Schmelczer 2026-05-16 20:45:42 +01:00
commit 560398fefb
110 changed files with 931 additions and 2645 deletions

View file

@ -7,8 +7,6 @@ struct Settings {
struct Counters {
aliveAgentCount: atomic<u32>,
padding0: atomic<u32>,
padding1: atomic<u32>,
};
@group(1) @binding(0) var<uniform> settings: Settings;
@ -22,10 +20,9 @@ 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>
@builtin(local_invocation_id) local_id: vec3<u32>
) {
let id = get_id(global_id, workgroup_count);
let id = get_id(global_id);
if local_id.x == 0u {
atomicStore(&workgroupAliveCount, 0u);
@ -65,10 +62,9 @@ fn main(
@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>
@builtin(local_invocation_id) local_id: vec3<u32>
) {
let id = get_id(global_id, workgroup_count);
let id = get_id(global_id);
if local_id.x == 0u {
workgroupCopyCount = atomicLoad(&counters.aliveAgentCount);

View file

@ -1,35 +0,0 @@
struct Settings {
agentCount: u32 // might be smaller than the length of the agents array
};
@group(1) @binding(0) var<uniform> settings: Settings;
struct Counters {
redAgentsAlive: atomic<u32>,
greenAgentsAlive: 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>,
@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;
}
if agents[id].colorIndex < 0.0 {
return;
}
if agents[id].colorIndex < 0.5 {
atomicAdd(&counters.redAgentsAlive, 1);
} else {
atomicAdd(&counters.greenAgentsAlive, 1);
}
}

View file

@ -1,37 +0,0 @@
@compute @workgroup_size(64)
fn main(
@builtin(global_invocation_id) global_id: vec3<u32>,
@builtin(num_workgroups) workgroup_count: vec3<u32>
) {
let id = get_id(global_id, workgroup_count);
if id >= arrayLength(&agents) {
return;
}
let clusterId = f32(id % 1000);
let random = textureSampleLevel(
noise,
noiseSampler,
vec2(f32(id % 1999) / 2000, f32(id) / 1999 / 2000),
0
);
let randomPosition = textureSampleLevel(
noise,
noiseSampler,
vec2(clusterId / 2000, clusterId / 2000),
0
);
agents[id] = Agent(
randomPosition.xz * state.size,
random.r * 3.14 * 2,
0,
vec2<f32>(-1.0, -1.0),
0.0,
0.0,
);
}

View file

@ -1,6 +1,5 @@
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';
@ -52,7 +51,7 @@ class FakeBuffer {
private readonly mappedRange: ArrayBuffer;
public readonly destroy = vi.fn();
public readonly mapAsync = vi.fn(async (_mode: number) => undefined);
public readonly mapAsync = vi.fn(async () => undefined);
public readonly getMappedRange = vi.fn(() => this.mappedRange);
public readonly unmap = vi.fn();
@ -72,9 +71,7 @@ class FakeComputePass {
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 setBindGroup = vi.fn(() => undefined);
public readonly dispatchWorkgroups = vi.fn((x: number, y = 1, z = 1) => {
this.device.dispatchCalls.push({
entryPoint: this.pipeline?.entryPoint ?? 'unset',
@ -111,12 +108,8 @@ class FakeCommandEncoder {
}
class FakeQueue {
public readonly writeBuffer = vi.fn(
(_buffer: GPUBuffer, _offset: number, _data: BufferSource) => undefined
);
public readonly submit = vi.fn(
(_commandBuffers: Iterable<GPUCommandBuffer>) => undefined
);
public readonly writeBuffer = vi.fn(() => undefined);
public readonly submit = vi.fn(() => undefined);
}
class FakeShaderModule {
@ -137,9 +130,7 @@ class FakeDevice {
private bufferIndex = 0;
public readonly createBindGroupLayout = vi.fn(
(_descriptor: GPUBindGroupLayoutDescriptor) => ({}) as GPUBindGroupLayout
);
public readonly createBindGroupLayout = vi.fn(() => ({}) as GPUBindGroupLayout);
public readonly createBuffer = vi.fn((descriptor: GPUBufferDescriptor) => {
const label =
['agents', 'compactedAgents', 'counters', 'countersStaging', 'uniforms'][
@ -155,15 +146,10 @@ class FakeDevice {
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 createBindGroup = vi.fn(() => ({}) as GPUBindGroup);
public readonly createPipelineLayout = vi.fn(() => ({}) as GPUPipelineLayout);
public readonly createShaderModule = vi.fn(
(_descriptor: GPUShaderModuleDescriptor) =>
new FakeShaderModule() as unknown as GPUShaderModule
() => new FakeShaderModule() as unknown as GPUShaderModule
);
public readonly createComputePipeline = vi.fn(
(descriptor: GPUComputePipelineDescriptor) => {
@ -183,18 +169,10 @@ 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
),
pipeline: new AgentGenerationPipeline(device as unknown as GPUDevice, 1024),
};
};

View file

@ -1,37 +1,29 @@
import { vec2 } from 'gl-matrix';
import { getWorkgroupCounts } from '../../../utils/graphics/get-workgroup-counts';
import { getWorkgroupCount } from '../../../utils/graphics/get-workgroup-count';
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 = 4;
private static readonly COUNTER_COUNT = 3;
private static readonly COUNTER_COUNT = 1;
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;
private readonly compactedAgentsCopyPipeline: GPUComputePipeline;
public readonly agentsBuffer: GPUBuffer;
private readonly compactedAgentsBuffer: GPUBuffer;
public readonly countersBuffer: GPUBuffer;
public readonly countersStagingBuffer: GPUBuffer;
private readonly countersBuffer: GPUBuffer;
private readonly countersStagingBuffer: GPUBuffer;
private readonly counterClearValues = new Uint32Array(
AgentGenerationPipeline.COUNTER_COUNT
);
@ -41,36 +33,10 @@ export class AgentGenerationPipeline {
public constructor(
private readonly device: GPUDevice,
private readonly commonState: CommonState,
private readonly maxAgentCountUpperLimit: number
) {
const emptyBindGroupLayout = device.createBindGroupLayout({ entries: [] });
this.bindGroupLayout = 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',
},
},
],
});
this.compactionBindGroupLayout = device.createBindGroupLayout({
entries: [
{
binding: 0,
@ -130,30 +96,6 @@ export class AgentGenerationPipeline {
this.bindGroup = this.device.createBindGroup({
layout: this.bindGroupLayout,
entries: [
{
binding: 0,
resource: {
buffer: this.uniforms,
},
},
{
binding: 1,
resource: {
buffer: this.agentsBuffer,
},
},
{
binding: 2,
resource: {
buffer: this.countersBuffer,
},
},
],
});
this.compactionBindGroup = this.device.createBindGroup({
layout: this.compactionBindGroupLayout,
entries: [
{
binding: 0,
@ -182,51 +124,21 @@ export class AgentGenerationPipeline {
],
});
this.firstGenerationPipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [commonState.bindGroupLayout, this.bindGroupLayout],
}),
compute: {
module: smartCompile(
device,
CommonState.shaderCode,
agentSchema,
firstGenerationShader
),
entryPoint: 'main',
},
});
this.countingPipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [commonState.bindGroupLayout, this.bindGroupLayout],
}),
compute: {
module: smartCompile(device, CommonState.shaderCode, agentSchema, countingShader),
entryPoint: 'main',
},
});
this.resizePipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [commonState.bindGroupLayout, this.bindGroupLayout],
bindGroupLayouts: [emptyBindGroupLayout, this.bindGroupLayout],
}),
compute: {
module: smartCompile(device, CommonState.shaderCode, agentSchema, resizeShader),
module: smartCompile(device, agentSchema, resizeShader),
entryPoint: 'main',
},
});
const compactionModule = smartCompile(
device,
CommonState.shaderCode,
agentSchema,
compactionShader
);
const compactionModule = smartCompile(device, agentSchema, compactionShader);
this.compactionPipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [commonState.bindGroupLayout, this.compactionBindGroupLayout],
bindGroupLayouts: [emptyBindGroupLayout, this.bindGroupLayout],
}),
compute: {
module: compactionModule,
@ -236,7 +148,7 @@ export class AgentGenerationPipeline {
this.compactedAgentsCopyPipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [commonState.bindGroupLayout, this.compactionBindGroupLayout],
bindGroupLayouts: [emptyBindGroupLayout, this.bindGroupLayout],
}),
compute: {
module: compactionModule,
@ -251,7 +163,8 @@ export class AgentGenerationPipeline {
? this.maxAgentCountUpperLimit
: Number.POSITIVE_INFINITY,
Math.floor(this.device.limits.maxBufferSize / AGENT_SIZE_IN_BYTES) - 1,
this.device.limits.maxComputeWorkgroupsPerDimension ** 3
this.device.limits.maxComputeWorkgroupsPerDimension *
AgentGenerationPipeline.WORKGROUP_SIZE
);
}
@ -276,82 +189,16 @@ export class AgentGenerationPipeline {
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
)
getWorkgroupCount(agentCount, AgentGenerationPipeline.WORKGROUP_SIZE)
);
passEncoder.end();
this.device.queue.submit([commandEncoder.finish()]);
}
public spawnFirstGeneration(): void {
const commandEncoder = this.device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
this.commonState.execute(passEncoder);
passEncoder.setPipeline(this.firstGenerationPipeline);
passEncoder.setBindGroup(1, this.bindGroup);
passEncoder.dispatchWorkgroups(
...getWorkgroupCounts(
this.device,
this.maxAgentCount,
AgentGenerationPipeline.WORKGROUP_SIZE
)
);
passEncoder.end();
this.device.queue.submit([commandEncoder.finish()]);
}
public async countAgents(agentCount: number): Promise<GenerationCounts> {
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.countingPipeline);
this.commonState.execute(passEncoder);
passEncoder.setBindGroup(1, this.bindGroup);
passEncoder.dispatchWorkgroups(
...getWorkgroupCounts(
this.device,
agentCount,
AgentGenerationPipeline.WORKGROUP_SIZE
)
);
passEncoder.end();
commandEncoder.copyBufferToBuffer(
this.countersBuffer,
0,
this.countersStagingBuffer,
0,
AgentGenerationPipeline.COUNTER_COUNT * Uint32Array.BYTES_PER_ELEMENT
);
this.device.queue.submit([commandEncoder.finish()]);
await this.countersStagingBuffer.mapAsync(GPUMapMode.READ);
const data = new Uint32Array(this.countersStagingBuffer.getMappedRange().slice(0));
this.countersStagingBuffer.unmap();
return {
evenGenerationCount: data[0],
oddGenerationCount: data[1],
};
}
public async compactAgents(agentCount: number): Promise<number> {
if (agentCount <= 0) {
return 0;
@ -366,27 +213,17 @@ export class AgentGenerationPipeline {
const commandEncoder = this.device.createCommandEncoder();
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(this.compactionPipeline);
this.commonState.execute(passEncoder);
passEncoder.setBindGroup(1, this.compactionBindGroup);
passEncoder.setBindGroup(1, this.bindGroup);
passEncoder.dispatchWorkgroups(
...getWorkgroupCounts(
this.device,
agentCount,
AgentGenerationPipeline.WORKGROUP_SIZE
)
getWorkgroupCount(agentCount, AgentGenerationPipeline.WORKGROUP_SIZE)
);
passEncoder.end();
const copyPassEncoder = commandEncoder.beginComputePass();
copyPassEncoder.setPipeline(this.compactedAgentsCopyPipeline);
this.commonState.execute(copyPassEncoder);
copyPassEncoder.setBindGroup(1, this.compactionBindGroup);
copyPassEncoder.setBindGroup(1, this.bindGroup);
copyPassEncoder.dispatchWorkgroups(
...getWorkgroupCounts(
this.device,
agentCount,
AgentGenerationPipeline.WORKGROUP_SIZE
)
getWorkgroupCount(agentCount, AgentGenerationPipeline.WORKGROUP_SIZE)
);
copyPassEncoder.end();

View file

@ -8,10 +8,9 @@ struct ResizeSettings {
@compute @workgroup_size(64)
fn main(
@builtin(global_invocation_id) global_id: vec3<u32>,
@builtin(num_workgroups) workgroup_count: vec3<u32>
@builtin(global_invocation_id) global_id: vec3<u32>
) {
let id = get_id(global_id, workgroup_count);
let id = get_id(global_id);
if id >= u32(resizeSettings.agentCount) {
return;

View file

@ -2,8 +2,6 @@ import { describe, expect, it } from 'vitest';
import { AGENT_FLOAT_COUNT, 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';
@ -60,15 +58,12 @@ describe('Agent TS/WGSL contract', () => {
});
it('keeps generation shader workgroup sizes aligned with agent indexing', () => {
[firstGenerationShader, countingShader, resizeShader, compactionShader].forEach(
(shader) => {
expect(shader).toMatch(/@workgroup_size\(64\)/);
}
);
[resizeShader, compactionShader].forEach((shader) => {
expect(shader).toMatch(/@workgroup_size\(64\)/);
});
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(agentSchema).toContain('return global_id.x;');
expect(compactionShader).toContain('let id = get_id(global_id);');
expect(compactionShader).toContain('if id < settings.agentCount');
});

View file

@ -9,6 +9,6 @@ struct Agent {
@group(1) @binding(1) var<storage, read_write> agents: array<Agent>;
fn get_id(global_id: vec3<u32>, workgroup_count: vec3<u32>) -> u32 {
return global_id.x + global_id.y * (workgroup_count.x * 64) + global_id.z * (workgroup_count.x * workgroup_count.y * 64);
}
fn get_id(global_id: vec3<u32>) -> u32 {
return global_id.x;
}

View file

@ -1,4 +0,0 @@
export interface GenerationCounts {
evenGenerationCount: number;
oddGenerationCount: number;
}

View file

@ -2,7 +2,7 @@ import {
createCachedFloat32BufferWrite,
writeFloat32BufferIfChanged,
} from '../../utils/graphics/cached-buffer-write';
import { getWorkgroupCounts } from '../../utils/graphics/get-workgroup-counts';
import { getWorkgroupCount } from '../../utils/graphics/get-workgroup-count';
import { smartCompile } from '../../utils/graphics/smart-compile';
import { CommonState } from '../common-state/common-state';
import agentSchema from './agent-generation/agent-schema.wgsl?raw';
@ -117,7 +117,7 @@ export class AgentPipeline {
this.commonState.execute(passEncoder);
passEncoder.setBindGroup(1, bindGroup);
passEncoder.dispatchWorkgroups(
...getWorkgroupCounts(this.device, this.agentCount, AgentPipeline.WORKGROUP_SIZE)
getWorkgroupCount(this.agentCount, AgentPipeline.WORKGROUP_SIZE)
);
passEncoder.end();
}

View file

@ -25,10 +25,9 @@ struct Settings {
@compute @workgroup_size(64)
fn main(
@builtin(global_invocation_id) global_id: vec3<u32>,
@builtin(num_workgroups) workgroup_count: vec3<u32>
@builtin(global_invocation_id) global_id: vec3<u32>
) {
let id = get_id(global_id, workgroup_count);
let id = get_id(global_id);
if id >= u32(settings.agentCount) {
return;