Nodes: Access Remaining Compute Builtins (#29469)

* init

* remove duplicated function

examples/webgpu_compute_sort_bitonic.html

@@ -54,7 +54,7 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { storageObject, If, vec3, not, uniform, uv, uint, float, Fn, vec2, abs, int, invocationLocalIndex, workgroupArray, uvec2, floor, instanceIndex, workgroupBarrier, atomicAdd, atomicStore } from 'three/tsl';
+			import { storageObject, If, vec3, not, uniform, uv, uint, float, Fn, vec2, abs, int, invocationLocalIndex, workgroupArray, uvec2, floor, instanceIndex, workgroupBarrier, atomicAdd, atomicStore, workgroupId } from 'three/tsl';
 
 			import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
 
@@ -264,8 +264,7 @@
 
 					// Get ids of indices needed to populate workgroup local buffer.
 					// Use .toVar() to prevent these values from being recalculated multiple times.
-					const workgroupId = instanceIndex.div( WORKGROUP_SIZE[ 0 ] ).toVar();
-					const localOffset = uint( WORKGROUP_SIZE[ 0 ] ).mul( 2 ).mul( workgroupId ).toVar();
+					const localOffset = uint( WORKGROUP_SIZE[ 0 ] ).mul( 2 ).mul( workgroupId.x ).toVar();
 
 					const localID1 = invocationLocalIndex.mul( 2 );
 					const localID2 = invocationLocalIndex.mul( 2 ).add( 1 );

src/nodes/TSL.js

@@ -142,6 +142,7 @@ export * from './geometry/RangeNode.js';
 
 // gpgpu
 export * from './gpgpu/ComputeNode.js';
+export * from './gpgpu/ComputeBuiltinNode.js';
 export * from './gpgpu/BarrierNode.js';
 export * from './gpgpu/WorkgroupInfoNode.js';
 export * from './gpgpu/AtomicFunctionNode.js';

src/nodes/core/IndexNode.js

@@ -28,20 +28,34 @@ class IndexNode extends Node {
 
 		if ( scope === IndexNode.VERTEX ) {
 
+			// The index of a vertex within a mesh.
 			propertyName = builder.getVertexIndex();
 
 		} else if ( scope === IndexNode.INSTANCE ) {
 
+			// The index of either a mesh instance or an invocation of a compute shader.
 			propertyName = builder.getInstanceIndex();
 
 		} else if ( scope === IndexNode.DRAW ) {
 
+			// The index of a draw call.
 			propertyName = builder.getDrawIndex();
 
 		} else if ( scope === IndexNode.INVOCATION_LOCAL ) {
 
+			// The index of a compute invocation within the scope of a workgroup load.
 			propertyName = builder.getInvocationLocalIndex();
 
+		} else if ( scope === IndexNode.INVOCATION_SUBGROUP ) {
+
+			// The index of a compute invocation within the scope of a subgroup.
+			propertyName = builder.getInvocationSubgroupIndex();
+
+		} else if ( scope === IndexNode.SUBGROUP ) {
+
+			// The index of the subgroup the current compute invocation belongs to.
+			propertyName = builder.getSubgroupIndex();
+
 		} else {
 
 			throw new Error( 'THREE.IndexNode: Unknown scope: ' + scope );
@@ -70,12 +84,16 @@ class IndexNode extends Node {
 
 IndexNode.VERTEX = 'vertex';
 IndexNode.INSTANCE = 'instance';
+IndexNode.SUBGROUP = 'subgroup';
 IndexNode.INVOCATION_LOCAL = 'invocationLocal';
+IndexNode.INVOCATION_SUBGROUP = 'invocationSubgroup';
 IndexNode.DRAW = 'draw';
 
 export default IndexNode;
 
 export const vertexIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.VERTEX );
 export const instanceIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.INSTANCE );
+export const subgroupIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.SUBGROUP );
+export const invocationSubgroupIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.INVOCATION_SUBGROUP );
 export const invocationLocalIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.INVOCATION_LOCAL );
 export const drawIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.DRAW );

src/nodes/gpgpu/ComputeBuiltinNode.js

