async function renderable(initialSimUniforms /*: Record*/) { const { device, context, format } = await gpu.init(w, h); const CONFIG_UNIFORMS_ORDER = Object.keys(initialSimUniforms); const module = device.createShaderModule({ label: 'vert & frag shader module', code: /* wgsl */ ` struct Uniforms { size: f32, elapsed: f32, random: f32, padding1: f32, // align next mat4x4 matrix: mat4x4f, // viewProj matrix resolution: vec2f, padding2: vec2f, // struct size multiple of 16 bytes } struct Agent { position: vec4f, direction: vec4f, } struct VertexOut { @builtin(position) position: vec4f, @location(1) norm_index: f32, @location(2) pos: vec3f, @location(3) direction: vec3f, @location(4) quad_pos: vec2f, }; ${ShaderFragment.hsl2rgb} @group(0) @binding(0) var uniforms: Uniforms; @group(0) @binding(1) var particles: array; @vertex fn vs( @builtin(instance_index) instance_index : u32, @builtin(vertex_index) vertex_index : u32, ) -> VertexOut { let quad = 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 position = particles[instance_index].position; let clip_pos = uniforms.matrix * vec4f(position.xyz, 1); let point_pos = vec4f((quad[vertex_index] - 0.5) * (uniforms.size / uniforms.resolution), 0, 0); let out_pos = clip_pos + point_pos; return VertexOut( out_pos, f32(instance_index) / f32(arrayLength(&particles)), position.xyz, particles[instance_index].direction.xyz, quad[vertex_index] - 0.5, ); } @fragment fn fs(vout: VertexOut) -> @location(0) vec4f { if (distance(vout.quad_pos, vec2f(0)) > 0.5) { discard; } // let color = vec4f( // hsl2rgb(vec3f((1.0 - dist) * 0.4 + 0.6, 1.0, 0.7)), // 1.0, // ); let position_color = vec4f( vout.pos.x * 0.5 + 0.5, vout.pos.y * 0.5 + 0.5, vout.pos.z * 0.5 + 0.5, 1.0, ); let direction_color = vec4f( vout.direction.x * 0.5 + 0.5, vout.direction.y * 0.5 + 0.5, vout.direction.z * 0.5 + 0.5, 1.0, ); let dist = distance(vout.pos, vec3f(0)); let color = position_color * direction_color * (1.0 - dist); // let color = position_color * direction_color * (1.0 - dist * 0.5); // let color = vec4f(hsl2rgb(vec3f(dist, 1.0, 0.65)), 1.0) * position_color * dist * 0.75; return color; } `, }); const computeShader = /* wgsl */ ` struct Uniforms { size: f32, elapsed: f32, random: f32, padding1: f32, // align next mat4x4 matrix: mat4x4f, resolution: vec2f, padding2: vec2f, // struct size multiple of 16 bytes } struct SimulationUniforms { ${CONFIG_UNIFORMS_ORDER.map(u => `${u}: f32,`).join('\n')} } struct Agent { position: vec4f, direction: vec4f, } struct Sensed { position: vec3f, direction: vec3f, } @group(0) @binding(0) var uniforms: Uniforms; @group(0) @binding(1) var particles_in: array; @group(0) @binding(2) var particles_out: array; @group(0) @binding(3) var trail_counts: array>; @group(0) @binding(4) var trail_texture_input: texture_3d; @group(0) @binding(6) var sim_uniforms: SimulationUniforms; ${ShaderFragment.hash14} // ported from threejs fn apply_quaternion(q: vec4f, v: vec3f) -> vec3f { // quaternion q is assumed to have unit length let vx = v.x; let vy = v.y; let vz = v.z; let qx = q.x; let qy = q.y; let qz = q.z; let qw = q.w; // t = 2 * cross( q.xyz, v ); let tx = 2 * (qy * vz - qz * vy); let ty = 2 * (qz * vx - qx * vz); let tz = 2 * (qx * vy - qy * vx); // v + q.w * t + cross( q.xyz, t ); let next = vec3f( vx + qw * tx + qy * tz - qz * ty, vy + qw * ty + qz * tx - qx * tz, vz + qw * tz + qx * ty - qy * tx, ); return next; } // the above function is faster than this one (prob bc of cross()) // fn apply_quaternion(q: vec4f, v: vec3f) -> vec3f { // let qv = vec3f(q.x, q.y, q.z); // let t = 2.0 * cross(qv, v); // return v + q.w * t + cross(qv, t); // } fn set_from_axis_angle(axis: vec3f, angle: f32) -> vec4f { let halfAngle = angle * 0.5; let s = sin(halfAngle); return vec4f(axis * s, cos(halfAngle)); } fn apply_axis_angle(axis: vec3f, angle: f32, v: vec3f) -> vec3f { if (angle == 0.0) { return v; } let q = set_from_axis_angle(axis, angle); return apply_quaternion(q, v); } fn get_next_position(center: vec3f, direction: vec3f) -> Sensed { let norm_direction = normalize(direction); let up = vec3f(direction.xyz); let right = normalize(cross(norm_direction, up)); let adjustedUp = cross(right, norm_direction); var theta_offset_index: f32 = 0.0; var phi_offset_index: f32 = 0.0; var selected_value: f32 = 0.0; let theta_offset_range = sim_uniforms.theta_offset_range; let phi_offset_range = sim_uniforms.phi_offset_range; let sensor_distance = sim_uniforms.sensor_distance; for (var off: f32 = 0.0; off < 9.0; off += 1.0) { let i = (off % 3.0) - 1.0; let j = floor(off / 3.0) - 1.0; let theta = theta_offset_range * i; let phi = phi_offset_range * j; let a = apply_axis_angle(right, theta, norm_direction); let b = apply_axis_angle(adjustedUp, phi, a); let local_direction = normalize(b) * sensor_distance; let position = center + local_direction; if (position.x >= -1.0 && position.x <= 1.0 && position.y >= -1.0 && position.y <= 1.0 && position.z >= -1.0 && position.z <= 1.0) { let texcoords = vec3u(floor((position * 0.5 + vec3f(0.5)) * 64.0)); let value = textureLoad(trail_texture_input, texcoords, 0).r; if (value > selected_value || (value == selected_value && hash14(vec4f(position.xyz, uniforms.random)) > 0.5)) { theta_offset_index = i; phi_offset_index = j; selected_value = value; } } } let theta = sim_uniforms.theta_offset_range_agent * theta_offset_index; let phi = sim_uniforms.phi_offset_range_agent * phi_offset_index; let a = apply_axis_angle(right, theta, norm_direction); let b = apply_axis_angle(adjustedUp, phi, a); let local_direction_distance = normalize(b) * sim_uniforms.step_distance; return Sensed(center + local_direction_distance, normalize(local_direction_distance)); } @compute @workgroup_size(${AGENT_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 agent = particles_in[index]; var next = get_next_position(agent.position.xyz, agent.direction.xyz); var next_pos = next.position; var next_dir = next.direction; // collide sphere if (dot(next_pos, next_pos) > 1.0) { let normal = normalize(next_pos); next_pos = normal; next_dir = next_dir - 2.0 * dot(next_dir, normal) * normal; } particles_out[index].position = vec4f(next_pos, 1.0); particles_out[index].direction = vec4f(next_dir, 0.0); let trail_counts_index = u32( floor((next_pos.x * 0.5 + 0.5) * 64.0) + floor((next_pos.y * 0.5 + 0.5) * 64.0) * 64.0 + floor((next_pos.z * 0.5 + 0.5) * 64.0) * (64.0 * 64.0) ); // update trail counts with new agent position atomicAdd(&trail_counts[trail_counts_index], 1); } `; 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' }, }, // // trail counts { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage' }, }, // trail texture IN { binding: 4, visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE, texture: { viewDimension: '3d', }, }, // trail texture OUT (blur + decay + counts) { binding: 5, visibility: GPUShaderStage.COMPUTE, storageTexture: { format: 'rgba8unorm', viewDimension: '3d', }, }, // sim uniforms { binding: 6, visibility: GPUShaderStage.VERTEX | GPUShaderStage.COMPUTE, buffer: {}, }, ], }); /** * app state uniforms */ // prettier-ignore const uniforms = new Float32Array([ initialSimUniforms.particle_size, // size 0, // elapsed (delta) time between frames Math.random(), // random 0, // padding1 ...util.MAT4X4_IDENTITY, // mat4x4 as 16 floats w, h, // resolution 0, 0, // padding2 ]); const uniformBuffer = device.createBuffer({ label: 'uniforms buffer', size: uniforms.byteLength, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(uniformBuffer, 0, uniforms); const matrixValue = uniforms.subarray(4, 4 + 16); /** * sim config uniforms */ const simUniforms = new Float32Array( CONFIG_UNIFORMS_ORDER.map(u => initialSimUniforms[u]), ); const simUniformBuffer = device.createBuffer({ label: 'sim uniforms buffer', size: simUniforms.byteLength, usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, }); device.queue.writeBuffer(simUniformBuffer, 0, simUniforms); /** * trails */ const trailCountsBuffer = gpu.storageBuffer( device, new Uint32Array(S * S * S), ); const dims = { w: S, h: S, d: S }; const { texture: trailTexturePing, cleanup: cleanupTrailTexture } = tex3d( device, dims, ); const trailTexturePong = device.createTexture({ dimension: '3d', size: [S, S, S], format: 'rgba8unorm', usage: GPUTextureUsage.STORAGE_BINDING | GPUTextureUsage.COPY_SRC | GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.RENDER_ATTACHMENT, }); /** * agent buffers setup */ // prettier-ignore const points = new Float32Array( util.flatArr(PARTICLE_COUNT, () => [ // ...randomPointInSphere(), 1.0, // position 0, 0, 0, 1.0, // position util.rands(), util.rands(), util.rands(), // direction 0, // pad ]), ); 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 } }, { binding: 3, resource: { buffer: trailCountsBuffer.buffer } }, { binding: 4, resource: trailTexturePing.createView({ dimension: '3d' }), }, { binding: 5, resource: trailTexturePong.createView({ dimension: '3d' }), }, { binding: 6, resource: { buffer: simUniformBuffer } }, ], }); const bindGroupB = device.createBindGroup({ label: 'bind group layout B', layout: bindGroupLayout, entries: [ { binding: 0, resource: { buffer: uniformBuffer } }, { binding: 1, resource: { buffer: pointsBufferB } }, // swap { binding: 2, resource: { buffer: pointsBufferA } }, { binding: 3, resource: { buffer: trailCountsBuffer.buffer } }, // swap { binding: 4, resource: trailTexturePong.createView({ dimension: '3d' }), }, { binding: 5, resource: trailTexturePing.createView({ dimension: '3d' }), }, { binding: 6, resource: { buffer: simUniformBuffer } }, ], }); 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', // @ts-expect-error targets: [ { format, blend: gpu.Blending.Additive, }, ], }, }); const computePipeline = device.createComputePipeline({ label: 'compute pipeline', layout: pipelineLayout, compute: { module: computeModule, entryPoint: 'cs', }, }); const trailUpdatePipeline = device.createComputePipeline({ label: 'trail update pipeline', layout: pipelineLayout, compute: { module: device.createShaderModule({ code: /* wgsl */ ` struct SimulationUniforms { ${CONFIG_UNIFORMS_ORDER.map(u => `${u}: f32,`).join('\n')} } const BLUR_RADIUS: i32 = 1; const TEX_SIZE: vec3i = vec3i(64); const COUNT: i32 = i32(pow(f32(BLUR_RADIUS * 2 + 1), 3.0)); // const offsets = array( // vec3i(-1, -1, -1), vec3i(0, -1, -1), vec3i(1, -1, -1), // vec3i(-1, 0, -1), vec3i(0, 0, -1), vec3i(1, 0, -1), // vec3i(-1, 1, -1), vec3i(0, 1, -1), vec3i(1, 1, -1), // vec3i(-1, -1, 0), vec3i(0, -1, 0), vec3i(1, -1, 0), // vec3i(-1, 0, 0), vec3i(0, 0, 0), vec3i(1, 0, 0), // vec3i(-1, 1, 0), vec3i(0, 1, 0), vec3i(1, 1, 0), // vec3i(-1, -1, 1), vec3i(0, -1, 1), vec3i(1, -1, 1), // vec3i(-1, 0, 1), vec3i(0, 0, 1), vec3i(1, 0, 1), // vec3i(-1, 1, 1), vec3i(0, 1, 1), vec3i(1, 1, 1) // ); const offsets = array( vec3i( 1, 0, 0), vec3i( 0, 1, 0), vec3i( 0, 0, 1), vec3i(-1, 0, 0), vec3i( 0, -1, 0), vec3i( 0, 0, -1), ); @group(0) @binding(3) var trail_lookup: array>; @group(0) @binding(4) var trail_texture_input: texture_3d; @group(0) @binding(5) var trail_texture_output: texture_storage_3d; @group(0) @binding(6) var sim_uniforms: SimulationUniforms; @compute @workgroup_size(${TRAIL_WORKGROUP_SIZE}) fn cs( @builtin(global_invocation_id) global_id: vec3u, ) { let texcoord = vec3i(global_id.xyz); let index = ( global_id.x + global_id.y * 64 + global_id.z * 64 * 64 ); var colorSum: vec4f = vec4f(0); for (var i = 0; i < 6; i += 1) { let samplePos = texcoord + offsets[i]; if (all(samplePos >= vec3i(0)) && all(samplePos < TEX_SIZE)) { colorSum = colorSum + textureLoad(trail_texture_input, samplePos, 0); } } let blurred: vec4f = colorSum / 6.0; var trail_count = f32(atomicLoad(&trail_lookup[index])); // trail_count = max(trail_count + 1, 5.0); var color = blurred - vec4f(sim_uniforms.decay_rate) + trail_count * sim_uniforms.deposit_amount * sim_uniforms.trail_count_dampener; color = clamp(color, vec4f(0), vec4f(1)); textureStore(trail_texture_output, texcoord, color); // reset atomicStore(&trail_lookup[index], 0); } `, }), entryPoint: 'cs', }, }); let step = 0; let last = 0; let elapsed; // let mouse = [0, 0] // camera const projection = mat4.perspective((90 * Math.PI) / 180, w / h, 0.1, 256); const view = mat4.lookAt( [0, 0, 1.5], // position [0, 0, 0], // target [0, 1, 0], // up ); const viewProjection = mat4.multiply(projection, view); mat4.copy(viewProjection, matrixValue); const renderPassConfig = { colorAttachments: [ { // clearValue: [0, 0, 0, 1], // clearValue: [0, 0, 0.29, 1], clearValue: [0.1, 0.1, 0.1, 1], loadOp: 'clear', storeOp: 'store', view: undefined, }, ], }; const handleZoom = e => { e.preventDefault() mat4.translate( viewProjection, [0, 0, e.wheelDelta * 0.0001], viewProjection, ); }; context.canvas.addEventListener('wheel', handleZoom); function updateSimUniforms(uniformName /*: string */, value /*: any */) { simUniforms.set([value], CONFIG_UNIFORMS_ORDER.indexOf(uniformName)); device.queue.writeBuffer(simUniformBuffer, 0, simUniforms); } let frame /*: number */; function render(_t /*: number */) { requestAnimationFrame(render); let t = performance.now() * 0.001; elapsed = t - last; last = t; uniforms.set([elapsed], 1); uniforms.set([Math.random()], 2); mat4.rotateY(viewProjection, t * 0.2, matrixValue); 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(AGENT_WORKGROUP_DISPATCH_COUNT); computePass.end(); const trailComputePass = encoder.beginComputePass(); trailComputePass.setPipeline(trailUpdatePipeline); // for (let i = 0; i < 2; i++) { trailComputePass.setBindGroup(0, pingpong[step % 2]); trailComputePass.dispatchWorkgroups(...TRAIL_WORKGROUP_DISPATCH_COUNT); // step++; // } trailComputePass.end(); step++; // render passes renderPassConfig.colorAttachments[0].view = context // @ts-expect-error .getCurrentTexture() .createView(); // @ts-expect-error const pass = encoder.beginRenderPass(renderPassConfig); pass.setPipeline(pipeline); pass.setBindGroup(0, pingpong[step % 2]); pass.draw(6, PARTICLE_COUNT); pass.end(); device.queue.submit([encoder.finish()]); } return { context, animate: () => requestAnimationFrame(render), stop: () => cancelAnimationFrame(frame), cleanup: () => { cancelAnimationFrame(frame); uniformBuffer.destroy(); trailCountsBuffer.buffer.destroy(); pointsBufferA.destroy(); pointsBufferB.destroy(); cleanupTrailTexture(); trailTexturePong.destroy(); device.destroy(); }, updateSimUniforms, } }