Public
Edited
Jun 5, 2024
Importers
1 star
renderable = {
const [w, h] = [640, 640]
const { device, context, format } = await gpu.init(w, h)
const PARTICLE_COUNT = 100_000
const PARTICLE_SIZE = 0.01
const WORKGROUP_SIZE = 64
const WORKGROUP_DISPATCH_COUNT = Math.ceil(PARTICLE_COUNT / WORKGROUP_SIZE)
const module = device.createShaderModule({
label: 'vert & frag shader module',
code: `
${shaderFragments.hsl2rgb}
struct Uniforms {
size: f32,
count: f32,
mouse: vec2f,
pad: f32,
elapsed: f32,
}
struct Particle {
position: vec2f,
velocity: vec2f,
}
struct VertexOut {
@builtin(position) position: vec4f,
@location(1) norm_index: f32,
@location(2) pos: vec2f,
@location(3) vel: vec2f,
};
@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<storage, read> particles: array<Particle>;
@vertex
fn vs(
@builtin(instance_index) instance_index : u32,
@builtin(vertex_index) vertex_index : u32,
) -> VertexOut {
let p = array(
vec2f(0.0, 0.0),
vec2f(1.0, 0.0),
vec2f(0.0, 1.0),
vec2f(0.0, 1.0),
vec2f(1.0, 0.0),
vec2f(1.0, 1.0),
);
let pos = particles[instance_index].position;
let vel = particles[instance_index].velocity.xy;
let xy = (p[vertex_index] - 0.5) * uniforms.size + pos;
return VertexOut(
vec4f(xy, 1, 1),
f32(instance_index) / f32(arrayLength(&particles)),
pos,
vel
);
}
const vec2_1 = vec2f(1);
const vec2_neg1 = vec2f(-1);
@fragment
fn fs(vout: VertexOut) -> @location(0) vec4f {
// return vec4f(hsl2rgb(vec3f(vout.norm_index, 1, 0.65)), 1) * 0.9;
// return vec4f((vout.vel + vec2_1) / 2, vout.norm_index, 1) * 0.5;
let vel = (vout.vel + vec2_1) / 2;
// return vec4f(vel.x, vout.norm_index, vel.y, 1) * ((vel.x + vel.y) / 3);
return vec4f(vel.x, vout.norm_index, vel.y, 1) * 0.5;
}
`,
});
const computeShader = `
struct Uniforms {
size: f32,
count: f32,
mouse: vec2f,
pad: f32,
elapsed: f32,
}
struct Particle {
position: vec2f,
velocity: vec2f,
}
@group(0) @binding(0) var<uniform> uniforms: Uniforms;
@group(0) @binding(1) var<storage> particles_in: array<Particle>;
@group(0) @binding(2) var<storage, read_write> particles_out: array<Particle>;
@compute
@workgroup_size(${WORKGROUP_SIZE})
fn cs(@builtin(global_invocation_id) global_invocation_id: vec3u) {
let index = global_invocation_id.x;
if (index > arrayLength(&particles_in)) {
return;
}
// let elapsed = 0.0015;
// let elapsed = min(0.0015, uniforms.elapsed * 0.00005);
let elapsed = uniforms.elapsed * 0.0001;
var pos = particles_in[index].position.xy;
var vel = particles_in[index].velocity.xy;
// mouse attraction
let direction = uniforms.mouse - pos;
// use this instead for mouse repulsion!
// let direction = pos - uniforms.mouse;
let distSquared = dot(direction, direction);
let radiusSquared = 0.25;
let magnitude = 500 * (1 - distSquared / radiusSquared);
var force = step(distSquared, radiusSquared) * magnitude * normalize(direction);
// gravity
force += vec2f(0.0, -1.5);
// accelerate
vel += elapsed * force;
// bounce off the sides
if (abs(pos.x) > 1) {
vel.x *= -1;
pos.x = clamp(pos.x, -1, 1);
}
if (abs(pos.y) > 1) {
vel.y *= -1;
pos.y = clamp(pos.y, -1, 1);
}
// damping
vel *= 0.995;
// move
pos += elapsed * vel;
// particles_out[index].position = clamp(pos, vec2f(-1, -1), vec2f(1, 1));
particles_out[index].position = pos;
particles_out[index].velocity = vel;
}
`
const computeModule = device.createShaderModule({
label: 'compute shader module',
code: computeShader
})
// define access to resources across all pipelines
const bindGroupLayout = device.createBindGroupLayout({
label: 'bind group layout',
entries: [
// uniforms
{
binding: 0,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
buffer: {}
},
// particles in
{
binding: 1,
visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE,
buffer: { type: 'read-only-storage' }
},
// particles out
{
binding: 2,
visibility: GPUShaderStage.COMPUTE,
buffer: { type: 'storage' }
},
]
})
// note about alignment...
// (https://surma.dev/things/webgpu/ ctrl-f "alignment")
const uniforms = new Float32Array([
// this works because the vec2f mouse coords must be aligned to a memory
// address of 8; particle size (f32, size=4, align=4) + particle count (f32) = 8
// which allows us to align the size=8 vec2<f32> align=8 to a memory address
// multiple of 8
PARTICLE_SIZE, // 4
PARTICLE_COUNT, // + 4
0, 0, // = address 8
//
0, // pad
0, // elapsed (delta) time between frames
])
const uniformBuffer = device.createBuffer({
label: 'uniforms buffer',
size: uniforms.byteLength,
usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
})
device.queue.writeBuffer(uniformBuffer, 0, uniforms);
// position x, position y, velocity x, velocity y
const points = new Float32Array(
util.arr(PARTICLE_COUNT, () => [
util.rands(), util.rands(), util.rands(), util.rands()
]).flat()
)
const pointsBufferA = device.createBuffer({
label: 'points storage buffer A',
size: points.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
})
device.queue.writeBuffer(pointsBufferA, 0, points);
const pointsBufferB = device.createBuffer({
label: 'points storage buffer B',
size: points.byteLength,
usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST
})
device.queue.writeBuffer(pointsBufferB, 0, points);
const bindGroupA = device.createBindGroup({
label: 'bind group layout A',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: uniformBuffer }},
{ binding: 1, resource: { buffer: pointsBufferA }},
{ binding: 2, resource: { buffer: pointsBufferB }},
]
});
const bindGroupB = device.createBindGroup({
label: 'bind group layout B',
layout: bindGroupLayout,
entries: [
{ binding: 0, resource: { buffer: uniformBuffer }},
{ binding: 1, resource: { buffer: pointsBufferB }},
{ binding: 2, resource: { buffer: pointsBufferA }},
]
});
const pingPong = [bindGroupA, bindGroupB]
const pipelineLayout = device.createPipelineLayout({
label: 'pipeline layout',
bindGroupLayouts: [bindGroupLayout]
})
const pipeline = device.createRenderPipeline({
label: 'pipeline',
layout: pipelineLayout,
vertex: {
module,
entryPoint: 'vs',
},
fragment: {
module,
entryPoint: 'fs',
targets: [{
format,
blend: {
color: {
operation: 'add',
srcFactor: 'one',
dstFactor: 'one-minus-src-alpha'
},
alpha: {
operation: 'add',
srcFactor: 'one',
dstFactor: 'one-minus-src-alpha'
},
},
}],
},
})
const computePipeline = device.createComputePipeline({
label: 'compute pipeline',
layout: pipelineLayout,
compute: {
module: computeModule,
entryPoint: 'cs'
}
})
let step = 0;
let last = 0;
let elapsed;
let mouse = [0, 0]
function render() {
let t = performance.now()
elapsed = t - last
last = t
uniforms.set(mouse, 2);
uniforms.set([elapsed], 5);
device.queue.writeBuffer(uniformBuffer, 0, uniforms);
const encoder = device.createCommandEncoder()
const computePass = encoder.beginComputePass()
computePass.setPipeline(computePipeline)
computePass.setBindGroup(0, pingPong[step % 2])
computePass.dispatchWorkgroups(WORKGROUP_DISPATCH_COUNT)
computePass.end()
step++;
const pass = encoder.beginRenderPass({
colorAttachments: [
{
clearValue: [0, 0, 0, 1],
loadOp: 'clear',
storeOp: 'store',
view: context.getCurrentTexture().createView()
},
],
})
pass.setPipeline(pipeline);
pass.setBindGroup(0, pingPong[step % 2]);
pass.draw(6, PARTICLE_COUNT);
pass.end();
device.queue.submit([encoder.finish()]);
}
context.canvas.addEventListener('pointermove', e => {
const rect = context.canvas.getBoundingClientRect();
const x = e.clientX - rect.left;
const y = e.clientY - rect.top;
mouse[0] = (x / w - 0.5) * 2
mouse[1] = (1 - y / h - 0.5) * 2
})
return { context, render }
}
shaderFragments = ({
hsl2rgb: `
fn hsl2rgb(hsl: vec3f) -> vec3f {
let c = vec3f(fract(hsl.x), clamp(hsl.yz, vec2f(0), vec2f(1)));
let rgb = clamp(abs((c.x * 6.0 + vec3f(0.0, 4.0, 2.0)) % 6.0 - 3.0) - 1.0, vec3f(0), vec3f(1));
return c.z + c.y * (rgb - 0.5) * (1.0 - abs(2.0 * c.z - 1.0));
}
`
})
gpu = ({
init: async (width = 512, height = 512) => {
const canvas = document.createElement('canvas');
canvas.width = width;
canvas.height = height;
const context = canvas.getContext('webgpu');
const adapter = await navigator.gpu.requestAdapter();
const device = await adapter.requestDevice();
const format = navigator.gpu.getPreferredCanvasFormat();
context.configure({ device, format, alphaMode: 'premultiplied', });
return { context, adapter, device, format }
},
})
util = ({
rand: (max = 1) => Math.random() * max,
// random signed
rands: () => (Math.random() - 0.5) * 2,
arr: (size, callback) => {
const arr = new Array(size)
if (typeof callback !== 'function') {
return arr.fill(callback)
}
return arr.fill(null).map((_, i) => callback(i))
},
shuffle: arr => {
return arr.sort((a, b) => Math.random() - 0.5)
},
// from karpathy:
randn: (mean, variance) => {
let V1, V2, S;
do {
const U1 = Math.random();
const U2 = Math.random();
V1 = 2 * U1 - 1;
V2 = 2 * U2 - 1;
S = V1 * V1 + V2 * V2;
} while (S > 1);
let X = Math.sqrt(-2 * Math.log(S) / S) * V1;
X = mean + Math.sqrt(variance) * X;
return X;
},
})
Purpose-built for displays of data
Observable is your go-to platform for exploring data and creating expressive data visualizations. Use reactive JavaScript notebooks for prototyping and a collaborative canvas for visual data exploration and dashboard creation.
