From 7d73099a65600dcc45ed7ded4a8d9a10e6b7c6c3 Mon Sep 17 00:00:00 2001
From: Christian Helgeson <62450112+cmhhelgeson@users.noreply.github.com>
Date: Tue, 24 Sep 2024 20:11:44 -0700
Subject: [PATCH] Nodes: Access Remaining Compute Builtins (#29469)

* init

* remove duplicated function
---
 examples/webgpu_compute_sort_bitonic.html     |   5 +-
 src/nodes/TSL.js                              |   1 +
 src/nodes/core/IndexNode.js                   |  18 ++++
 src/nodes/gpgpu/ComputeBuiltinNode.js         | 100 ++++++++++++++++++
 src/renderers/webgpu/nodes/WGSLNodeBuilder.js |  23 +++-
 5 files changed, 143 insertions(+), 4 deletions(-)
 create mode 100644 src/nodes/gpgpu/ComputeBuiltinNode.js

diff --git a/examples/webgpu_compute_sort_bitonic.html b/examples/webgpu_compute_sort_bitonic.html
index 11fddf3ad1d424..4af05eab070385 100644
--- a/examples/webgpu_compute_sort_bitonic.html
+++ b/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 );
diff --git a/src/nodes/TSL.js b/src/nodes/TSL.js
index 1246c1b4692a39..a75b5f42b6f00b 100644
--- a/src/nodes/TSL.js
+++ b/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';
diff --git a/src/nodes/core/IndexNode.js b/src/nodes/core/IndexNode.js
index e77e38a726baf5..a30debe8e00728 100644
--- a/src/nodes/core/IndexNode.js
+++ b/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 );
diff --git a/src/nodes/gpgpu/ComputeBuiltinNode.js b/src/nodes/gpgpu/ComputeBuiltinNode.js
new file mode 100644
index 00000000000000..e2312095930585
--- /dev/null
+++ b/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' );
+
diff --git a/src/renderers/webgpu/nodes/WGSLNodeBuilder.js b/src/renderers/webgpu/nodes/WGSLNodeBuilder.js
index a8594e6efe4773..4fc673c8fda90f 100644
--- a/src/renderers/webgpu/nodes/WGSLNodeBuilder.js
+++ b/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' ) {