Merge pull request #20390 from takahirox/ComputeShader

WebGPU compute shader support prototype.

Author: Mugen87

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,17 +1,20 @@
+import WebGPUStorageBuffer from './WebGPUStorageBuffer.js';
 import WebGPUUniformsGroup from './WebGPUUniformsGroup.js';
 import { FloatUniform, Matrix3Uniform, Matrix4Uniform } from './WebGPUUniform.js';
 import WebGPUSampler from './WebGPUSampler.js';
 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;