WebGPURenderer: BatchMesh support for Instanced rendering with sorting, frustum culling (#28753)

* WebGPURenderer: Full BatchMesh Support in both backend

* cleanup

* fix circular deps

* webgpu wip

* cleanup webgup

* webgpu

* cleanup

* cleanup

* cleanup

* add the ability to specify nodeType in TextureNode and CubeTextureNode

* drawIndex

* cleanup

* fix uint override to indirectTexture in BatchNode

* update batch mesh example

* fix glslnodebuilder

* cleanup and fix glsl

* more cleanup

* cleanup and feedbacks

* minor cleanup

---------

Co-authored-by: sunag <[email protected]>

examples/jsm/nodes/Nodes.js

@@ -11,7 +11,7 @@ export { default as BypassNode, bypass } from './core/BypassNode.js';
 export { default as CacheNode, cache } from './core/CacheNode.js';
 export { default as ConstNode } from './core/ConstNode.js';
 export { default as ContextNode, context, label } from './core/ContextNode.js';
-export { default as IndexNode, vertexIndex, instanceIndex } from './core/IndexNode.js';
+export { default as IndexNode, vertexIndex, instanceIndex, drawIndex } from './core/IndexNode.js';
 export { default as LightingModel } from './core/LightingModel.js';
 export { default as Node, addNodeClass, createNodeFromType } from './core/Node.js';
 export { default as VarNode, temp } from './core/VarNode.js';

examples/jsm/nodes/accessors/BatchNode.js

@@ -1,11 +1,11 @@
 import Node, { addNodeClass } from '../core/Node.js';
 import { normalLocal } from './NormalNode.js';
 import { positionLocal } from './PositionNode.js';
-import { nodeProxy, vec3, mat3, mat4, int, ivec2, float } from '../shadernode/ShaderNode.js';
+import { nodeProxy, vec3, mat3, mat4, int, ivec2, float, tslFn } from '../shadernode/ShaderNode.js';
 import { textureLoad } from './TextureNode.js';
 import { textureSize } from './TextureSizeNode.js';
-import { attribute } from '../core/AttributeNode.js';
 import { tangentLocal } from './TangentNode.js';
+import { instanceIndex, drawIndex } from '../core/IndexNode.js';
 
 class BatchNode extends Node {
 
@@ -28,14 +28,37 @@ class BatchNode extends Node {
 
 		if ( this.batchingIdNode === null ) {
 
-			this.batchingIdNode = attribute( 'batchId' );
+			if ( builder.getDrawIndex() === null ) {
+
+				this.batchingIdNode = instanceIndex;
+
+			} else {
+
+				this.batchingIdNode = drawIndex;
+
+			}
 
 		}
 
+		const getIndirectIndex = tslFn( ( [ id ] ) => {
+
+			const size = textureSize( textureLoad( this.batchMesh._indirectTexture ), 0 );
+			const x = int( id ).remainder( int( size ) );
+			const y = int( id ).div( int( size ) );
+			return textureLoad( this.batchMesh._indirectTexture, ivec2( x, y ), null, 'uvec4' ).x;
+
+		} ).setLayout( {
+			name: 'getIndirectIndex',
+			type: 'uint',
+			inputs: [
+				{ name: 'id', type: 'int' }
+			]
+		} );
+
 		const matriceTexture = this.batchMesh._matricesTexture;
 
 		const size = textureSize( textureLoad( matriceTexture ), 0 );
-		const j = float( int( this.batchingIdNode ) ).mul( 4 ).toVar();
+		const j = float( getIndirectIndex( int( this.batchingIdNode ) ) ).mul( 4 ).toVar();
 
 		const x = int( j.mod( size ) );
 		const y = int( j ).div( int( size ) );

examples/jsm/nodes/accessors/TextureNode.js

@@ -7,6 +7,7 @@ import { addNodeClass } from '../core/Node.js';
 import { maxMipLevel } from '../utils/MaxMipLevelNode.js';
 import { addNodeElement, nodeProxy, vec3, nodeObject } from '../shadernode/ShaderNode.js';
 import { NodeUpdateType } from '../core/constants.js';
+import { IntType, UnsignedIntType } from 'three';
 
 class TextureNode extends UniformNode {
 
@@ -64,6 +65,16 @@ class TextureNode extends UniformNode {
 
 		if ( this.value.isDepthTexture === true ) return 'float';
 
+		if ( this.value.type === UnsignedIntType ) {
+
+			return 'uvec4';
+
+		} else if ( this.value.type === IntType ) {
+
+			return 'ivec4';
+
+		}
+
 		return 'vec4';
 
 	}

examples/jsm/nodes/core/IndexNode.js

@@ -29,6 +29,10 @@ class IndexNode extends Node {
 
 			propertyName = builder.getInstanceIndex();
 
+		} else if ( scope === IndexNode.DRAW ) {
+
+			propertyName = builder.getDrawIndex();
+
 		} else {
 
 			throw new Error( 'THREE.IndexNode: Unknown scope: ' + scope );
@@ -57,10 +61,12 @@ class IndexNode extends Node {
 
 IndexNode.VERTEX = 'vertex';
 IndexNode.INSTANCE = 'instance';
+IndexNode.DRAW = 'draw';
 
 export default IndexNode;
 
 export const vertexIndex = nodeImmutable( IndexNode, IndexNode.VERTEX );
 export const instanceIndex = nodeImmutable( IndexNode, IndexNode.INSTANCE );
+export const drawIndex = nodeImmutable( IndexNode, IndexNode.DRAW );
 
 addNodeClass( 'IndexNode', IndexNode );

examples/jsm/nodes/core/NodeBuilder.js

@@ -451,6 +451,12 @@ class NodeBuilder {
 
 	}
 
+	getDrawIndex() {
+
+		console.warn( 'Abstract function.' );
+
+	}
+
 	getFrontFacing() {
 
 		console.warn( 'Abstract function.' );

examples/jsm/renderers/webgl/WebGLBackend.js

@@ -11,6 +11,7 @@ import WebGLExtensions from './utils/WebGLExtensions.js';
 import WebGLCapabilities from './utils/WebGLCapabilities.js';
 import { GLFeatureName } from './utils/WebGLConstants.js';
 import { WebGLBufferRenderer } from './WebGLBufferRenderer.js';
+import { warnOnce } from '../../../../src/utils.js';
 
 //
 
@@ -51,6 +52,8 @@ class WebGLBackend extends Backend {
 		this.trackTimestamp = ( parameters.trackTimestamp === true );
 
 		this.extensions.get( 'EXT_color_buffer_float' );
+		this.extensions.get( 'WEBGL_multi_draw' );
+
 		this.disjoint = this.extensions.get( 'EXT_disjoint_timer_query_webgl2' );
 		this.parallel = this.extensions.get( 'KHR_parallel_shader_compile' );
 		this._currentContext = null;
@@ -701,9 +704,9 @@ class WebGLBackend extends Backend {
 
 		if ( object.isBatchedMesh ) {
 
-			if ( object._multiDrawInstances !== null ) {
+			if ( ! this.hasFeature( 'WEBGL_multi_draw' ) ) {
 
-				renderer.renderMultiDrawInstances( object._multiDrawStarts, object._multiDrawCounts, object._multiDrawCount, object._multiDrawInstances );
+				warnOnce( 'THREE.WebGLRenderer: WEBGL_multi_draw not supported.' );
 
 			} else {
 

examples/jsm/renderers/webgl/nodes/GLSLNodeBuilder.js

@@ -157,7 +157,7 @@ ${ flowData.code }
 			pboTexture.needsUpdate = true;
 			pboTexture.isPBOTexture = true;
 
-			const pbo = new TextureNode( pboTexture );
+			const pbo = new TextureNode( pboTexture, null, null );
 			pbo.setPrecision( 'high' );
 
 			attribute.pboNode = pbo;
@@ -241,14 +241,14 @@ ${ flowData.code }
 
 			//
 
-			const typePrefix = attribute.array.constructor.name.toLowerCase().charAt( 0 );
 
 			let prefix = 'vec4';
-			if ( typePrefix === 'u' ) {
+
+			if ( attribute.pbo.type === UnsignedIntType ) {
 
 				prefix = 'uvec4';
 
-			} else if ( typePrefix === 'i' ) {
+			} else if ( attribute.pbo.type === IntType ) {
 
 				prefix = 'ivec4';
 
@@ -358,13 +358,16 @@ ${ flowData.code }
 
 				let typePrefix = '';
 
-				if ( texture.isPBOTexture === true ) {
+				if ( texture.isDataTexture === true ) {
+
 
-					const prefix = texture.source.data.data.constructor.name.toLowerCase().charAt( 0 );
+					if ( texture.type === UnsignedIntType ) {
 
-					if ( prefix === 'u' || prefix === 'i' ) {
+						typePrefix = 'u';
 
-						typePrefix = prefix;
+					} else if ( texture.type === IntType ) {
+
+						typePrefix = 'i';
 
 					}
 
@@ -594,6 +597,20 @@ ${ flowData.code }
 
 	}
 
+	getDrawIndex() {
+
+		const extensions = this.renderer.backend.extensions;
+
+		if ( extensions.has( 'WEBGL_multi_draw' ) ) {
+
+			return 'uint( gl_DrawID )';
+
+		}
+
+		return null;
+
+	}
+
 	getFrontFacing() {
 
 		return 'gl_FrontFacing';
@@ -612,6 +629,27 @@ ${ flowData.code }
 
 	}
 
+	getExtensions( shaderStage ) {
+
+		let extensions = '';
+
+		if ( shaderStage === 'vertex' ) {
+
+			const ext = this.renderer.backend.extensions;
+			const isBatchedMesh = this.object.isBatchedMesh;
+
+			if ( isBatchedMesh && ext.has( 'WEBGL_multi_draw' ) ) {
+
+				extensions += '#extension GL_ANGLE_multi_draw : require\n';
+
+			}
+
+		}
+
+		return extensions;
+
+	}
+
 	isAvailable( name ) {
 
 		let result = supports[ name ];
@@ -620,11 +658,11 @@ ${ flowData.code }
 
 			if ( name === 'float32Filterable' ) {
 
-				const extentions = this.renderer.backend.extensions;
+				const extensions = this.renderer.backend.extensions;
 
-				if ( extentions.has( 'OES_texture_float_linear' ) ) {
+				if ( extensions.has( 'OES_texture_float_linear' ) ) {
 
-					extentions.get( 'OES_texture_float_linear' );
+					extensions.get( 'OES_texture_float_linear' );
 					result = true;
 
 				} else {
@@ -688,7 +726,8 @@ ${vars}
 
 		return `#version 300 es
 
-${ this.getSignature() }
+// extensions 
+${shaderData.extensions}
 
 // precision
 ${ defaultPrecisions }
@@ -809,6 +848,7 @@ void main() {
 
 			const stageData = shadersData[ shaderStage ];
 
+			stageData.extensions = this.getExtensions( shaderStage );
 			stageData.uniforms = this.getUniforms( shaderStage );
 			stageData.attributes = this.getAttributes( shaderStage );
 			stageData.varyings = this.getVaryings( shaderStage );

examples/jsm/renderers/webgl/utils/WebGLConstants.js

@@ -1,5 +1,6 @@
 export const GLFeatureName = {
 
+	'WEBGL_multi_draw': 'WEBGL_multi_draw',
 	'WEBGL_compressed_texture_astc': 'texture-compression-astc',
 	'WEBGL_compressed_texture_etc': 'texture-compression-etc2',
 	'WEBGL_compressed_texture_etc1': 'texture-compression-etc1',

examples/jsm/renderers/webgpu/WebGPUBackend.js

@@ -918,7 +918,21 @@ class WebGPUBackend extends Backend {
 		const instanceCount = this.getInstanceCount( renderObject );
 		if ( instanceCount === 0 ) return;
 
-		if ( hasIndex === true ) {
+		if ( object.isBatchedMesh === true ) {
+
+			const starts = object._multiDrawStarts;
+			const counts = object._multiDrawCounts;
+			const drawCount = object._multiDrawCount;
+
+			const bytesPerElement = index.bytesPerElement || 1;
+
+			for ( let i = 0; i < drawCount; i ++ ) {
+
+				passEncoderGPU.drawIndexed( counts[ i ] / bytesPerElement, 1, starts[ i ] / 4, 0, i );
+
+			}
+
+		} else if ( hasIndex === true ) {
 
 			const indexCount = ( drawRange.count !== Infinity ) ? drawRange.count : index.count;
 

examples/jsm/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -611,6 +611,12 @@ ${ flowData.code }
 
 	}
 
+	getDrawIndex() {
+
+		return null;
+
+	}
+
 	getFrontFacing() {
 
 		return this.getBuiltin( 'front_facing', 'isFront', 'bool' );

examples/webgpu_mesh_batch.html

@@ -155,8 +155,8 @@
 		function initBatchedMesh() {
 
 			const geometryCount = api.count;
-			const vertexCount = api.count * 512;
-			const indexCount = api.count * 1024;
+			const vertexCount = geometries.length * 512;
+			const indexCount = geometries.length * 1024;
 
 			const euler = new THREE.Euler();
 			const matrix = new THREE.Matrix4();
@@ -168,9 +168,14 @@
 
 			ids.length = 0;
 
+			const geometryIds = [
+				mesh.addGeometry( geometries[ 0 ] ),
+				mesh.addGeometry( geometries[ 1 ] ),
+				mesh.addGeometry( geometries[ 2 ] ),
+			];
 			for ( let i = 0; i < api.count; i ++ ) {
 
-				const id = mesh.addGeometry( geometries[ i % geometries.length ] );
+				const id = mesh.addInstance( geometryIds[ i % geometryIds.length ] );
 				mesh.setMatrixAt( id, randomizeMatrix( matrix ) );
 
 				const rotationMatrix = new THREE.Matrix4();