This commit is contained in:
Andras Schmelczer 2026-05-18 08:11:52 +01:00
commit 15e99380b5
84 changed files with 1659 additions and 4726 deletions

View file

@ -16,12 +16,17 @@ struct Counters {
var<workgroup> workgroupAliveCount: atomic<u32>;
var<workgroup> workgroupCompactedOffset: u32;
fn dead_agent() -> Agent {
return Agent(vec2<f32>(0.0, 0.0), 0.0, -1.0, vec2<f32>(-1.0, -1.0), 0.0, 0.0);
}
@compute @workgroup_size(64)
fn main(
@builtin(global_invocation_id) global_id: vec3<u32>,
@builtin(local_invocation_id) local_id: vec3<u32>
@builtin(local_invocation_id) local_id: vec3<u32>,
@builtin(num_workgroups) num_workgroups: vec3<u32>
) {
let id = get_id(global_id);
let id = get_id(global_id, num_workgroups);
if local_id.x == 0u {
atomicStore(&workgroupAliveCount, 0u);
@ -30,7 +35,7 @@ fn main(
workgroupBarrier();
var localCompactedIndex = 0u;
var agent = Agent(vec2<f32>(0.0, 0.0), 0.0, -1.0, vec2<f32>(-1.0, -1.0), 0.0, 0.0);
var agent = dead_agent();
var isAlive = false;
if id < settings.agentCount {
agent = agents[id];
@ -57,3 +62,20 @@ fn main(
compactedAgents[workgroupCompactedOffset + localCompactedIndex] = agent;
}
}
@compute @workgroup_size(64)
fn clearCompactedTail(
@builtin(global_invocation_id) global_id: vec3<u32>,
@builtin(num_workgroups) num_workgroups: vec3<u32>
) {
let id = get_id(global_id, num_workgroups);
if id >= settings.agentCount {
return;
}
let aliveAgentCount = atomicLoad(&counters.aliveAgentCount);
if id >= aliveAgentCount {
compactedAgents[id] = dead_agent();
}
}

View file

@ -1,210 +0,0 @@
import { describe, expect, it, vi } from 'vitest';
import { AGENT_SIZE_IN_BYTES } from './agent';
import { AgentGenerationPipeline } from './agent-generation-pipeline';
const installGpuConstants = () => {
Object.defineProperties(globalThis, {
GPUBufferUsage: {
configurable: true,
value: {
MAP_READ: 1,
COPY_DST: 2,
COPY_SRC: 4,
STORAGE: 8,
UNIFORM: 16,
},
},
GPUMapMode: {
configurable: true,
value: {
READ: 1,
},
},
GPUShaderStage: {
configurable: true,
value: {
COMPUTE: 1,
},
},
});
};
type CopyCall = {
source: GPUBuffer;
sourceOffset: number;
destination: GPUBuffer;
destinationOffset: number;
size: number;
};
type DispatchCall = {
entryPoint: string;
workgroups: [number, number, number];
};
type FakePipeline = {
entryPoint: string;
};
class FakeBuffer {
private readonly mappedRange: ArrayBuffer;
public readonly destroy = vi.fn();
public readonly mapAsync = vi.fn(async () => undefined);
public readonly getMappedRange = vi.fn(() => this.mappedRange);
public readonly unmap = vi.fn();
public constructor(
public readonly label: string,
size: number,
mappedValue = 0
) {
this.mappedRange = new ArrayBuffer(Math.max(size, Uint32Array.BYTES_PER_ELEMENT));
new Uint32Array(this.mappedRange)[0] = mappedValue;
}
}
class FakeComputePass {
private pipeline: FakePipeline | null = null;
public readonly setPipeline = vi.fn((pipeline: GPUComputePipeline) => {
this.pipeline = pipeline as unknown as FakePipeline;
});
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',
workgroups: [x, y, z],
});
});
public readonly end = vi.fn();
public constructor(private readonly device: FakeDevice) {}
}
class FakeCommandEncoder {
public readonly beginComputePass = vi.fn(() => new FakeComputePass(this.device));
public readonly copyBufferToBuffer = vi.fn(
(
source: GPUBuffer,
sourceOffset: number,
destination: GPUBuffer,
destinationOffset: number,
size: number
) => {
this.device.copyCalls.push({
source,
sourceOffset,
destination,
destinationOffset,
size,
});
}
);
public readonly finish = vi.fn(() => ({}) as GPUCommandBuffer);
public constructor(private readonly device: FakeDevice) {}
}
class FakeQueue {
public readonly writeBuffer = vi.fn(() => undefined);
public readonly submit = vi.fn(() => undefined);
}
class FakeShaderModule {
public readonly getCompilationInfo = vi.fn(async () => ({
messages: [],
}));
}
class FakeDevice {
public readonly copyCalls: Array<CopyCall> = [];
public readonly dispatchCalls: Array<DispatchCall> = [];
public readonly createdComputeEntryPoints: Array<string> = [];
public readonly limits = {
maxBufferSize: 1024 * 1024 * 1024,
maxComputeWorkgroupsPerDimension: 65_535,
};
public readonly queue = new FakeQueue();
private bufferIndex = 0;
public readonly createBindGroupLayout = vi.fn(() => ({}) as GPUBindGroupLayout);
public readonly createBuffer = vi.fn((descriptor: GPUBufferDescriptor) => {
const label =
['agents', 'compactedAgents', 'counters', 'countersStaging', 'uniforms'][
this.bufferIndex
] ?? `buffer${this.bufferIndex}`;
this.bufferIndex += 1;
const isMappedReadBuffer = (Number(descriptor.usage) & GPUBufferUsage.MAP_READ) !== 0;
return new FakeBuffer(
label,
Number(descriptor.size),
isMappedReadBuffer ? this.compactedCount : 0
) as unknown as GPUBuffer;
});
public readonly createBindGroup = vi.fn(() => ({}) as GPUBindGroup);
public readonly createPipelineLayout = vi.fn(() => ({}) as GPUPipelineLayout);
public readonly createShaderModule = vi.fn(
() => new FakeShaderModule() as unknown as GPUShaderModule
);
public readonly createComputePipeline = vi.fn(
(descriptor: GPUComputePipelineDescriptor) => {
const pipeline = {
entryPoint: descriptor.compute.entryPoint ?? 'main',
};
this.createdComputeEntryPoints.push(pipeline.entryPoint);
return pipeline as unknown as GPUComputePipeline;
}
);
public readonly createCommandEncoder = vi.fn(() => new FakeCommandEncoder(this));
public constructor(private readonly compactedCount: number) {}
}
const createPipeline = (compactedCount: number) => {
installGpuConstants();
const device = new FakeDevice(compactedCount);
return {
device,
pipeline: new AgentGenerationPipeline(device as unknown as GPUDevice, 1024),
};
};
describe('AgentGenerationPipeline compaction', () => {
it('swaps compacted agents into the active buffer without a copy-back dispatch', async () => {
const agentCount = 10;
const { device, pipeline } = createPipeline(3);
await expect(pipeline.compactAgents(agentCount)).resolves.toBe(3);
expect(device.createdComputeEntryPoints).not.toContain('copyCompactedAgents');
expect(device.dispatchCalls.map((call) => call.entryPoint)).toEqual(['main']);
expect(device.copyCalls.map((call) => call.size)).toEqual([
Uint32Array.BYTES_PER_ELEMENT,
]);
expect(
device.copyCalls.some((call) => call.size === agentCount * AGENT_SIZE_IN_BYTES)
).toBe(false);
expect(device.queue.submit).toHaveBeenCalledTimes(1);
pipeline.destroy();
});
it('does not encode work for empty compaction requests', async () => {
const { device, pipeline } = createPipeline(0);
await expect(pipeline.compactAgents(0)).resolves.toBe(0);
expect(device.dispatchCalls).toEqual([]);
expect(device.copyCalls).toEqual([]);
expect(device.queue.submit).not.toHaveBeenCalled();
pipeline.destroy();
});
});

