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

@ -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;
}