renderable = {
const w = SIZE_CONST;
const h = SIZE_CONST;
const agentCount = 10_000;
const { device } = await gpu.device();
const sample = gpu.context(w, h, '2d');
sample.fillStyle = '#000';
sample.fillRect(0, 0, w, h);
sample.fillStyle = '#fff';
sample.beginPath();
sample.arc(w / 2, h / 2, w / 4, 0, util.TWO_PI);
sample.fill();
const uniforms = new Float32Array([0, ...Object.values(config)]);
const uniformBuffer = gpu.uniformBuffer(device, uniforms);
const agent = new Float32Array(
util.flatArr(agentCount, () => [
util.randi(w),
util.randi(h),
0,
util.rand(util.TWO_PI),
]),
);
const agentBufferPing = gpu.storageBuffer(device, agent);
const agentBufferPong = gpu.storageBuffer(device, agent);
const trailCountsBuffer = gpu.storageBuffer(device, new Uint32Array(w * h));
const trailTexturePong = gpu.createTexture(device, {
width: sample.canvas.width,
height: sample.canvas.height,
usage: gpu.Usage.StorageTexture,
});
const trailTexturePing = gpu.canvasTexture(device, sample.canvas, {
usage: gpu.Usage.StorageTexture,
});
const bindGroupLayout = device.createBindGroupLayout({
label: 'bind group layout',
entries: [
// agents in
{
binding: 0,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
buffer: { type: 'read-only-storage' },
},
// agents out
{
binding: 1,
visibility: GPUShaderStage.COMPUTE,
buffer: { type: 'storage' },
},
// trail counts
{
binding: 2,
visibility: GPUShaderStage.COMPUTE,
buffer: { type: 'storage' },
},
// trail texture blur
{
binding: 3,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
texture: {},
},
// trail texture
{
binding: 4,
visibility: GPUShaderStage.COMPUTE,
storageTexture: {
format: 'rgba8unorm',
},
},
// uniforms!
{
binding: 5,
visibility: GPUShaderStage.FRAGMENT | GPUShaderStage.COMPUTE,
buffer: {},
},
],
});
const agentShader = `
const TWO_PI: f32 = 6.283185307179586;
struct Uniforms {
time: f32,
sensor_angle: f32,
sensor_distance: f32,
rotation_angle: f32,
step_distance: f32,
decay_rate: f32,
deposit_amount: f32,
blur_strength: f32,
}
struct Agent {
pos: vec2f,
pad: f32,
heading: f32
}
@group(0) @binding(0) var<storage> agents_in: array<Agent>;
@group(0) @binding(1) var<storage, read_write> agents_out: array<Agent>;
@group(0) @binding(2) var<storage, read_write> trail_counts: array<atomic<u32>>;
@group(0) @binding(3) var trails_in: texture_2d<f32>;
@group(0) @binding(5) var<uniform> uniforms: Uniforms;
fn grid_index_from_agent_pos(pos: vec2u) -> u32 {
return (pos.x % ${w}) + (pos.y % ${w}) * ${w};
}
// On generating random numbers, with help of y= [(a+x)sin(bx)] mod 1", W.J.J. Rey, 22nd European Meeting of Statisticians 1998
fn rand11(n: f32) -> f32 { return fract(sin(n) * 43758.5453123); }
fn rand22(n: vec2f) -> f32 { return fract(sin(dot(n, vec2f(12.9898, 4.1414))) * 43758.5453); }
fn shft(x: f32, size: f32) -> f32 {
return x + select(0.0, size, x < 0) + select(0.0, -size, x >= size);
}
fn next_agent_state(agent: Agent) -> Agent {
let p = agent.pos;
let a = agent.heading;
let sa = uniforms.sensor_angle;
let sd = uniforms.sensor_distance;
let ra = uniforms.rotation_angle;
let step_dist = uniforms.step_distance;
let xc = p.x + cos(a) * sd;
let yc = p.y + sin(a) * sd;
let xl = p.x + cos(a - sa) * sd;
let yl = p.y + sin(a - sa) * sd;
let xr = p.x + cos(a + sa) * sd;
let yr = p.y + sin(a + sa) * sd;
// good enough
let C = textureLoad(trails_in, vec2u(u32(xc % ${w}), u32(yc % ${h})), 0).r;
let L = textureLoad(trails_in, vec2u(u32(xl % ${w}), u32(yl % ${h})), 0).r;
let R = textureLoad(trails_in, vec2u(u32(xr % ${w}), u32(yr % ${h})), 0).r;
var d = 0.0;
if (C > L && C > R) {
d = 0.0;
} else if (C < L && C < R) {
d = select(-1.0, 1.0, rand11(uniforms.time) > 0.5);
} else if (L < R) {
d = 1.0;
} else if (R < L) {
d = -1.0;
}
let da = ra * d;
let next_angle = shft(a + da, TWO_PI);
let next_pos = vec2f(
shft(p.x + cos(next_angle) * step_dist, ${w}),
shft(p.y + sin(next_angle) * step_dist, ${h}),
);
return Agent(next_pos, 0, next_angle);
}
@compute
@workgroup_size(${w}, 1, 1)
fn cs(
@builtin(global_invocation_id) gid: vec3u,
@builtin(local_invocation_id) lid: vec3u,
) {
let agent_index = gid.x;
if (agent_index > arrayLength(&agents_in)) {
return;
}
let agent = agents_in[agent_index];
let next_agent = next_agent_state(agent);
agents_out[agent_index].pos = next_agent.pos;
agents_out[agent_index].heading = next_agent.heading;
// add to trail counts
let trail_index = grid_index_from_agent_pos(vec2u(u32(next_agent.pos.x), u32(next_agent.pos.y)));
atomicAdd(&trail_counts[trail_index], 1);
}
`;
// FIXME this is backwards - we should compose all assets first, then run agent compute
const trailCombineShader = `
struct Uniforms {
time: f32,
sensor_angle: f32,
sensor_distance: f32,
rotation_angle: f32,
step_distance: f32,
decay_rate: f32,
deposit_amount: f32,
blur_strength: f32,
}
struct Agent {
pos: vec2f,
pad: f32,
heading: f32
}
@group(0) @binding(2) var<storage, read_write> trail_counts: array<atomic<u32>>;
@group(0) @binding(3) var tex_in: texture_2d<f32>;
@group(0) @binding(4) var tex_out: texture_storage_2d<rgba8unorm, write>;
@group(0) @binding(5) var<uniform> uniforms: Uniforms;
// from https://threejs.org/examples/webgpu_compute_texture_pingpong.html
fn blur(image: texture_2d<f32>, uv: vec2i) -> vec4f {
var color = vec4f(0.0);
color += textureLoad(image, uv + vec2i(-1, 1), 0);
color += textureLoad(image, uv + vec2i(-1, -1), 0);
color += textureLoad(image, uv + vec2i( 0, 0), 0);
color += textureLoad(image, uv + vec2i( 1, -1), 0);
color += textureLoad(image, uv + vec2i( 1, 1), 0);
return color / 5.0;
}
@compute
@workgroup_size(${w}, 1, 1)
fn cs(
@builtin(global_invocation_id) gid: vec3u,
@builtin(local_invocation_id) lid: vec3u,
) {
let y = floor(f32(gid.x) / ${w}.0);
let x = f32(lid.x);
let index = gid.x;
let tex_pos = vec2u(u32(x), u32(y));
let prev = textureLoad(tex_in, tex_pos, 0);
let value = clamp(f32(atomicLoad(&trail_counts[index])) * uniforms.deposit_amount, 0.0, 1.0);
// should be doing "val * 1 - decay_rate" but this looks smoother
let next = vec3(clamp(prev.r + value - uniforms.decay_rate, 0.0, 1.0));
// blur for next pass (this prob isn't in the correct order)
var b = blur(tex_in, vec2i(i32(x), i32(y))) * uniforms.blur_strength;
textureStore(tex_out, tex_pos, vec4f(vec3f(b.r + next), 1.0));
// reset
atomicStore(&trail_counts[index], 0);
}
`;
const agentBindGroupPing = device.createBindGroup({
label: 'agent bind group ping',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: agentBufferPing.buffer } },
{ binding: 1, resource: { buffer: agentBufferPong.buffer } },
{ binding: 2, resource: { buffer: trailCountsBuffer.buffer } },
{ binding: 3, resource: trailTexturePing.createView() },
{ binding: 4, resource: trailTexturePong.createView() },
{ binding: 5, resource: { buffer: uniformBuffer.buffer } },
],
});
const agentBindGroupPong = device.createBindGroup({
label: 'agent bind group pong',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: agentBufferPong.buffer } },
{ binding: 1, resource: { buffer: agentBufferPing.buffer } },
{ binding: 2, resource: { buffer: trailCountsBuffer.buffer } },
{ binding: 3, resource: trailTexturePong.createView() },
{ binding: 4, resource: trailTexturePing.createView() },
{ binding: 5, resource: { buffer: uniformBuffer.buffer } },
],
});
const pipelineLayout = device.createPipelineLayout({
label: 'pipeline layout',
bindGroupLayouts: [bindGroupLayout],
});
const agentPipeline = device.createComputePipeline({
label: 'agent pipeline',
layout: pipelineLayout,
compute: {
module: device.createShaderModule({ code: agentShader }),
entryPoint: 'cs',
},
});
const trailCombinePipeline = device.createComputePipeline({
label: 'agent pipeline',
layout: pipelineLayout,
compute: {
module: device.createShaderModule({ code: trailCombineShader }),
entryPoint: 'cs',
},
});
const debugViewer = await renderTexture(device, trailTexturePing);
const debugAgents = await drawAgents(device, {
w,
h,
layout: pipelineLayout,
});
let step = 0;
function render() {
uniforms.set([performance.now()], 0);
uniformBuffer.write();
const encoder = device.createCommandEncoder();
const agentPass = encoder.beginComputePass();
agentPass.setPipeline(agentPipeline);
agentPass.setBindGroup(
0,
[agentBindGroupPing, agentBindGroupPong][step % 2],
);
agentPass.dispatchWorkgroups(h, 1, 1);
agentPass.end();
const trailCombinePass = encoder.beginComputePass();
trailCombinePass.setPipeline(trailCombinePipeline);
trailCombinePass.setBindGroup(
0,
[agentBindGroupPing, agentBindGroupPong][step % 2],
);
trailCombinePass.dispatchWorkgroups(h, 1, 1);
trailCombinePass.end();
step++;
debugViewer.render();
debugAgents.render([agentBindGroupPing, agentBindGroupPong][step % 2]);
device.queue.submit([encoder.finish()]);
}
render();
debugViewer.render();
const display = [debugViewer.context.canvas, debugAgents.context.canvas];
display.forEach(d => {
d.style.width = `${d.width * 2}px`
d.style.height = `${d.height * 2}px`
})
return { render, display };
}