WebGPU compute shader support prototype.

examples/files.js

@@ -323,7 +323,8 @@ var files = {
 	],
 	"webgpu": [
 		"webgpu_sandbox",
-		"webgpu_rtt"
+		"webgpu_rtt",
+		"webgpu_compute",
 	],
 	"webaudio": [
 		"webaudio_orientation",

examples/jsm/renderers/webgpu/WebGPUAttributes.js

@@ -27,13 +27,25 @@ class WebGPUAttributes {
 
 	}
 
-	update( attribute, isIndex = false ) {
+	update( attribute, isIndex = false, usage = null ) {
 
 		let data = this.buffers.get( attribute );
 
 		if ( data === undefined ) {
 
-			const usage = ( isIndex === true ) ? GPUBufferUsage.INDEX : GPUBufferUsage.VERTEX;
+			if ( usage === null ) {
+
+				usage = ( isIndex === true ) ? GPUBufferUsage.INDEX : GPUBufferUsage.VERTEX;
+
+			}
+
+			data = this._createBuffer( attribute, usage );
+
+			this.buffers.set( attribute, data );
+
+		} else if ( usage && usage !== data.usage ) {
+
+			data.buffer.destroy();
 
 			data = this._createBuffer( attribute, usage );
 
@@ -68,7 +80,8 @@ class WebGPUAttributes {
 
 		return {
 			version: attribute.version,
-			buffer: buffer
+			buffer: buffer,
+			usage: usage
 		};
 
 	}

examples/jsm/renderers/webgpu/WebGPUBindings.js

@@ -1,3 +1,4 @@
+import WebGPUStorageBuffer from './WebGPUStorageBuffer.js';
 import WebGPUUniformsGroup from './WebGPUUniformsGroup.js';
 import { FloatUniform, Matrix3Uniform, Matrix4Uniform } from './WebGPUUniform.js';
 import WebGPUSampler from './WebGPUSampler.js';
@@ -5,13 +6,15 @@ import { WebGPUSampledTexture } from './WebGPUSampledTexture.js';
 
 class WebGPUBindings {
 
-	constructor( device, info, properties, textures, pipelines ) {
+	constructor( device, info, properties, textures, pipelines, computePipelines, attributes ) {
 
 		this.device = device;
 		this.info = info;
 		this.properties = properties;
 		this.textures = textures;
 		this.pipelines = pipelines;
+		this.computePipelines = computePipelines;
+		this.attributes = attributes;
 
 		this.uniformsData = new WeakMap();
 
@@ -72,6 +75,31 @@ class WebGPUBindings {
 
 	}
 
+	getForCompute( param ) {
+
+		let data = this.uniformsData.get( param );
+
+		if ( data === undefined ) {
+
+			const pipeline = this.computePipelines.get( param );
+			const bindings = param.bindings !== undefined ? param.bindings.slice() : [];
+			const bindLayout = pipeline.getBindGroupLayout( 0 );
+			const bindGroup = this._createBindGroup( bindings, bindLayout );
+
+			data = {
+				layout: bindLayout,
+				group: bindGroup,
+				bindings: bindings
+			};
+
+			this.uniformsData.set( param, data );
+
+		}
+
+		return data;
+
+	}
+
 	update( object, camera ) {
 
 		const textures = this.textures;
@@ -116,6 +144,11 @@ class WebGPUBindings {
 
 				updateMap.set( binding, frame );
 
+			} else if ( binding.isStorageBuffer ) {
+
+				const attribute = binding.attribute;
+				this.attributes.update( attribute, false, binding.usage );
+
 			} else if ( binding.isSampler ) {
 
 				const material = object.material;
@@ -198,6 +231,19 @@ class WebGPUBindings {
 
 				entries.push( { binding: bindingPoint, resource: { buffer: binding.bufferGPU } } );
 
+			} else if ( binding.isStorageBuffer ) {
+
+				if ( binding.bufferGPU === null ) {
+
+					const attribute = binding.attribute;
+
+					this.attributes.update( attribute, false, binding.usage );
+					binding.bufferGPU = this.attributes.get( attribute ).buffer;
+
+				}
+
+				entries.push( { binding: bindingPoint, resource: { buffer: binding.bufferGPU } } );
+
 			} else if ( binding.isSampler ) {
 
 				if ( binding.samplerGPU === null ) {

examples/jsm/renderers/webgpu/WebGPUComputePipelines.js

@@ -0,0 +1,72 @@
+class WebGPUComputePipelines {
+
+	constructor( device, glslang ) {
+
+		this.device = device;
+		this.glslang = glslang;
+
+		this.pipelines = new WeakMap();
+		this.shaderModules = {
+			compute: new WeakMap()
+		};
+
+	}
+
+	get( param ) {
+
+		let pipeline = this.pipelines.get( param );
+
+		if ( pipeline === undefined ) {
+
+			const device = this.device;
+			const shader = {
+				computeShader: param.shader
+			};
+
+			// shader modules
+
+			const glslang = this.glslang;
+
+			let moduleCompute = this.shaderModules.compute.get( shader );
+
+			if ( moduleCompute === undefined ) {
+
+				const byteCodeCompute = glslang.compileGLSL( shader.computeShader, 'compute' );
+
+				moduleCompute = device.createShaderModule( { code: byteCodeCompute } );
+
+				this.shaderModules.compute.set( shader, moduleCompute );
+
+			}
+
+			//
+
+			const computeStage = {
+				module: moduleCompute,
+				entryPoint: 'main'
+			};
+	
+			pipeline = device.createComputePipeline( {
+				computeStage: computeStage
+			} );
+
+			this.pipelines.set( param, pipeline );
+
+		}
+
+		return pipeline;
+
+	}
+
+	dispose() {
+
+		this.pipelines = new WeakMap();
+		this.shaderModules = {
+			compute: new WeakMap()
+		};
+
+	}
+
+}
+
+export default WebGPUComputePipelines;

examples/jsm/renderers/webgpu/WebGPURenderer.js

@@ -5,6 +5,7 @@ import WebGPUGeometries from './WebGPUGeometries.js';
 import WebGPUInfo from './WebGPUInfo.js';
 import WebGPUProperties from './WebGPUProperties.js';
 import WebGPURenderPipelines from './WebGPURenderPipelines.js';
+import WebGPUComputePipelines from './WebGPUComputePipelines.js';
 import WebGPUBindings from './WebGPUBindings.js';
 import WebGPURenderLists from './WebGPURenderLists.js';
 import WebGPUTextures from './WebGPUTextures.js';
@@ -99,6 +100,7 @@ class WebGPURenderer {
 		this._bindings = null;
 		this._objects = null;
 		this._renderPipelines = null;
+		this._computePipelines = null;
 		this._renderLists = null;
 		this._textures = null;
 		this._background = null;
@@ -174,7 +176,8 @@ class WebGPURenderer {
 		this._textures = new WebGPUTextures( device, this._properties, this._info, compiler );
 		this._objects = new WebGPUObjects( this._geometries, this._info );
 		this._renderPipelines = new WebGPURenderPipelines( this, this._properties, device, compiler, parameters.sampleCount );
-		this._bindings = new WebGPUBindings( device, this._info, this._properties, this._textures, this._renderPipelines );
+		this._computePipelines = new WebGPUComputePipelines( device, compiler );
+		this._bindings = new WebGPUBindings( device, this._info, this._properties, this._textures, this._renderPipelines, this._computePipelines, this._attributes );
 		this._renderLists = new WebGPURenderLists();
 		this._background = new WebGPUBackground( this );
 
@@ -510,6 +513,7 @@ class WebGPURenderer {
 		this._objects.dispose();
 		this._properties.dispose();
 		this._renderPipelines.dispose();
+		this._computePipelines.dispose();
 		this._bindings.dispose();
 		this._info.dispose();
 		this._renderLists.dispose();
@@ -529,6 +533,34 @@ class WebGPURenderer {
 
 	}
 
+	compute( computeParams ) {
+
+		const device = this._device;
+		const cmdEncoder = device.createCommandEncoder( {} );
+		const passEncoder = cmdEncoder.beginComputePass();
+
+		for ( const param of computeParams ) {
+
+			// pipeline
+
+			const pipeline = this._computePipelines.get( param );
+			passEncoder.setPipeline( pipeline );
+
+			// bind group
+
+			const bindGroup = this._bindings.getForCompute( param ).group;
+			this._bindings.update( param );
+			passEncoder.setBindGroup( 0, bindGroup );
+
+			passEncoder.dispatch( param.num );
+
+		}
+
+		passEncoder.endPass();
+		device.defaultQueue.submit( [ cmdEncoder.finish() ] );
+
+	}
+
 	getRenderTarget() {
 
 		return this._renderTarget;

examples/jsm/renderers/webgpu/WebGPUStorageBuffer.js

@@ -0,0 +1,22 @@
+import WebGPUBinding from './WebGPUBinding.js';
+
+class WebGPUStorageBuffer extends WebGPUBinding {
+
+	constructor ( name, attribute ) {
+
+		super( name );
+
+		this.type = 'storage-buffer';
+
+		this.usage = GPUBufferUsage.VERTEX | GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST;
+
+		this.attribute = attribute;
+		this.bufferGPU = null; // set by the renderer
+
+		Object.defineProperty( this, 'isStorageBuffer', { value: true } );
+
+	}
+
+}
+
+export default WebGPUStorageBuffer;

examples/webgpu_compute.html

@@ -0,0 +1,210 @@
+<html lang="en">
+	<head>
+		<title>WebGPU Compute</title>
+		<meta charset="utf-8">
+		<meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
+		<link type="text/css" rel="stylesheet" href="main.css">
+	</head>
+	<body>
+		<div id="info">
+			<a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> WebGPU - Compute<br/>(Chrome Canary with flag: --enable-unsafe-webgpu)
+		</div>
+
+		<script type="module">
+
+			import * as THREE from '../build/three.module.js';
+
+			import WebGPURenderer from './jsm/renderers/webgpu/WebGPURenderer.js';
+			import WebGPU from './jsm/renderers/webgpu/WebGPU.js';
+
+			import WebGPUStorageBuffer from './jsm/renderers/webgpu/WebGPUStorageBuffer.js';
+			import WebGPUUniformsGroup from './jsm/renderers/webgpu/WebGPUUniformsGroup.js';
+			import { Vector2Uniform } from './jsm/renderers/webgpu/WebGPUUniform.js';
+
+			let camera, scene, renderer;
+			let pointer;
+
+			const computeParams = [];
+
+			init().then( animate ).catch( onInitError );
+
+			async function init() {
+
+				if ( WebGPU.isAvailable() === false ) {
+
+					document.body.appendChild( WebGPU.getErrorMessage() );
+
+					throw 'No WebGPU support';
+
+				}
+
+				camera = new THREE.OrthographicCamera( - 1.0, 1.0, 1.0, - 1.0, 0, 1 );
+				camera.position.z = 1;
+
+				scene = new THREE.Scene();
+				scene.background = new THREE.Color( 0x000000 );
+
+				const particleNum = 100000;
+				const particleSize = 3;
+
+				const particleArray = new Float32Array( particleNum * particleSize );
+				const velocityArray = new Float32Array( particleNum * particleSize );
+
+				for ( let i = 0; i < particleArray.length; i += 3 ) {
+
+					const r = Math.random() * 0.01 + 0.0005;
+					const degree = Math.random() * 360;
+					velocityArray[ i * 3 + 0 ] = r * Math.sin( degree * Math.PI / 180 );
+					velocityArray[ i * 3 + 1 ] = r * Math.cos( degree * Math.PI / 180 );
+
+				}
+
+				const particleBuffer = new WebGPUStorageBuffer( 'particle', new THREE.BufferAttribute( particleArray, 3 ) );
+				const velocityBuffer = new WebGPUStorageBuffer( 'velocity', new THREE.BufferAttribute( velocityArray, 3 ) );
+
+				pointer = new THREE.Vector2( - 10.0, - 10.0 ); // Out of bounds first
+
+				const pointerGroup = new WebGPUUniformsGroup( 'mouseUniforms' ).addUniform(
+					new Vector2Uniform( 'pointer', pointer )
+				);
+
+				const computeBindings = [
+					particleBuffer,
+					velocityBuffer,
+					pointerGroup
+				];
+
+				const computeShader = `#version 450
+					#define PARTICLE_NUM ${particleNum}
+					#define PARTICLE_SIZE ${particleSize}
+					#define ROOM_SIZE 1.0
+					#define POINTER_SIZE 0.1
+
+					// Limitation for now: the order should be the same as bindings order
+
+					layout(set = 0, binding = 0) buffer Particle {
+						float particle[ PARTICLE_NUM * PARTICLE_SIZE ];
+					} particle;
+
+					layout(set = 0, binding = 1) buffer Velocity {
+						float velocity[ PARTICLE_NUM * PARTICLE_SIZE ];
+					} velocity;
+
+					layout(set = 0, binding = 2) uniform MouseUniforms {
+						vec2 pointer;
+					} mouseUniforms;
+
+					void main() {
+						uint index = gl_GlobalInvocationID.x;
+						if ( index >= PARTICLE_NUM ) { return; }
+
+						vec3 position = vec3(
+							particle.particle[ index * 3 + 0 ] + velocity.velocity[ index * 3 + 0 ],
+							particle.particle[ index * 3 + 1 ] + velocity.velocity[ index * 3 + 1 ],
+							particle.particle[ index * 3 + 2 ] + velocity.velocity[ index * 3 + 2 ]
+						);
+
+						if ( abs( position.x ) >= ROOM_SIZE ) {
+
+							velocity.velocity[ index * 3 + 0 ] = - velocity.velocity[ index * 3 + 0 ];
+
+						}
+
+						if ( abs( position.y ) >= ROOM_SIZE ) {
+
+							velocity.velocity[ index * 3 + 1 ] = - velocity.velocity[ index * 3 + 1 ];
+
+						}
+
+						if ( abs( position.z ) >= ROOM_SIZE ) {
+
+							velocity.velocity[ index * 3 + 2 ] = - velocity.velocity[ index * 3 + 2 ];
+
+						}
+
+						float dx = mouseUniforms.pointer.x - position.x;
+						float dy = mouseUniforms.pointer.y - position.y;
+						float distanceFromPointer = sqrt( dx * dx + dy * dy );
+
+						if ( distanceFromPointer <= POINTER_SIZE ) {
+
+							position.x = 0.0;
+							position.y = 0.0;
+							position.z = 0.0;
+
+						}
+
+						particle.particle[ index * 3 + 0 ] = position.x;
+						particle.particle[ index * 3 + 1 ] = position.y;
+						particle.particle[ index * 3 + 2 ] = position.z;
+					}
+				`;
+
+				computeParams.push( {
+					num: particleNum,
+					shader: computeShader,
+					bindings: computeBindings
+				} );
+
+				// Use a compute shader to animate the point cloud's vertex data.
+
+				const pointsGeometry = new THREE.BufferGeometry().setAttribute(
+					'position', particleBuffer.attribute
+				);
+				const pointsMaterial = new THREE.PointsMaterial();
+				const mesh = new THREE.Points( pointsGeometry, pointsMaterial );
+				scene.add( mesh );
+
+				renderer = new WebGPURenderer();
+				renderer.setPixelRatio( window.devicePixelRatio );
+				renderer.setSize( window.innerWidth, window.innerHeight );
+				document.body.appendChild( renderer.domElement );
+
+				window.addEventListener( 'resize', onWindowResize, false );
+				window.addEventListener( 'mousemove', onMouseMove, false );
+
+				return renderer.init();
+
+			}
+
+			function onWindowResize() {
+
+				camera.updateProjectionMatrix();
+
+				renderer.setSize( window.innerWidth, window.innerHeight );
+
+			}
+
+			function onMouseMove( event ) {
+
+				const x = event.clientX;
+				const y = event.clientY;
+
+				const width = window.innerWidth;
+				const height = window.innerHeight;
+
+				pointer.set(
+					( x / width - 0.5 ) * 2.0,
+					( - y / height + 0.5 ) * 2.0
+				);
+			
+			}
+
+			function animate() {
+
+				requestAnimationFrame( animate );
+
+				renderer.compute( computeParams );
+				renderer.render( scene, camera );
+
+			}
+
+			function onInitError( error ) {
+
+				console.error( error );
+
+			}
+
+		</script>
+	</body>
+</html>