@@ -0,0 +1,100 @@
+import Node from '../core/Node.js';
+import { nodeObject } from '../tsl/TSLBase.js';
+
+class ComputeBuiltinNode extends Node {
+
+	static get type() {
+
+		return 'ComputeBuiltinNode';
+
+	}
+
+	constructor( builtinName, nodeType ) {
+
+		super( nodeType );
+
+		this._builtinName = builtinName;
+
+	}
+
+	getHash( builder ) {
+
+		return this.getBuiltinName( builder );
+
+	}
+
+	getNodeType( /*builder*/ ) {
+
+		return this.nodeType;
+
+	}
+
+	setBuiltinName( builtinName ) {
+
+		this._builtinName = builtinName;
+
+		return this;
+
+	}
+
+	getBuiltinName( /*builder*/ ) {
+
+		console.log( this._builtinName );
+
+		return this._builtinName;
+
+	}
+
+	hasBuiltin( builder ) {
+
+		builder.hasBuiltin( this._builtinName );
+
+	}
+
+	generate( builder, output ) {
+
+		const builtinName = this.getBuiltinName( builder );
+		const nodeType = this.getNodeType( builder );
+
+		if ( builder.shaderStage === 'compute' ) {
+
+			return builder.format( builtinName, nodeType, output );
+
+		} else {
+
+			console.warn( `ComputeBuiltinNode: Compute built-in value ${builtinName} can not be accessed in the ${builder.shaderStage} stage` );
+			return builder.generateConst( nodeType );
+
+		}
+
+	}
+
+	serialize( data ) {
+
+		super.serialize( data );
+
+		data.global = this.global;
+		data._builtinName = this._builtinName;
+
+	}
+
+	deserialize( data ) {
+
+		super.deserialize( data );
+
+		this.global = data.global;
+		this._builtinName = data._builtinName;
+
+	}
+
+}
+
+export default ComputeBuiltinNode;
+
+const computeBuiltin = ( name, nodeType ) => nodeObject( new ComputeBuiltinNode( name, nodeType ) );
+
+export const numWorkgroups = /*@__PURE__*/ computeBuiltin( 'numWorkgroups', 'uvec3' );
+export const workgroupId = /*@__PURE__*/ computeBuiltin( 'workgroupId', 'uvec3' );
+export const localId = /*@__PURE__*/ computeBuiltin( 'localId', 'uvec3' );
+export const subgroupSize = /*@__PURE__*/ computeBuiltin( 'subgroupSize', 'uint' );
+

src/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -584,6 +584,12 @@ class WGSLNodeBuilder extends NodeBuilder {
 
 	}
 
+	hasBuiltin( name, shaderStage = this.shaderStage ) {
+
+		return ( this.builtins[ shaderStage ] !== undefined && this.builtins[ shaderStage ].has( name ) );
+
+	}
+
 	getVertexIndex() {
 
 		if ( this.shaderStage === 'vertex' ) {
@@ -656,11 +662,19 @@ ${ flowData.code }
 
 	}
 
+	getInvocationSubgroupIndex() {
+
+		this.enableSubGroups();
+
+		return this.getBuiltin( 'subgroup_invocation_id', 'invocationSubgroupIndex', 'u32', 'attribute' );
+
+	}
+
 	getSubgroupIndex() {
 
 		this.enableSubGroups();
 
-		return this.getBuiltin( 'subgroup_invocation_id', 'subgroupIndex', 'u32', 'attribute' );
+		return this.getBuiltin( 'subgroup_id', 'subgroupIndex', 'u32', 'attribute' );
 
 	}
 
@@ -819,6 +833,13 @@ ${ flowData.code }
 			this.getBuiltin( 'local_invocation_id', 'localId', 'vec3<u32>', 'attribute' );
 			this.getBuiltin( 'num_workgroups', 'numWorkgroups', 'vec3<u32>', 'attribute' );
 
+			if ( this.renderer.hasFeature( 'subgroups' ) ) {
+
+				this.enableDirective( 'subgroups', shaderStage );
+				this.getBuiltin( 'subgroup_size', 'subgroupSize', 'u32', 'attribute' );
+
+			}
+
 		}
 
 		if ( shaderStage === 'vertex' || shaderStage === 'compute' ) {