View file

@ -1,14 +1,19 @@
import { vec2 } from 'gl-matrix';
import { getWorkgroupCount } from '../../../utils/graphics/get-workgroup-count';
import { smartCompile } from '../../../utils/graphics/smart-compile';
import { AGENT_SIZE_IN_BYTES } from './agent';
import {
AGENT_MAX_DISPATCHABLE_COUNT,
dispatchAgentWorkgroups,
} from '../agent-dispatch';
import compactionShader from './agent-compaction.wgsl?raw';
import resizeShader from './agent-resize.wgsl?raw';
import agentSchema from './agent-schema.wgsl?raw';
export const AGENT_FLOAT_COUNT = 8;
export const AGENT_SIZE_IN_BYTES =
AGENT_FLOAT_COUNT * Float32Array.BYTES_PER_ELEMENT;
export class AgentGenerationPipeline {
private static readonly WORKGROUP_SIZE = 64;
private static readonly UNIFORM_COUNT = 4;
private static readonly COUNTER_COUNT = 1;
@ -21,9 +26,11 @@ export class AgentGenerationPipeline {
private readonly resizePipeline: GPUComputePipeline;
private readonly compactionPipeline: GPUComputePipeline;
private readonly clearCompactedTailPipeline: GPUComputePipeline;
private activeAgentsBuffer: GPUBuffer;
private inactiveAgentsBuffer: GPUBuffer;
private allocatedMaxAgentCount: number;
private readonly countersBuffer: GPUBuffer;
private readonly countersStagingBuffer: GPUBuffer;
private readonly counterClearValues = new Uint32Array(
@ -35,8 +42,10 @@ export class AgentGenerationPipeline {
public constructor(
private readonly device: GPUDevice,
private readonly maxAgentCountUpperLimit: number
initialMaxAgentCount: number,
private readonly maxAgentCountUpperLimit = Number.POSITIVE_INFINITY
) {
this.allocatedMaxAgentCount = this.clampMaxAgentCount(initialMaxAgentCount);
const emptyBindGroupLayout = device.createBindGroupLayout({ entries: [] });
this.bindGroupLayout = device.createBindGroupLayout({
entries: [
@ -72,7 +81,7 @@ export class AgentGenerationPipeline {
});
this.activeAgentsBuffer = this.createAgentsBuffer();
this.inactiveAgentsBuffer = this.createAgentsBuffer();
this.inactiveAgentsBuffer = this.createInactivePlaceholderBuffer();
this.countersBuffer = this.device.createBuffer({
size: AgentGenerationPipeline.COUNTER_COUNT * Uint32Array.BYTES_PER_ELEMENT,
@ -110,6 +119,16 @@ export class AgentGenerationPipeline {
entryPoint: 'main',
},
});
this.clearCompactedTailPipeline = device.createComputePipeline({
layout: device.createPipelineLayout({
bindGroupLayouts: [emptyBindGroupLayout, this.bindGroupLayout],
}),
compute: {
module: compactionModule,
entryPoint: 'clearCompactedTail',
},
});
}
public get agentsBuffer(): GPUBuffer {
@ -118,23 +137,86 @@ export class AgentGenerationPipeline {
private createAgentsBuffer(): GPUBuffer {
return this.device.createBuffer({
size: this.maxAgentCount * AGENT_SIZE_IN_BYTES,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
size: this.allocatedMaxAgentCount * AGENT_SIZE_IN_BYTES,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST | GPUBufferUsage.COPY_SRC,
});
}
// 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;
}
public get maxSupportedAgentCount(): number {
return this.clampMaxAgentCount(Number.POSITIVE_INFINITY);
}
public ensureMaxAgentCount(
requestedMaxAgentCount: number,
activeAgentCount: number
): number {
const nextMaxAgentCount = this.clampMaxAgentCount(requestedMaxAgentCount);
if (nextMaxAgentCount <= this.allocatedMaxAgentCount) {
return this.allocatedMaxAgentCount;
}
const previousActiveAgentsBuffer = this.activeAgentsBuffer;
const previousMaxAgentCount = this.allocatedMaxAgentCount;
this.allocatedMaxAgentCount = nextMaxAgentCount;
this.activeAgentsBuffer = this.createAgentsBuffer();
const copyAgentCount = Math.min(
Math.max(0, Math.floor(activeAgentCount)),
previousMaxAgentCount,
nextMaxAgentCount
);
if (copyAgentCount > 0) {
const commandEncoder = this.device.createCommandEncoder();
commandEncoder.copyBufferToBuffer(
previousActiveAgentsBuffer,
0,
this.activeAgentsBuffer,
0,
copyAgentCount * AGENT_SIZE_IN_BYTES
);
this.device.queue.submit([commandEncoder.finish()]);
}
// 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.
previousActiveAgentsBuffer.destroy();
return this.allocatedMaxAgentCount;
}
private clampMaxAgentCount(value: number): number {
const requestedMaxAgentCount =
value === Number.POSITIVE_INFINITY
? Number.POSITIVE_INFINITY
: Number.isFinite(value)
? Math.floor(value)
: 0;
return Math.min(
Number.isFinite(this.maxAgentCountUpperLimit)
? this.maxAgentCountUpperLimit
: Number.POSITIVE_INFINITY,
Math.floor(this.device.limits.maxBufferSize / AGENT_SIZE_IN_BYTES) - 1,
Math.floor(this.device.limits.maxBufferSize / AGENT_SIZE_IN_BYTES),
Math.floor(
((this.device.limits as GPUSupportedLimits).maxStorageBufferBindingSize ??
this.device.limits.maxBufferSize) / AGENT_SIZE_IN_BYTES
) - 1,
this.device.limits.maxComputeWorkgroupsPerDimension *
AgentGenerationPipeline.WORKGROUP_SIZE
),
AGENT_MAX_DISPATCHABLE_COUNT,
Math.max(0, requestedMaxAgentCount)
);
}
@ -161,9 +243,7 @@ export class AgentGenerationPipeline {
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(this.resizePipeline);
passEncoder.setBindGroup(1, this.getBindGroup());
passEncoder.dispatchWorkgroups(
getWorkgroupCount(agentCount, AgentGenerationPipeline.WORKGROUP_SIZE)
);
dispatchAgentWorkgroups(passEncoder, agentCount);
passEncoder.end();
this.device.queue.submit([commandEncoder.finish()]);
@ -174,6 +254,12 @@ 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.countersBuffer, 0, this.counterClearValues);
this.device.queue.writeBuffer(this.uniforms, 0, this.agentCountUniformValues);
@ -182,9 +268,9 @@ export class AgentGenerationPipeline {
const passEncoder = commandEncoder.beginComputePass();
passEncoder.setPipeline(this.compactionPipeline);
passEncoder.setBindGroup(1, this.getBindGroup());
passEncoder.dispatchWorkgroups(
getWorkgroupCount(agentCount, AgentGenerationPipeline.WORKGROUP_SIZE)
);
dispatchAgentWorkgroups(passEncoder, agentCount);
passEncoder.setPipeline(this.clearCompactedTailPipeline);
dispatchAgentWorkgroups(passEncoder, agentCount);
passEncoder.end();
commandEncoder.copyBufferToBuffer(
@ -196,6 +282,14 @@ 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(
@ -204,7 +298,6 @@ export class AgentGenerationPipeline {
1
)[0];
this.countersStagingBuffer.unmap();
this.swapAgentBuffers();
return compactedCount;
}

View file

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

View file

@ -1,89 +0,0 @@
import { describe, expect, it } from 'vitest';
import { AGENT_FLOAT_COUNT, AGENT_SIZE_IN_BYTES } from './agent';
import compactionShader from './agent-compaction.wgsl?raw';
import resizeShader from './agent-resize.wgsl?raw';
import agentSchema from './agent-schema.wgsl?raw';
const wgslFloatCountByType: Record<string, number> = {
f32: 1,
'vec2<f32>': 2,
};
const getAgentStructFields = () => {
const match = /struct Agent\s*\{(?<body>[\s\S]*?)\n\}/.exec(agentSchema);
if (!match?.groups?.body) {
throw new Error('Agent struct was not found in agent-schema.wgsl');
}
return match.groups.body
.split('\n')
.map((line) => line.trim().replace(/,$/, ''))
.filter(Boolean)
.map((line) => {
const fieldMatch = /^(?<name>\w+):\s*(?<type>[^,]+)$/.exec(line);
if (!fieldMatch?.groups) {
throw new Error(`Unsupported Agent field syntax: ${line}`);
}
return {
name: fieldMatch.groups.name,
type: fieldMatch.groups.type,
};
});
};
describe('Agent TS/WGSL contract', () => {
it('keeps the TypeScript float count aligned with the WGSL Agent struct', () => {
const fields = getAgentStructFields();
const wgslFloatCount = fields.reduce((sum, field) => {
const count = wgslFloatCountByType[field.type];
if (!count) {
throw new Error(`Unsupported WGSL Agent field type: ${field.type}`);
}
return sum + count;
}, 0);
expect(fields.map((field) => field.name)).toEqual([
'position',
'angle',
'colorIndex',
'targetPosition',
'targetAngle',
'introDelay',
]);
expect(wgslFloatCount).toBe(AGENT_FLOAT_COUNT);
expect(AGENT_SIZE_IN_BYTES).toBe(AGENT_FLOAT_COUNT * Float32Array.BYTES_PER_ELEMENT);
});
it('keeps generation shader workgroup sizes aligned with agent indexing', () => {
[resizeShader, compactionShader].forEach((shader) => {
expect(shader).toMatch(/@workgroup_size\(64\)/);
});
expect(agentSchema).toContain('return global_id.x;');
expect(compactionShader).toContain('let id = get_id(global_id);');
expect(compactionShader).toContain('if id < settings.agentCount');
});
it('keeps compaction as a ping-pong write without copy-back shader work', () => {
expect(compactionShader).not.toContain('fn copyCompactedAgents');
expect(compactionShader).not.toContain('agents[id] = compactedAgents[id];');
expect(compactionShader).toContain(
'compactedAgents[workgroupCompactedOffset + localCompactedIndex] = agent;'
);
});
it('uses workgroup-local counting before allocating global compacted ranges', () => {
expect(compactionShader).toContain(
'var<workgroup> workgroupAliveCount: atomic<u32>;'
);
expect(compactionShader).toContain(
'localCompactedIndex = atomicAdd(&workgroupAliveCount, 1u);'
);
expect(
compactionShader.match(/atomicAdd\(&counters\.aliveAgentCount/g) ?? []
).toHaveLength(1);
});
});

View file

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

View file

@ -1,2 +0,0 @@
export const AGENT_FLOAT_COUNT = 8;
export const AGENT_SIZE_IN_BYTES = AGENT_FLOAT_COUNT * Float32Array.BYTES_PER_ELEMENT;

View file

@ -2,15 +2,14 @@ import {
createCachedFloat32BufferWrite,
writeFloat32BufferIfChanged,
} from '../../utils/graphics/cached-buffer-write';
import { getWorkgroupCount } from '../../utils/graphics/get-workgroup-count';
import { smartCompile } from '../../utils/graphics/smart-compile';
import { CommonState } from '../common-state/common-state';
import { dispatchAgentWorkgroups } from './agent-dispatch';
import agentSchema from './agent-generation/agent-schema.wgsl?raw';
import { AgentSettings } from './agent-settings';
import shader from './agent.wgsl?raw';
export class AgentPipeline {
private static readonly WORKGROUP_SIZE = 64;
private static readonly UNIFORM_COUNT = 33;
private readonly bindGroupLayout: GPUBindGroupLayout;
@ -151,9 +150,7 @@ export class AgentPipeline {
passEncoder.setPipeline(this.pipeline);
this.commonState.execute(passEncoder);
passEncoder.setBindGroup(1, bindGroup);
passEncoder.dispatchWorkgroups(
getWorkgroupCount(this.agentCount, AgentPipeline.WORKGROUP_SIZE)
);
dispatchAgentWorkgroups(passEncoder, this.agentCount);
passEncoder.end();
}

View file

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