async function main() { const adapter = await navigator.gpu?.requestAdapter(); const device = await adapter?.requestDevice(); if (!device) { fail('need a browser that supports WebGPU'); return; } const canvas = document.querySelector('canvas'); const context = canvas.getContext('webgpu'); const presentationFormat = navigator.gpu.getPreferredCanvasFormat(); context.configure({device, format: presentationFormat}); let Nparticles = 2**13; let dt = 0.04; let mass = 1.0/Nparticles; let workgroup_size = 64; let Nworkgroups = Nparticles/workgroup_size; const computeModule = device.createShaderModule({ label: 'nbody compute module', code: ` struct Particle { pos: vec3f, vel: vec3f, }; @group(0) @binding(0) var particles: array; @compute @workgroup_size(${workgroup_size}) fn kick( @builtin(global_invocation_id) id: vec3u, @builtin(local_invocation_index) local_invocation_index: u32, @builtin(num_workgroups) num_workgroups: vec3, @builtin(workgroup_id) workgroup_id : vec3, ) { let pi = workgroup_id.x * ${workgroup_size} + local_invocation_index; // global_invocation_index for (var pj = 0u; pj<${Nparticles}; pj++){ if (pi!=pj){ let rel_pos = particles[pj].pos-particles[pi].pos; let softening = 0.01; let d = sqrt(rel_pos.x*rel_pos.x+rel_pos.y*rel_pos.y+rel_pos.z*rel_pos.z)+softening; particles[pi].vel += ${mass} * ${dt} * rel_pos/(d*d*d) ; } } } @compute @workgroup_size(${workgroup_size}) fn drift( @builtin(global_invocation_id) id: vec3u, @builtin(local_invocation_index) local_invocation_index: u32, @builtin(num_workgroups) num_workgroups: vec3, @builtin(workgroup_id) workgroup_id : vec3, ) { let pi = workgroup_id.x * ${workgroup_size} + local_invocation_index; // global_invocation_index particles[pi].pos += 0.5*${dt} * particles[pi].vel; } `, }); const renderModule = device.createShaderModule({ label: 'nbody render module', code: ` struct Particle { pos: vec3f, vel: vec3f, }; @group(0) @binding(0) var particles: array; @group(0) @binding(1) var aspect: f32; struct InOut { @builtin(position) position: vec4f, @location(0) tex: vec2f, @location(1) vel: f32, }; @vertex fn vs( @builtin(vertex_index) vertexIndex : u32 ) -> InOut { let quad = array( vec2f( -1, -1), vec2f( -1, 1), vec2f( 1, -1), vec2f( 1, 1), vec2f( 1, -1), vec2f( -1, 1) ); var pos = 0.05 * quad[vertexIndex%6].xy + 0.25 * particles[vertexIndex/6].pos.xy; pos.x *= aspect; return InOut(vec4f(pos,0,1), quad[vertexIndex%6], length(particles[vertexIndex/6].vel)); } @fragment fn fs(inOut: InOut) -> @location(0) vec4f { let phi = atan2(inOut.tex.y,inOut.tex.x); var r = 1.0-length(inOut.tex); let a = pow(clamp(sin(phi*6.0),0.0,1.0),5.0/r); let c = inOut.vel/2; return vec4f(1,1-c,1-c,pow(1.1*r,8.0) + r*a); } `, }); const computeBindGroupLayout = device.createBindGroupLayout({ entries: [ { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: 'storage', minBindingSize: 0, }, }, ], }); const computePipelineLayout = device.createPipelineLayout({ bindGroupLayouts: [ computeBindGroupLayout ], }); const kickComputePipeline = device.createComputePipeline({ label: 'kick compute pipeline', layout: computePipelineLayout, compute: { entryPoint: "kick", module: computeModule, }, }); const driftComputePipeline = device.createComputePipeline({ label: 'drift compute pipeline', layout: computePipelineLayout, compute: { entryPoint: "drift", module: computeModule, }, }); const renderPipeline = device.createRenderPipeline({ label: 'render pipeline', layout: "auto", vertex: { module: renderModule, }, fragment: { module: renderModule, targets: [ { format: presentationFormat, blend: { color: { srcFactor: 'src-alpha', dstFactor: 'one-minus-src-alpha' }, alpha: { srcFactor: 'src-alpha', dstFactor: 'one-minus-src-alpha' }, }, } ], }, }); const initialConditions = new Float32Array(8*Nparticles); for (var i = 0; i { for (const entry of entries) { const canvas = entry.target; width = entry.contentBoxSize[0].inlineSize; height = entry.contentBoxSize[0].blockSize; canvas.width = Math.max(1, Math.min(width, device.limits.maxTextureDimension2D)); canvas.height = Math.max(1, Math.min(height, device.limits.maxTextureDimension2D)); render(); } }); observer.observe(canvas); } function fail(msg) { // eslint-disable-next-line no-alert alert(msg); } main();