WebGpu-Wasm-Raytracing-And-Particles/particlesWasmWebGpu.js

import Shader 				from "../framework/WebGpu.js";

import Matrix4 				from "../framework/Matrix4.js";

import Vector3 				from "../framework/Vector3.js";

import Camera 				from "../framework/Camera.js";

import EventManager 		from "../framework/eventManager.js";

import ShaderInpector 		from "../framework/ShaderInpector.js";

import WasmModule 			from "./wasm_particles.js";



export class ParticleSimulation {

	canvas;

	device;

	camera;

	eventManager 			= new EventManager();

	frameCount 				= 0;

	wasm;

	framebufferPtr;

	framebufferUint8;

	frameBufferTexture;

	fbSampler;

	width;

	height;

	renderShader;

	time 					= 0;

	lastTime 				= 0;

	lastFrameTime;
	

	setCanvas ( canvas ) {

		this.canvas = canvas;

		this.eventManager.setCanvas( canvas );

	}

	createPlane ( width, height, repeatU, repeatV ) {

		const vertices 		= new Float32Array( new Array(
			-width / 2, -height / 2, 0,
			 width / 2, -height / 2, 0,
			-width / 2,  height / 2, 0,
			 width / 2, -height / 2, 0,
			 width / 2,  height / 2, 0,
			-width / 2,  height / 2, 0
		) );

		const uvs 			= new Float32Array( new Array(
			0,         0,
			repeatU,   0,
			0,         repeatV,
			repeatU,   0,
			repeatU,   repeatV,
			0,         repeatV
		) );

		const result 		= {
			vertices:	vertices,
			normals:	null,
			uvs:		uvs
		};

		return result;

	}

	async setup ( offscreenCanvas, width, height ) {

		this.fpsElement 		= document.getElementById( "fps" );

		this.width 				= width;

		this.height 			= height;

		offscreenCanvas.width 	= width;

		offscreenCanvas.height 	= height;

		this.canvas 			= offscreenCanvas;

		const context 			= offscreenCanvas.getContext( "webgpu" );

		this.camera 			= new Camera( new Array( 0, 0, 1115 ), new Array( 0, -0.3, 0 ), new Array( 0, 1, 0 ) );

		this.eventManager.setup( offscreenCanvas, this.camera );

		const adapter 			= await navigator.gpu.requestAdapter();

		if ( adapter === null ) {

			throw new Error( "Failed to get GPU adapter" );

		}

		this.device 			= await adapter.requestDevice();

		this.wasm 				= await WasmModule();

		this.wasm._init_particles();

		this.particleCount 		= this.wasm._get_particle_count();

		this.particlePtr 		= this.wasm._get_particles_ptr();


		const presentationFormat = navigator.gpu.getPreferredCanvasFormat();

		context.configure({
			device: 		this.device,
			format: 		presentationFormat,
			alphaMode: 		"opaque"
		});

		this.particleSize 			= 8 * 4;

		this.particleBufferSize 	= this.particleCount * this.particleSize;

		this.renderShader 			= new Shader( this.device );

		this.renderShader.setCanvas( this.canvas );

		this.particleFloatArray = new Float32Array(
			this.wasm.HEAPU8.buffer,
			this.wasm._get_particles_ptr(),
			this.particleCount * 8
		);

		await this.renderShader.setup( "../shaders/simplePoint.wgsl" );

		this.renderShader.setVariable("particles", this.particleFloatArray);

		this.lastFrameTime 		= performance.now();

		this.bindRenderLoop();
	}



	updateTimeDelta () {

		const now 				= performance.now();

		this.deltaTimeValue 	= ( now - this.lastFrameTime ) / 1000;

		this.time 				+= this.deltaTimeValue;

		this.lastFrameTime 		= now;

		this.frameCount++;

		if ( now - this.lastTime >= 1000 ) {

			const fps 		= this.frameCount / ( ( now - this.lastTime ) / 1000 );

			const fpsText 	= "FPS: " + fps.toFixed( 1 );

			this.fpsElement.textContent = fpsText;

			this.frameCount = 0;

			this.lastTime 	= now;

		}

	}

	async render () {

		this.updateTimeDelta();

		this.wasm._update_particles( this.deltaTimeValue );

		let buffer = this.renderShader.buffers.get( "particles" );

		this.device.queue.writeBuffer(
			buffer,
			0,
			this.particleFloatArray.buffer,
			this.particleFloatArray.byteOffset,
			this.particleFloatArray.byteLength
		);

		//this.renderShader.setVariable("particles", this.particleFloatArray);


		const viewMatrix 				= this.camera.getViewMatrix();

		const projectionMatrix 			= Matrix4.createProjectionMatrix( this.camera, this.canvas );

		const viewProjectionMatrix 		= Matrix4.multiply( projectionMatrix, viewMatrix );

		const cameraWorldMatrix 		= Matrix4.invert( viewMatrix );

		const cameraPosition 			= Matrix4.getColumn( cameraWorldMatrix, 3 );

		const cameraRight				= Matrix4.getColumn( cameraWorldMatrix, 0 );

		const cameraUp					= Matrix4.getColumn( cameraWorldMatrix, 1 );


		this.renderShader.setVariable( "cameraRight", cameraRight );

		this.renderShader.setVariable( "cameraUp",    cameraUp );

		this.renderShader.setVariable("viewProjectionMatrix", viewProjectionMatrix);

		this.renderShader.renderToCanvas( 6, this.particleCount, 0 ); // 6 vertices per quad

		//this.fxaaShader.renderToCanvas(6, 1, 0); 

	}

	renderTaa() {

		const prevIndex = this.currentHistoryIndex;
		const nextIndex = 1 - prevIndex;

		// Set framebuffer into history
		this.device.queue.writeTexture(
			{ texture: this.historyTextures[nextIndex] },
			this.framebufferUint8,
			{ bytesPerRow: this.width * 4 },
			{ width: this.width, height: this.height, depthOrArrayLayers: 1 }
		);

		// Apply TAA blend
		this.taaShader.setVariable("currentFrame", this.historyTextures[nextIndex]);
		this.taaShader.setVariable("historyFrame", this.historyTextures[prevIndex]);

		this.taaShader.renderToCanvas(6, 1, 0);

		// Swap history
		this.currentHistoryIndex = nextIndex;

	}


	async setupTaa() {

		this.taaShader = new Shader(this.device);

		this.taaShader.setCanvas(this.canvas);

		const quad = this.createPlane(2, 2, 1, 1);

		this.taaShader.setAttribute("position", quad.vertices);

		this.taaShader.setAttribute("uv", quad.uvs);

		this.taaShader.topology = "triangle-list";

		await this.taaShader.setup("../shaders/taa.wgsl");

		this.taaShader.setVariable("mySampler", this.fbSampler);

		this.taaShader.setVariable("blendAmount", new Float32Array([0.1])); // you can tweak this

		this.historyTextures = [
			this.frameBufferTexture,
			this.device.createTexture({
				size: [this.width, this.height, 1],
				format: "rgba8unorm",
				usage:
					GPUTextureUsage.COPY_DST |
					GPUTextureUsage.TEXTURE_BINDING |
					GPUTextureUsage.RENDER_ATTACHMENT
			})
		];

		this.currentHistoryIndex = 0;

	}


	bindRenderLoop () {

		requestAnimationFrame( this.loop.bind( this ) );

	}



	async loop () {

		await this.render();

		this.bindRenderLoop();

	}
}