[pull] dev from mrdoob:dev

src/nodes/core/IndexNode.js

@@ -9,7 +9,7 @@ import { varying } from './VaryingNode.js';
  * - `vertexIndex`: The index of a vertex within a mesh.
  * - `instanceIndex`: The index of either a mesh instance or an invocation of a compute shader.
  * - `drawIndex`: The index of a draw call.
- * - `invocationLocalIndex`: The index of a compute invocation within the scope of a workgroup load.
+ * - `invocationLocalIndex`: The index of a compute invocation within the scope of a workgroup.
  * - `invocationSubgroupIndex`: The index of a compute invocation within the scope of a subgroup.
  * - `subgroupIndex`: The index of a compute invocation's subgroup within its workgroup.
  *
@@ -26,7 +26,7 @@ class IndexNode extends Node {
 	/**
 	 * Constructs a new index node.
 	 *
-	 * @param {('vertex'|'instance'|'subgroup'|'invocationLocal'|'invocationGlobal'|'invocationSubgroup'|'draw')} scope - The scope of the index node.
+	 * @param {('vertex'|'instance'|'subgroup'|'invocationLocal'|'invocationSubgroup'|'draw')} scope - The scope of the index node.
 	 */
 	constructor( scope ) {
 
@@ -125,7 +125,29 @@ export default IndexNode;
 export const vertexIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.VERTEX );
 
 /**
- * TSL object that represents the index of either a mesh instance or an invocation of a compute shader.
+ * TSL object that contextually represents specific index data depending on the shader stage.
+ *
+ * Within the vertex and fragment stages, `instanceIndex` will represent the index of the current mesh instance being evaluated by the shader.
+ * In these stages, use `instanceIndex` to modify a mesh based on its instance or to select per-instance data.
+ *
+ * ```js
+ * // instanceIndex will equal the current mesh's instance index between 0-500
+ * const material = new THREE.BasicNodeMaterial();
+ * material.positionNode = vec3( instanceIndex.mod( 10 ), instanceIndex.div( 10 ), 0 );
+ * const mesh = new THREE.InstancedMesh( geometry, material, 500 );
+ * ```
+ *
+ * Within the compute stage, `instanceIndex` will represent the global index of a compute invocation within the 3-dimensional compute workgroup load.
+ * In this stage, use `instanceIndex` to modify or select data at a given index within a buffer, or derive values from the index itself.
+ *
+ * ```js
+ * // instanceIndex will equal value between 0 - 255
+ * const computeFn = Fn() => {
+ *
+ * 	storageBuffer.element( instanceIndex ).assign( instanceIndex );
+ *
+ * } )().compute( 255 )
+ * ```
  *
  * @tsl
  * @type {IndexNode}
@@ -134,6 +156,21 @@ export const instanceIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.I
 
 /**
  * TSL object that represents the index of the subgroup the current compute invocation belongs to.
+ * Subgroup indices are local to the workgroups to which they belong.
+ *
+ * ```js
+ * // Execute 12 compute threads with a workgroup size of 9. Example assumes a subgroup size of 3.
+ * const computeFn = Fn( () => {
+ *
+ * 	storageBufferOne.element( instanceIndex ).assign( subgroupIndex );
+ * 	storageBufferTwo.element( instanceIndex ).assign( workgroupId.x );
+ *
+ * } )().compute( 12, [ 9 ] );
+ *
+ * // instanceIndex =  [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 ];
+ * // Buffer One ( Subgroup Index ) =  [ 0, 0, 0, 1, 1, 1, 2, 2, 2, 0, 0, 0 ];
+ * // Buffer Two ( Workgroup ID )   =  [ 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1 ];
+ * ```
  *
  * @tsl
  * @type {IndexNode}
@@ -143,13 +180,41 @@ export const subgroupIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.S
 /**
  * TSL object that represents the index of a compute invocation within the scope of a subgroup.
  *
+ * ```js
+ * // Execute 12 compute threads with a workgroup size of 12. Example assumes a subgroup size of 3.
+ * const computeFn = Fn( () => {
+ *
+ * 	storageBufferOne.element( instanceIndex ).assign( invocationSubgroupIndex );
+ * 	storageBufferTwo.element( instanceIndex ).assign( subgroupIndex );
+ *
+ * } )().compute( 12, [ 12 ] );
+ *
+ * // instanceIndex =  [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 ];
+ * // Buffer One ( Invocation Subgroup Index ) =  [ 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2 ];
+ * // Buffer Two ( Subgroup Index )            =  [ 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3 ];
+ * ```
+ *
  * @tsl
  * @type {IndexNode}
  */
 export const invocationSubgroupIndex = /*@__PURE__*/ nodeImmutable( IndexNode, IndexNode.INVOCATION_SUBGROUP );
 
 /**
- * TSL object that represents the index of a compute invocation within the scope of a workgroup load.
+ * TSL object that represents the index of a compute invocation within the scope of a workgroup.
+ *
+ * ```js
+ * // Execute 12 compute threads with a workgroup size of 4.
+ * const computeFn = Fn( () => {
+ *
+ * 	storageBufferOne.element( instanceIndex ).assign( invocationLocalIndex );
+ * 	storageBufferTwo.element( instanceIndex ).assign( workgroupId.x );
+ *
+ * } )().compute( 12, [ 4 ] );
+ *
+ * // instanceIndex =  [ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 ];
+ * // Buffer One ( Invocation Local Index ) =     [ 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3 ];
+ * // Buffer Two ( Workgroup ID )           =     [ 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2 ];
+ * ```
  *
  * @tsl
  * @type {IndexNode}

src/nodes/gpgpu/SubgroupFunctionNode.js

@@ -462,29 +462,30 @@ export const subgroupMax = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode,
  *
  * @tsl
  * @method
+ * @param {boolean} e - The predicate provided by the current invocation.
  * @return {bool} The result of the computation.
  */
-export const subgroupAll = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_ALL ).setParameterLength( 0 );
+export const subgroupAll = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_ALL ).setParameterLength( 1 );
 
 /**
  * Returns true if e is true for any active invocation in the subgroup
  *
  * @tsl
  * @method
+ * @param {boolean} e - The predicate provided by the current invocation.
  * @return {bool} The result of the computation.
  */
