From 1f26ab7ef0586cdc413e38a9300a2525498e3933 Mon Sep 17 00:00:00 2001 From: Forrest Sun Date: Sat, 4 May 2024 12:39:35 -0700 Subject: [PATCH] add webgpu node based vfx renderer demo --- examples/index.html | 14 +- examples/nodeBasedVFXDemo.js | 3 +- examples/webgpu.html | 72 +++++ examples/webgpu.js | 182 +++++++++++++ rollup.config.js | 14 + src/WebGPUDriver.ts | 8 - src/WebGPURenderer.ts | 418 ++++++++++++++++++++++++++++++ src/index.ts | 1 + src/nodes/NodeValueType.ts | 42 +++ src/nodes/WebGPUCompiler.ts | 70 ++++- src/nodes/index.ts | 1 + src/shaders/wgsl/particle.wgsl.ts | 60 +++++ types/index.d.ts | 2 +- 13 files changed, 854 insertions(+), 33 deletions(-) create mode 100644 examples/webgpu.html create mode 100644 examples/webgpu.js delete mode 100644 src/WebGPUDriver.ts create mode 100644 src/WebGPURenderer.ts create mode 100644 src/shaders/wgsl/particle.wgsl.ts diff --git a/examples/index.html b/examples/index.html index 8181cf8..91d7766 100644 --- a/examples/index.html +++ b/examples/index.html @@ -2,7 +2,7 @@ - three.quarks – particle system library for three.js + three.quarks – particle system / VFX library for three.js @@ -13,13 +13,7 @@ Three.quarks Particle System Examples - - - + + + + + + + \ No newline at end of file diff --git a/examples/webgpu.js b/examples/webgpu.js new file mode 100644 index 0000000..783df88 --- /dev/null +++ b/examples/webgpu.js @@ -0,0 +1,182 @@ +import {WebGPURenderer, initWebGPU, NodeGraph, Node, NodeTypes, NodeValueType, Wire, WebGPUCompiler} from "three.quarks"; +import {PerspectiveCamera, Vector3} from "three"; + + +const renderCode = ` +//////////////////////////////////////////////////////////////////////////////// +// Vertex shader +//////////////////////////////////////////////////////////////////////////////// +struct RenderParams { + modelViewMatrix : mat4x4f, + projectionMatrix : mat4x4f, +} + +struct Particle { + color: vec4, + position: vec3, + velocity: vec3, + size: f32, + rotation: f32, + life: f32, + age: f32, +} + +@group(0) @binding(0) var render_params : RenderParams; +@group(0) @binding(1) var particles: array; +@group(0) @binding(2) var particleIndices: array; + +struct VertexInput { + /*@location(0) position : vec3f, + @location(1) color : vec4f,*/ + @builtin(instance_index) instanceIndex: u32, + @location(0) quad_pos : vec2f, // -1..+1 +} + +struct VertexOutput { + @builtin(position) position : vec4f, + @location(0) color : vec4f, + @location(1) quad_pos : vec2f, // -1..+1 +} + +@vertex +fn vs_main(in : VertexInput) -> VertexOutput { + var particle = particles[particleIndices[in.instanceIndex]]; + var alignedPosition = ( in.quad_pos.xy ) * particle.size; + var rotatedPosition = vec2f( + cos( particle.rotation ) * alignedPosition.x - sin( particle.rotation ) * alignedPosition.y, + sin( particle.rotation ) * alignedPosition.x + cos( particle.rotation ) * alignedPosition.y, + ); + + var mvPosition = render_params.modelViewMatrix * vec4f(particle.position, 1.0); + + var out : VertexOutput; + out.position = render_params.projectionMatrix * vec4(mvPosition.xy + rotatedPosition, mvPosition.zw); + out.color = particle.color; + out.quad_pos = in.quad_pos; + return out; +} + +//////////////////////////////////////////////////////////////////////////////// +// Fragment shader +//////////////////////////////////////////////////////////////////////////////// +@fragment +fn fs_main(in : VertexOutput) -> @location(0) vec4f { + var color = in.color; + // Apply a circular particle alpha mask + // color.a = color.a * max(1.0 - length(in.quad_pos), 0.0); + return color; +} + +`; +initWebGPU().then((context) => { + const graph = new NodeGraph('test'); + + const sizeVal = new Node(NodeTypes['number'], 0); + const size = new Node(NodeTypes['particleProperty'], 0, {property: 'size', type: NodeValueType.Number}); + sizeVal.inputs[0] = {getValue: () => 1}; + graph.addNode(sizeVal); + graph.addNode(size); + graph.addWire(new Wire(sizeVal, 0, size, 0)); + + const rotationValue = new Node(NodeTypes['number'], 0); + const rotation = new Node(NodeTypes['particleProperty'], 0, {property: 'rotation', type: NodeValueType.Number}); + rotationValue.inputs[0] = {getValue: () => 1}; + graph.addNode(rotationValue); + graph.addNode(rotation); + graph.addWire(new Wire(rotationValue, 0, rotation, 0)); + + const colorVec = new Node(NodeTypes['vec4'], 0); + const color = new Node(NodeTypes['particleProperty'], 0, {property: 'color', type: NodeValueType.Vec4}); + colorVec.inputs[0] = {getValue: () => 1}; + colorVec.inputs[1] = {getValue: () => 1}; + colorVec.inputs[2] = {getValue: () => 1}; + colorVec.inputs[3] = {getValue: () => 0.5}; + graph.addNode(colorVec); + graph.addNode(color); + graph.addWire(new Wire(colorVec, 0, color, 0)); + + const age = new Node(NodeTypes['particleProperty'], 0, {property: 'age', type: NodeValueType.Number}); + graph.addNode(age); + const cos = new Node(NodeTypes['cos'], 0); + graph.addNode(cos); + const pos = new Node(NodeTypes['vec3'], 0); + pos.inputs[0] = {getValue: () => 0}; + pos.inputs[1] = {getValue: () => 1}; + pos.inputs[2] = {getValue: () => 1}; + graph.addNode(pos); + graph.addWire(new Wire(age, 0, cos, 0)); + graph.addWire(new Wire(cos, 0, pos, 0)); + + const pos2 = new Node(NodeTypes['vec3'], 0); + pos2.inputs[0] = {getValue: () => 1}; + pos2.inputs[1] = {getValue: () => 0.5}; + pos2.inputs[2] = {getValue: () => 0.5}; + graph.addNode(pos2); + const add = new Node(NodeTypes['add'], 2); + graph.addNode(add); + graph.addWire(new Wire(pos, 0, add, 0)); + graph.addWire(new Wire(pos2, 0, add, 1)); + + const ppos = new Node(NodeTypes['particleProperty'], 0, {property: 'position', type: NodeValueType.Vec3}); + graph.addNode(ppos); + const pvel = new Node(NodeTypes['particleProperty'], 0, {property: 'velocity', type: NodeValueType.Vec3}); + graph.addNode(pvel); + graph.addWire(new Wire(add, 0, ppos, 0)); + graph.addWire(new Wire(pos2, 0, pvel, 0)); + + const compiler = new WebGPUCompiler(); + const particle = { position: new Vector3(), velocity: new Vector3(), age: 10 }; + const graphContext = { particle: particle }; + + const code = compiler.build(graph, graphContext); + console.log(code); + console.log(compiler.particleInstanceByteSize); + + const debug = false; + const count = 64; + const canvas = document.getElementById('renderer-canvas'); + const renderer = new WebGPURenderer( + context, + canvas, + count, + compiler.particleInstanceByteSize, + code, + debug, + renderCode + ); + const camera = new PerspectiveCamera( + (2 * Math.PI) / 5 * 180 / Math.PI, + canvas.width / canvas.height, + 1, + 100.0 + ); + camera.position.z = 10; + camera.rotateX(Math.PI * 0.05); + camera.updateMatrixWorld(true); + camera.updateProjectionMatrix(); + console.log(camera.matrixWorldInverse.elements); + console.log(camera.projectionMatrix.elements); + console.log(renderer.mvp.multiplyMatrices(camera.projectionMatrix, camera.matrixWorldInverse).elements); + + let frameCount = 0; + function animate() { + renderer.frame(camera); + + if (frameCount % 20 === 0 && debug) { + renderer.cpuReadableBuffer.mapAsync( + GPUMapMode.READ, + 0, // Offset + count * compiler.particleInstanceByteSize, // Length + ).then(() => { + const copyArrayBuffer = renderer.cpuReadableBuffer.getMappedRange(0, count * compiler.particleInstanceByteSize); + const data = copyArrayBuffer.slice(0); + renderer.cpuReadableBuffer.unmap(); + console.log(new Float32Array(data)) + }); + } + frameCount ++; + requestAnimationFrame(animate); + }; + requestAnimationFrame(animate); + //document.body.appendChild(renderer.domElement); +}); \ No newline at end of file diff --git a/rollup.config.js b/rollup.config.js index 94aafa1..dd08d34 100644 --- a/rollup.config.js +++ b/rollup.config.js @@ -16,11 +16,25 @@ const production = process.env.NODE_ENV === 'production'; const globals = {three: 'THREE'}; const extensions = ['.js', '.jsx', '.ts', '.tsx']; +function wgslPlugin() { + return { + name: 'wgsl-plugin', + transform(code, id) { + if (id.endsWith('.wgsl')) { + return { + code: `export default \`${code}\`;`, + map: { mappings: '' }, + }; + } + }, + }; +} export const lib = { main: { input: 'src/index.ts', external: Object.keys(globals), plugins: [ + //wgslPlugin(), resolve({ extensions: extensions, customResolveOptions: { diff --git a/src/WebGPUDriver.ts b/src/WebGPUDriver.ts deleted file mode 100644 index ee137a2..0000000 --- a/src/WebGPUDriver.ts +++ /dev/null @@ -1,8 +0,0 @@ -async function initialize() { - const adapter = await navigator.gpu?.requestAdapter(); - const device = await adapter?.requestDevice(); - if (!device) { - fail('need a browser that supports WebGPU'); - return; - } -} diff --git a/src/WebGPURenderer.ts b/src/WebGPURenderer.ts new file mode 100644 index 0000000..8e75129 --- /dev/null +++ b/src/WebGPURenderer.ts @@ -0,0 +1,418 @@ +import {ParticleSystem} from "./ParticleSystem"; +import {Camera, Matrix4, PerspectiveCamera} from "three"; +import particleWGSL from "./shaders/wgsl/particle.wgsl"; + +interface WebGPUDeviceContext { + adapter: GPUAdapter; + device: GPUDevice; +} + +export async function initWebGPU(): Promise { + const adapter = await navigator.gpu?.requestAdapter(); + if (! + adapter + ) { + throw 'need a browser that supports WebGPU'; + } + const device = await adapter?.requestDevice(); + if (!device) { + throw 'need a browser that supports WebGPU'; + } + return { + adapter, + device, + } +} + +export class WebGPURenderer { + + // simulation + private numLiveParticles: number = 0; + simulationUBOBuffer!: GPUBuffer; + simulationParams!: { + simulate: boolean, + deltaTime: number, + } + private particlesBuffer!: GPUBuffer; + private particleIndexBuffer!: GPUBuffer; + private computeBindGroup!: GPUBindGroup; + private computePipeline!: GPUComputePipeline; + + // rendering + private context!: GPUCanvasContext; + private mvp: Matrix4 = new Matrix4(); + private uniformBuffer!: GPUBuffer; + private quadVertexBuffer!: GPUBuffer; + private renderBindGroup!: GPUBindGroup; + private renderPassDescriptor!: GPURenderPassDescriptor; + private renderPipeline!: GPURenderPipeline; + + // cpu staging buffer + public cpuReadableBuffer!: GPUBuffer; + + + constructor( + public deviceContext: WebGPUDeviceContext, + public canvas: HTMLCanvasElement, + public numParticles: number, + public particleInstanceByteSize: number, + code: string, + public debug: boolean = false, + public renderCode: string = '') { + this.initBuffers(); + this.initRenderPipeline(); + this.initSimulationPipeline(code); + } + + initBuffers() { + + this.numLiveParticles = this.numParticles; + this.particlesBuffer = this.deviceContext.device.createBuffer({ + size: this.numParticles * this.particleInstanceByteSize, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC, + }); + + this.particleIndexBuffer = this.deviceContext.device.createBuffer({ + size: this.numParticles * 4, + usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST, + }) + + if (this.debug) { + this.cpuReadableBuffer = this.deviceContext.device.createBuffer({ + size: this.numParticles * this.particleInstanceByteSize, + usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST, + }); + } + + const particleIndexArr = new Uint32Array(this.numParticles); + for (let i = 0; i < this.numParticles; i++) { + particleIndexArr[i] = i; + } + if(this.debug) { + console.log(particleIndexArr); + } + this.deviceContext.device.queue.writeBuffer(this.particleIndexBuffer, 0, particleIndexArr); + } + + initRenderPipeline() { + const particlePositionOffset = 0; + const particleColorOffset = 4 * 4; + const devicePixelRatio = window.devicePixelRatio; + this.canvas.width = this.canvas.clientWidth * devicePixelRatio; + this.canvas.height = this.canvas.clientHeight * devicePixelRatio; + this.context = this.canvas.getContext('webgpu') as GPUCanvasContext; + const presentationFormat = navigator.gpu.getPreferredCanvasFormat(); + this.context.configure({ + device: this.deviceContext.device, + format: presentationFormat, + alphaMode: 'premultiplied', + }); + const shaderModule = this.deviceContext.device.createShaderModule({ + code: this.renderCode === '' ? particleWGSL : this.renderCode, + }); + this.renderPipeline = this.deviceContext.device.createRenderPipeline({ + layout: 'auto', + vertex: { + module: shaderModule, + entryPoint: 'vs_main', + buffers: [ + /*{ + // instanced particles buffer + arrayStride: this.particleInstanceByteSize, + stepMode: 'instance', + attributes: [ + { + // position + shaderLocation: 0, + offset: particlePositionOffset, + format: 'float32x3', + }, + { + // color + shaderLocation: 1, + offset: particleColorOffset, + format: 'float32x4', + }, + ], + },*/ + { + // quad vertex buffer + arrayStride: 2 * 4, // vec2f + stepMode: 'vertex', + attributes: [ + { + // vertex positions + shaderLocation: 0, + offset: 0, + format: 'float32x2', + }, + ], + }, + ], + }, + fragment: { + module: shaderModule, + entryPoint: 'fs_main', + targets: [ + { + format: presentationFormat, + blend: { + color: { + srcFactor: 'src-alpha', + dstFactor: 'one', + operation: 'add', + }, + alpha: { + srcFactor: 'zero', + dstFactor: 'one', + operation: 'add', + }, + }, + }, + ], + }, + primitive: { + topology: 'triangle-list', + }, + + depthStencil: { + depthWriteEnabled: false, + depthCompare: 'less', + format: 'depth24plus', + }, + }); + + const depthTexture = this.deviceContext.device.createTexture({ + size: [this.canvas.width, this.canvas.height], + format: 'depth24plus', + usage: GPUTextureUsage.RENDER_ATTACHMENT, + }); + + const uniformBufferSize = + 4 * 4 * 4 + // modelViewMatrix : mat4x4f + 4 * 4 * 4 // ProjectionMatrix : mat4x4f; + this.uniformBuffer = this.deviceContext.device.createBuffer({ + size: uniformBufferSize, + usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + }); + + this.renderBindGroup = this.deviceContext.device.createBindGroup({ + layout: this.renderPipeline.getBindGroupLayout(0), + entries: [ + { + binding: 0, + resource: { + buffer: this.uniformBuffer, + }, + }, + { + binding: 1, + resource: { + buffer: this.particlesBuffer, + }, + }, + { + binding: 2, + resource: { + buffer: this.particleIndexBuffer, + }, + }, + ], + }); + + this.renderPassDescriptor = { + colorAttachments: [ + { + view: undefined as any, // Assigned later + clearValue: [0, 0, 0, 1], + loadOp: 'clear', + storeOp: 'store', + }, + ], + depthStencilAttachment: { + view: depthTexture.createView(), + + depthClearValue: 1.0, + depthLoadOp: 'clear', + depthStoreOp: 'store', + }, + }; + + ////////////////////////////////////////////////////////////////////////////// + // Quad vertex buffer + ////////////////////////////////////////////////////////////////////////////// + this.quadVertexBuffer = this.deviceContext.device.createBuffer({ + size: 6 * 2 * 4, // 6x vec2f + usage: GPUBufferUsage.VERTEX, + mappedAtCreation: true, + }); + // prettier-ignore + const vertexData = [ + -1.0, -1.0, +1.0, -1.0, -1.0, +1.0, -1.0, +1.0, +1.0, -1.0, +1.0, +1.0, + ]; + new Float32Array(this.quadVertexBuffer.getMappedRange()).set(vertexData); + this.quadVertexBuffer.unmap(); + + + ////////////////////////////////////////////////////////////////////////////// + // Texture + ////////////////////////////////////////////////////////////////////////////// + /*let texture: GPUTexture; + let textureWidth = 1; + let textureHeight = 1; + let numMipLevels = 1; + { + const response = await fetch('../../assets/img/webgpu.png'); + const imageBitmap = await createImageBitmap(await response.blob()); + + // Calculate number of mip levels required to generate the probability map + while ( + textureWidth < imageBitmap.width || + textureHeight < imageBitmap.height + ) { + textureWidth *= 2; + textureHeight *= 2; + numMipLevels++; + } + texture = this.deviceContext.device.createTexture({ + size: [imageBitmap.width, imageBitmap.height, 1], + mipLevelCount: numMipLevels, + format: 'rgba8unorm', + usage: + GPUTextureUsage.TEXTURE_BINDING | + GPUTextureUsage.STORAGE_BINDING | + GPUTextureUsage.COPY_DST | + GPUTextureUsage.RENDER_ATTACHMENT, + }); + this.deviceContext.device.queue.copyExternalImageToTexture( + {source: imageBitmap}, + {texture: texture}, + [imageBitmap.width, imageBitmap.height] + ); + }*/ + } + + + private initSimulationPipeline(code: string) { + + const simulationUBOBufferSize = + 1 * 4 + // deltaTime + 3 * 4 + // padding + 4 * 4 + // seed + 0; + this.simulationUBOBuffer = this.deviceContext.device.createBuffer({ + size: simulationUBOBufferSize, + usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST, + }); + this.simulationParams = { + simulate: true, + deltaTime: 0.01, + }; + this.computePipeline = this.deviceContext.device.createComputePipeline({ + layout: 'auto', + compute: { + module: this.deviceContext.device.createShaderModule({ + code: code, + }), + entryPoint: 'simulate', + }, + }); + this.computeBindGroup = this.deviceContext.device.createBindGroup({ + layout: this.computePipeline.getBindGroupLayout(0), + entries: [ + { + binding: 0, + resource: { + buffer: this.simulationUBOBuffer, + }, + }, + { + binding: 1, + resource: { + buffer: this.particlesBuffer, + offset: 0, + size: this.numParticles * this.particleInstanceByteSize, + }, + }, + { + binding: 2, + resource: { + buffer: this.particleIndexBuffer, + offset: 0, + size: this.numParticles * 4, + }, + }, + /*{ + binding: 2, + resource: texture.createView(), + },*/ + ], + }); + } + + frame = async (camera: Camera) => { + + this.deviceContext.device.queue.writeBuffer( + this.simulationUBOBuffer, + 0, + new Float32Array([ + this.simulationParams.simulate ? this.simulationParams.deltaTime : 0.0, + 0.0, + 0.0, + 0.0, // padding + Math.random() * 100, + Math.random() * 100, // seed.xy + 1 + Math.random(), + 1 + Math.random(), // seed.zw + ]) + ); + + this.mvp.multiplyMatrices(camera.projectionMatrix, camera.matrixWorldInverse); + // prettier-ignore + this.deviceContext.device.queue.writeBuffer( + this.uniformBuffer, + 0, + new Float32Array([ + // modelViewProjectionMatrix + camera.matrixWorldInverse.elements[0], camera.matrixWorldInverse.elements[1], camera.matrixWorldInverse.elements[2], camera.matrixWorldInverse.elements[3], + camera.matrixWorldInverse.elements[4], camera.matrixWorldInverse.elements[5], camera.matrixWorldInverse.elements[6], camera.matrixWorldInverse.elements[7], + camera.matrixWorldInverse.elements[8], camera.matrixWorldInverse.elements[9], camera.matrixWorldInverse.elements[10], camera.matrixWorldInverse.elements[11], + camera.matrixWorldInverse.elements[12], camera.matrixWorldInverse.elements[13], camera.matrixWorldInverse.elements[14], camera.matrixWorldInverse.elements[15], + camera.projectionMatrix.elements[0], camera.projectionMatrix.elements[1], camera.projectionMatrix.elements[2], camera.projectionMatrix.elements[3], + camera.projectionMatrix.elements[4], camera.projectionMatrix.elements[5], camera.projectionMatrix.elements[6], camera.projectionMatrix.elements[7], + camera.projectionMatrix.elements[8], camera.projectionMatrix.elements[9], camera.projectionMatrix.elements[10], camera.projectionMatrix.elements[11], + camera.projectionMatrix.elements[12], camera.projectionMatrix.elements[13], camera.projectionMatrix.elements[14], camera.projectionMatrix.elements[15], + ]) + ); + + const swapChainTexture = this.context.getCurrentTexture(); + // prettier-ignore + this.renderPassDescriptor.colorAttachments[0]!.view = swapChainTexture.createView(); + + const commandEncoder = this.deviceContext.device.createCommandEncoder(); + { + const passEncoder = commandEncoder.beginComputePass(); + passEncoder.setPipeline(this.computePipeline); + passEncoder.setBindGroup(0, this.computeBindGroup); + passEncoder.dispatchWorkgroups(Math.ceil(this.numParticles / 64)); + passEncoder.end(); + } + { + const passEncoder = commandEncoder.beginRenderPass(this.renderPassDescriptor); + passEncoder.setPipeline(this.renderPipeline); + passEncoder.setBindGroup(0, this.renderBindGroup); + //passEncoder.set + //passEncoder.setVertexBuffer(0, this.particlesBuffer); + passEncoder.setVertexBuffer(0, this.quadVertexBuffer); + passEncoder.draw(6, this.numLiveParticles, 0, 0); + + + passEncoder.end(); + } + if (this.debug) { + commandEncoder.copyBufferToBuffer(this.particlesBuffer, 0, this.cpuReadableBuffer, 0, this.numParticles * this.particleInstanceByteSize); + } + + this.deviceContext.device.queue.submit([commandEncoder.finish()]); + } +} \ No newline at end of file diff --git a/src/index.ts b/src/index.ts index b55dd34..eaa7aec 100644 --- a/src/index.ts +++ b/src/index.ts @@ -6,6 +6,7 @@ export * from './SpriteBatch'; export * from './TrailBatch'; export * from './BatchedRenderer'; export * from './BatchedParticleRenderer'; +export * from './WebGPURenderer'; export * from './QuarksLoader'; export * from './TypeUtil'; export * from './Plugin'; diff --git a/src/nodes/NodeValueType.ts b/src/nodes/NodeValueType.ts index 43dfcb0..ac61e92 100644 --- a/src/nodes/NodeValueType.ts +++ b/src/nodes/NodeValueType.ts @@ -11,6 +11,48 @@ export enum NodeValueType { EventStream = 7, } +export const getAlignOfNodeValueType = (type: NodeValueType): number => { + switch (type) { + case NodeValueType.Boolean: + return 4; + case NodeValueType.Number: + return 4; + case NodeValueType.Vec2: + return 8; + case NodeValueType.Vec3: + return 16; + case NodeValueType.Vec4: + return 16; + case NodeValueType.AnyType: + return 0; + case NodeValueType.NullableAnyType: + return 0; + default: + return 0; + } +}; + +export const getSizeOfNodeValueType = (type: NodeValueType): number => { + switch (type) { + case NodeValueType.Boolean: + return 1; + case NodeValueType.Number: + return 4; + case NodeValueType.Vec2: + return 8; + case NodeValueType.Vec3: + return 12; + case NodeValueType.Vec4: + return 16; + case NodeValueType.AnyType: + return 0; + case NodeValueType.NullableAnyType: + return 0; + default: + return 0; + } +}; + export const genDefaultForNodeValueType = (type: NodeValueType): any => { switch (type) { case NodeValueType.Boolean: diff --git a/src/nodes/WebGPUCompiler.ts b/src/nodes/WebGPUCompiler.ts index e0c5d06..eae5fd8 100644 --- a/src/nodes/WebGPUCompiler.ts +++ b/src/nodes/WebGPUCompiler.ts @@ -3,7 +3,7 @@ import {ExecutionContext, NodeType} from './NodeDef'; import {NodeGraph} from './NodeGraph'; import {Adapter, ConstInput, Node, Wire} from './Node'; import {Vector2, Vector3, Vector4} from 'three'; -import {NodeValueType} from './NodeValueType'; +import {getAlignOfNodeValueType, getSizeOfNodeValueType, NodeValueType} from './NodeValueType'; type buildFunction = (node: Node, inputs: string[], context: ExecutionContext) => string; @@ -11,7 +11,14 @@ interface NodeBuilder { buildBySigIndex: buildFunction[]; } -const nodeBuilders: {[key: string]: NodeBuilder} = { +const nodeBuilders: { [key: string]: NodeBuilder } = { + number: { + buildBySigIndex: [ + (node, inputs, context) => { + return `${inputs[0]}`; + }, + ], + }, vec2: { buildBySigIndex: [ (node, inputs, context) => { @@ -324,6 +331,22 @@ export class WebGPUCompiler extends BaseCompiler { } protected nodeResult = new Map(); + particleInstanceByteSize: number = 0; + hasRandom = false; + + private calculateMemoryLayout(types: NodeValueType[]) { + let offset = 0; + + for (let i = 0; i < types.length; i++) { + const type = types[i]; + const align = getAlignOfNodeValueType(type); + const size = getSizeOfNodeValueType(type); + offset = Math.ceil(offset / align) * align; + offset += size; + } + offset = Math.ceil(offset / 16) * 16; + return offset; + } private buildPredefinedStructs(statements: string[], graph: NodeGraph, context: ExecutionContext) { statements.push('struct SimulationParams {'); @@ -332,19 +355,31 @@ export class WebGPUCompiler extends BaseCompiler { statements.push('}'); statements.push('struct Particle {'); - graph.outputNodes.forEach((node) => { - if ( + statements.push(' color: vec4,'); + statements.push(' position: vec3,'); + const nodes = graph.outputNodes + .filter((node) => node.definition.name === 'particleProperty' && node.data.property !== 'position' && node.data.property !== 'color' && node.data.property !== 'life' && - node.data.property !== 'age' - ) { - statements.push(` ${node.data.property}: ${this.getTypeFromNodeType(node.data.type)},`); - } + node.data.property !== 'size' && + node.data.property !== 'rotation' && + node.data.property !== 'age') + .sort((a, b) => getSizeOfNodeValueType(a.data.type) - getSizeOfNodeValueType(b.data.type)); + this.particleInstanceByteSize = this.calculateMemoryLayout( + [NodeValueType.Vec4, NodeValueType.Vec3] + .concat(nodes.map((node) => node.data.type)) + .concat([NodeValueType.Number, NodeValueType.Number, NodeValueType.Number, NodeValueType.Number]) + ); + //this.particleInstanceByteSize += getSizeOfNodeValueType(node.data.type); + //this.particleInstanceByteSize += 4 * 4 * 2 + 4 + 4; + + nodes.forEach((node) => { + statements.push(` ${node.data.property}: ${this.getTypeFromNodeType(node.data.type)},`); }); - statements.push(' position: vec3,'); - statements.push(' color: vec3,'); + statements.push(' size: f32,'); + statements.push(' rotation: f32,'); statements.push(' life: f32,'); statements.push(' age: f32,'); statements.push('}'); @@ -360,13 +395,16 @@ export class WebGPUCompiler extends BaseCompiler { statements.push('@binding(0) @group(0) var sim_params : SimulationParams;'); statements.push('@binding(1) @group(0) var data : Particles;'); - statements.push('@binding(2) @group(0) var indexList : array;'); + statements.push('@binding(2) @group(0) var indexList : array;'); statements.push('@compute @workgroup_size(64)'); statements.push('fn simulate(@builtin(global_invocation_id) global_invocation_id : vec3) {'); statements.push(' let invo_id = global_invocation_id.x;'); - statements.push(' init_rand(idx, sim_params.seed);'); statements.push(' let idx = indexList[invo_id];'); + //statements.push(' let idx = invo_id;'); + if (this.hasRandom) { + statements.push(' init_rand(invo_id, sim_params.seed);'); + } statements.push(' var particle = data.particles[idx];'); } @@ -389,7 +427,11 @@ export class WebGPUCompiler extends BaseCompiler { private buildDataFlow(graph: NodeGraph, statements: string[], context: ExecutionContext) { for (let i = 0; i < graph.nodesInOrder.length; i++) { const currentNode = graph.nodesInOrder[i]; - let nodeBuilder = nodeBuilders[currentNode.definition.name].buildBySigIndex[currentNode.signatureIndex]; + const nodeBuilder = nodeBuilders[currentNode.definition.name]; + if (nodeBuilder === undefined) { + throw new Error(`Node ${currentNode.id} ${currentNode.definition.name} has no builder`); + } + let nodeBuilderFunc = nodeBuilder.buildBySigIndex[currentNode.signatureIndex]; const inputs: string[] = []; for (let j = 0; j < currentNode.inputs.length; j++) { if (currentNode.inputs[j] instanceof Wire) { @@ -408,7 +450,7 @@ export class WebGPUCompiler extends BaseCompiler { } } } - const result = nodeBuilder(currentNode, inputs, context); + const result = nodeBuilderFunc(currentNode, inputs, context); if (currentNode.outputs.length === 1) { if (currentNode.definition.type === NodeType.Storage) { if (inputs.length > 0) { diff --git a/src/nodes/index.ts b/src/nodes/index.ts index 7c7342e..e29a4dd 100644 --- a/src/nodes/index.ts +++ b/src/nodes/index.ts @@ -5,3 +5,4 @@ export * from './NodeGraph'; export * from './NodeDef'; export * from './NodeValueType'; export * from './NodeVFX'; +export * from './WebGPUCompiler'; diff --git a/src/shaders/wgsl/particle.wgsl.ts b/src/shaders/wgsl/particle.wgsl.ts new file mode 100644 index 0000000..198cba8 --- /dev/null +++ b/src/shaders/wgsl/particle.wgsl.ts @@ -0,0 +1,60 @@ +export default ` +//////////////////////////////////////////////////////////////////////////////// +// Vertex shader +//////////////////////////////////////////////////////////////////////////////// +struct RenderParams { + modelViewProjectionMatrix : mat4x4f, + right : vec3, + up : vec3 +} + +struct Particle { + color: vec4, + position: vec3, + velocity: vec3, + life: f32, + age: f32, +} + +@group(0) @binding(0) var render_params : RenderParams; +@group(0) @binding(1) var particles: array; +@group(0) @binding(2) var particleIndices: array; + +struct VertexInput { + /*@location(0) position : vec3f, + @location(1) color : vec4f,*/ + @location(0) quad_pos : vec2f, // -1..+1 + @builtin(instance_index) instanceIndex: u32, +} + +struct VertexOutput { + @builtin(position) position : vec4f, + @location(0) color : vec4f, + @location(1) quad_pos : vec2f, // -1..+1 +} + +@vertex +fn vs_main(in : VertexInput) -> VertexOutput { + var quad_pos = mat2x3f(render_params.right, render_params.up) * in.quad_pos; + + var particle = particles[particleIndices[in.instanceIndex]]; + var position = particle.position + quad_pos * 0.01; + var out : VertexOutput; + out.position = render_params.modelViewProjectionMatrix * vec4f(position, 1.0); + out.color = particle.color; + out.quad_pos = in.quad_pos; + return out; +} + +//////////////////////////////////////////////////////////////////////////////// +// Fragment shader +//////////////////////////////////////////////////////////////////////////////// +@fragment +fn fs_main(in : VertexOutput) -> @location(0) vec4f { + var color = vec4(1.0, 0.0, 0.0, 1.0); //in.color; + // Apply a circular particle alpha mask + color.a = color.a * max(1.0 - length(in.quad_pos), 0.0); + return color; +} + +`; \ No newline at end of file diff --git a/types/index.d.ts b/types/index.d.ts index 348145d..a89156b 100644 --- a/types/index.d.ts +++ b/types/index.d.ts @@ -1019,7 +1019,7 @@ interface GPURenderPassDescriptor extends GPUObjectDescriptorBase { * Due to compatible usage list|usage compatibility, no color attachment * may alias another attachment or any resource used inside the render pass. */ - colorAttachments: Iterable; + colorAttachments: ArrayLike; /** * The {@link GPURenderPassDepthStencilAttachment} value that defines the depth/stencil * attachment that will be output to and tested against when executing this render pass.