From 0389a79dfd2a8e4b5151fbddf9b84478cfe274c8 Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Wed, 2 Sep 2026 13:46:12 -0400 Subject: [PATCH 1/3] TSL Bug: Fix `subgroupAll()`/`subgroupAny()` parameter length (#34432) --- src/nodes/gpgpu/SubgroupFunctionNode.js | 6 +- .../addons/tsl/GPUSubgroupAllAny.tests.js | 72 +++++++++++++++++++ test/unit/three.addons.unit.js | 1 + 3 files changed, 77 insertions(+), 2 deletions(-) create mode 100644 test/unit/addons/tsl/GPUSubgroupAllAny.tests.js diff --git a/src/nodes/gpgpu/SubgroupFunctionNode.js b/src/nodes/gpgpu/SubgroupFunctionNode.js index b4f3682320c9b8..b6136de9546f05 100644 --- a/src/nodes/gpgpu/SubgroupFunctionNode.js +++ b/src/nodes/gpgpu/SubgroupFunctionNode.js @@ -462,18 +462,20 @@ 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. diff --git a/test/unit/addons/tsl/GPUSubgroupAllAny.tests.js b/test/unit/addons/tsl/GPUSubgroupAllAny.tests.js new file mode 100644 index 00000000000000..8a21fbfda327b0 --- /dev/null +++ b/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)` ); + + } + + } ); + + } ); + +} ); diff --git a/test/unit/three.addons.unit.js b/test/unit/three.addons.unit.js index bd5f2526427394..4f8e26ce010120 100644 --- a/test/unit/three.addons.unit.js +++ b/test/unit/three.addons.unit.js @@ -21,6 +21,7 @@ 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/TSLDeterminant.tests.js'; import './addons/tsl/TSLFaceForward.tests.js'; import './addons/tsl/TSLGainPcurve.tests.js'; From eee00211d0a7574d1621b483947a51b142402b8b Mon Sep 17 00:00:00 2001 From: Ben Houston Date: Wed, 2 Sep 2026 13:46:59 -0400 Subject: [PATCH 2/3] TSL Bug: Fix subgroup broadcast first param length (#34433) Co-authored-by: Michael Herzog --- src/nodes/gpgpu/SubgroupFunctionNode.js | 3 +- .../tsl/GPUSubgroupBroadcastFirst.tests.js | 58 +++++++++++++++++++ test/unit/three.addons.unit.js | 1 + 3 files changed, 60 insertions(+), 2 deletions(-) create mode 100644 test/unit/addons/tsl/GPUSubgroupBroadcastFirst.tests.js diff --git a/src/nodes/gpgpu/SubgroupFunctionNode.js b/src/nodes/gpgpu/SubgroupFunctionNode.js index b6136de9546f05..2300523538354c 100644 --- a/src/nodes/gpgpu/SubgroupFunctionNode.js +++ b/src/nodes/gpgpu/SubgroupFunctionNode.js @@ -483,10 +483,9 @@ export const subgroupAny = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, * @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. diff --git a/test/unit/addons/tsl/GPUSubgroupBroadcastFirst.tests.js b/test/unit/addons/tsl/GPUSubgroupBroadcastFirst.tests.js new file mode 100644 index 00000000000000..bfcd12ffdfdaff --- /dev/null +++ b/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)` ); + + } + + } ); + + } ); + +} ); diff --git a/test/unit/three.addons.unit.js b/test/unit/three.addons.unit.js index 4f8e26ce010120..4ee61b6d8b0214 100644 --- a/test/unit/three.addons.unit.js +++ b/test/unit/three.addons.unit.js @@ -22,6 +22,7 @@ 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'; From 8cbb02f922cb68c4793524e64f7a7928474f0b08 Mon Sep 17 00:00:00 2001 From: Christian Helgeson <62450112+cmhhelgeson@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:11:53 -0700 Subject: [PATCH 3/3] TSL: IndexNode - Update JSDoc (#34413) Co-authored-by: Michael Herzog --- src/nodes/core/IndexNode.js | 73 +++++++++++++++++++++++++++++++++++-- 1 file changed, 69 insertions(+), 4 deletions(-) diff --git a/src/nodes/core/IndexNode.js b/src/nodes/core/IndexNode.js index 38c4ed15962b26..1da19c93cb7de5 100644 --- a/src/nodes/core/IndexNode.js +++ b/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}