-export const subgroupAny = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_ANY ).setParameterLength( 0 );
+export const subgroupAny = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_ANY ).setParameterLength( 1 );
 
 /**
  * Broadcasts e from the active invocation with the lowest subgroup_invocation_id in the subgroup to all other active invocations.
  *
  * @tsl
  * @method
  * @param {number} e - The value to broadcast from the lowest subgroup invocation.
- * @param {number} id - The subgroup invocation to broadcast from.
  * @return {number} The broadcast value.
  */
-export const subgroupBroadcastFirst = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_BROADCAST_FIRST ).setParameterLength( 2 );
+export const subgroupBroadcastFirst = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_BROADCAST_FIRST ).setParameterLength( 1 );
 
 /**
  * Swaps e between invocations in the quad in the X direction.

test/unit/addons/tsl/GPUSubgroupAllAny.tests.js

@@ -0,0 +1,72 @@
+import {
+	Fn,
+	instanceIndex, invocationSubgroupIndex,
+	instancedArray,
+	subgroupAll, subgroupAny,
+	uint, bool
+} from 'three/tsl';
+import { rawComputeTest, readUintBuffer } from './gpu-raw-test-utils.js';
+
+// Regression test for a bug fixed alongside this file: subgroupAll()/
+// subgroupAny() were declared with `setParameterLength(0)` in
+// SubgroupFunctionNode.js, so calling either with the one boolean predicate
+// argument the WGSL spec requires (`subgroupAll(e: bool) -> bool`) was
+// rejected outright ("parameter length exceeds limit") -- as shipped,
+// neither function was callable for its documented purpose. See
+// GPUSubgroup.tests.js for the rest of this codebase's subgroup coverage and
+// the general approach (closed-form expected values derived from each
+// invocation's own lane id / subgroup size, read back from the same
+// dispatch).
+
+const WORKGROUP_SIZE = 64;
+const WORKGROUP_COUNT = 4;
+const DISPATCH_COUNT = WORKGROUP_SIZE * WORKGROUP_COUNT;
+
+// See GPUSubgroup.tests.js's matching comment: a plain numeric `count` in
+// `.compute(count, ws)` auto-inserts a bounds-check branch that violates
+// WGSL's "subgroup functions need uniform control flow" rule.
+const DISPATCH_SIZE = [ WORKGROUP_COUNT, 1, 1 ];
+
+export default QUnit.module( 'TSL', () => {
+
+	QUnit.module( 'subgroup functions', () => {
+
+		rawComputeTest( 'subgroupAll and subgroupAny reflect a lane-0-only predicate', { requiredFeature: 'subgroups' }, async ( { assert, renderer } ) => {
+
+			const allOut = instancedArray( DISPATCH_COUNT, 'uint' );
+			const anyOut = instancedArray( DISPATCH_COUNT, 'uint' );
+
+			const kernel = Fn( () => {
+
+				// Lane 0 always exists in every active subgroup, and (for a
+				// non-divergent compute shader like this one) every lane in
+				// the subgroup is active -- so both predicates below have an
+				// unambiguous, group-size-independent expected result.
+				const isLaneZero = invocationSubgroupIndex.equal( uint( 0 ) );
+
+				// Not every lane is lane 0 (unless the subgroup has exactly
+				// 1 lane, which subgroupSize being >= 1 doesn't rule out --
+				// but this sandbox's subgroupSize is 32, so this is false).
+				allOut.element( instanceIndex ).assign( subgroupAll( bool( isLaneZero.not() ) ).select( uint( 1 ), uint( 0 ) ) );
+				// Some lane (lane 0 itself) is lane 0 -- always true.
+				anyOut.element( instanceIndex ).assign( subgroupAny( bool( isLaneZero ) ).select( uint( 1 ), uint( 0 ) ) );
+
+			} )().compute( DISPATCH_SIZE, [ WORKGROUP_SIZE ] );
+
+			await renderer.computeAsync( kernel );
+
+			const allData = await readUintBuffer( renderer, allOut );
+			const anyData = await readUintBuffer( renderer, anyOut );
+
+			for ( let i = 0; i < DISPATCH_COUNT; i ++ ) {
+
+				assert.strictEqual( allData[ i ], 0, `invocation ${ i }: subgroupAll(laneId != 0) should be false (lane 0 fails it)` );
+				assert.strictEqual( anyData[ i ], 1, `invocation ${ i }: subgroupAny(laneId == 0) should be true (lane 0 satisfies it)` );
+
+			}
+
+		} );
+
+	} );
+
+} );

test/unit/addons/tsl/GPUSubgroupBroadcastFirst.tests.js

@@ -0,0 +1,58 @@
+import {
+	Fn,
+	instanceIndex, invocationSubgroupIndex,
+	instancedArray, subgroupBroadcastFirst,
+	uint
+} from 'three/tsl';
+import { rawComputeTest, readUintBuffer } from './gpu-raw-test-utils.js';
+
+// Regression test for a bug fixed alongside this file: subgroupBroadcastFirst()
+// was declared with `setParameterLength(2)` in SubgroupFunctionNode.js, but
+// per the WGSL spec it takes exactly one argument
+// (`subgroupBroadcastFirst(e: T) -> T` -- no lane id, unlike
+// subgroupBroadcast). Calling it correctly (one argument) made three.js
+// auto-pad a bogus second argument, producing invalid WGSL ("no matching
+// call to 'subgroupBroadcastFirst(f32, abstract-float)'"). See
+// GPUSubgroup.tests.js for the rest of this codebase's subgroup coverage and
+// the general approach.
+
+const WORKGROUP_SIZE = 64;
+const WORKGROUP_COUNT = 4;
+const DISPATCH_COUNT = WORKGROUP_SIZE * WORKGROUP_COUNT;
+
+// See GPUSubgroup.tests.js's matching comment: a plain numeric `count` in
+// `.compute(count, ws)` auto-inserts a bounds-check branch that violates
+// WGSL's "subgroup functions need uniform control flow" rule.
+const DISPATCH_SIZE = [ WORKGROUP_COUNT, 1, 1 ];
+
+export default QUnit.module( 'TSL', () => {
+
+	QUnit.module( 'subgroup functions', () => {
+
+		rawComputeTest( 'subgroupBroadcastFirst reads the first active lane\'s value', { requiredFeature: 'subgroups' }, async ( { assert, renderer } ) => {
+
+			const output = instancedArray( DISPATCH_COUNT, 'uint' );
+
+			const kernel = Fn( () => {
+
+				const value = invocationSubgroupIndex.add( uint( 100 ) ); // 100, 101, 102, ...
+
+				output.element( instanceIndex ).assign( subgroupBroadcastFirst( value ) );
+
+			} )().compute( DISPATCH_SIZE, [ WORKGROUP_SIZE ] );
+
+			await renderer.computeAsync( kernel );
+
+			const data = await readUintBuffer( renderer, output );
+
+			for ( let i = 0; i < DISPATCH_COUNT; i ++ ) {
+
+				assert.strictEqual( data[ i ], 100, `invocation ${ i }: subgroupBroadcastFirst(value) should read back the first lane's value (100)` );
+
+			}
+
+		} );
+
+	} );
+
+} );

test/unit/three.addons.unit.js

@@ -21,6 +21,8 @@ import './addons/tsl/GPUComputeBuiltins.tests.js';
 import './addons/tsl/GPUBarriers.tests.js';
 import './addons/tsl/GPUSubgroup.tests.js';
 import './addons/tsl/GPUWorkgroupAtomic.tests.js';
+import './addons/tsl/GPUSubgroupAllAny.tests.js';
+import './addons/tsl/GPUSubgroupBroadcastFirst.tests.js';
 import './addons/tsl/TSLDeterminant.tests.js';
 import './addons/tsl/TSLFaceForward.tests.js';
 import './addons/tsl/TSLGainPcurve.tests.js';