WebGpu-Wasm-Raytracing-And-Particles/rayTracingWebGpuWasm.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_renderer.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();

		const presentationFormat = navigator.gpu.getPreferredCanvasFormat();

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

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

		this.renderShader.setCanvas( this.canvas );

		this.wasm 				= await WasmModule();

		this.wasm._set_resolution( width, height );

		this.framebufferPtr 	= this.wasm._get_framebuffer_ptr();

		this.framebufferUint8 	= new Uint8Array(
			this.wasm.HEAPU8.buffer,
			this.framebufferPtr,
			width * height * 4
		);

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

		this.fbSampler 			= this.device.createSampler({
			magFilter: 		"nearest",
			minFilter: 		"nearest",
			addressModeU: 	"clamp-to-edge",
			addressModeV: 	"clamp-to-edge"
		});

		const plane 			= this.createPlane( 1800, 1800, 1, 1 );

		this.renderShader.setAttribute( "position", plane.vertices );

		this.renderShader.setAttribute( "uv", plane.uvs );

		this.vertexCount 		= plane.vertices.length / 3;

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

		await this.renderShader.setup( "../shaders/triangle-list-texture-solid.wgsl" );

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

		this.renderShader.setVariable( "myTexture", this.frameBufferTexture );


		await this.setupFxaa();

		await this.setupTaa();


		this.lastFrameTime 		= performance.now();

		this.bindRenderLoop();


	}

	async setupFxaa() {

		this.fxaaShader = new Shader(this.device);
		
		this.fxaaShader.setCanvas(this.canvas);

		const fullscreenQuad = this.createPlane(2, 2, 1, 1); // covers NDC [-1,1]

		this.fxaaShader.setAttribute("position", fullscreenQuad.vertices);

		this.fxaaShader.setAttribute("uv", fullscreenQuad.uvs);

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

		await this.fxaaShader.setup("../shaders/fxaa.wgsl");

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

		this.fxaaShader.setVariable("myTexture", this.frameBufferTexture);

		this.fxaaShader.setVariable("resolution", new Float32Array([this.width, this.height]));


	}

	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_framebuffer( this.time );

		this.device.queue.writeTexture(
			{ texture: this.frameBufferTexture },
			this.framebufferUint8,
			{
				bytesPerRow: this.width * 4
			},
			{
				width: this.width,
				height: this.height,
				depthOrArrayLayers: 1
			}
		);
/*
		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 );

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

		this.renderShader.renderToCanvas( this.vertexCount, 1, 0 );
*/


		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();

	}
}