Nodes - Document ComputeBuiltinNode and AtomicFunctionNode (#30162)

* document computeBuiltin and WorkgroupInfoNode

* workgroupInfoNode

src/nodes/gpgpu/AtomicFunctionNode.js

@@ -10,6 +10,8 @@ import { nodeProxy } from '../tsl/TSLCore.js';
  * Accordingly, even if multiple atomic functions are modifying an atomic variable at once
  * atomic operations will not interfer with each other.
  *
+ * This node can only be used with a WebGPU backend.
+ *
  * @augments TempNode
  */
 class AtomicFunctionNode extends TempNode {

src/nodes/gpgpu/BarrierNode.js

@@ -6,6 +6,8 @@ import { nodeProxy } from '../tsl/TSLCore.js';
 /**
  * Represents a GPU control barrier that synchronizes compute operations within a given scope.
  *
+ * This node can only be used with a WebGPU backend.
+ *
  * @augments Node
  */
 class BarrierNode extends Node {

src/nodes/gpgpu/ComputeBuiltinNode.js

@@ -4,7 +4,10 @@ import { nodeObject } from '../tsl/TSLBase.js';
 /** @module ComputeBuiltinNode **/
 
 /**
- * TODO
+ * `ComputeBuiltinNode` represents a compute-scope builtin value that expose information
+ * about the currently running dispatch and/or the device it is running on.
+ *
+ * This node can only be used with a WebGPU backend.
  *
  * @augments Node
  */
@@ -149,6 +152,24 @@ const computeBuiltin = ( name, nodeType ) => nodeObject( new ComputeBuiltinNode(
 
 /**
  * TSL function for creating a `numWorkgroups` builtin node.
+ * Represents the number of workgroups dispatched by the compute shader.
+ * ```js
+ * // Run 512 invocations/threads with a workgroup size of 128.
+ * const computeFn = Fn(() => {
+ *
+ *     // numWorkgroups.x = 4
+ *     storageBuffer.element(0).assign(numWorkgroups.x)
+ *
+ * })().compute(512, [128]);
+ *
+ * // Run 512 invocations/threads with the default workgroup size of 64.
+ * const computeFn = Fn(() => {
+ *
+ *     // numWorkgroups.x = 8
+ *     storageBuffer.element(0).assign(numWorkgroups.x)
+ *
+ * })().compute(512);
+ * ```
  *
  * @function
  * @returns {ComputeBuiltinNode<uvec3>}
@@ -157,6 +178,26 @@ export const numWorkgroups = /*@__PURE__*/ computeBuiltin( 'numWorkgroups', 'uve
 
 /**
  * TSL function for creating a `workgroupId` builtin node.
+ * Represents the 3-dimensional index of the workgroup the current compute invocation belongs to.
+ * ```js
+ * // Execute 12 compute threads with a workgroup size of 3.
+ * const computeFn = Fn( () => {
+ *
+ * 	If( workgroupId.x.modInt( 2 ).equal( 0 ), () => {
+ *
+ * 		storageBuffer.element( instanceIndex ).assign( instanceIndex );
+ *
+ * 	} ).Else( () => {
+ *
+ * 		storageBuffer.element( instanceIndex ).assign( 0 );
+ *
+ * 	} );
+ *
+ * } )().compute( 12, [ 3 ] );
+ *
+ * // workgroupId.x =  [0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3];
+ * // Buffer Output =  [0, 1, 2, 0, 0, 0, 6, 7, 8, 0, 0, 0];
+ * ```
  *
  * @function
  * @returns {ComputeBuiltinNode<uvec3>}
@@ -164,7 +205,8 @@ export const numWorkgroups = /*@__PURE__*/ computeBuiltin( 'numWorkgroups', 'uve
 export const workgroupId = /*@__PURE__*/ computeBuiltin( 'workgroupId', 'uvec3' );
 
 /**
- * TSL function for creating a `localId` builtin node.
+ * TSL function for creating a `localId` builtin node. A non-linearized 3-dimensional
+ * representation of the current invocation's position within a 3D workgroup grid.
  *
  * @function
  * @returns {ComputeBuiltinNode<uvec3>}
@@ -172,7 +214,8 @@ export const workgroupId = /*@__PURE__*/ computeBuiltin( 'workgroupId', 'uvec3'
 export const localId = /*@__PURE__*/ computeBuiltin( 'localId', 'uvec3' );
 
 /**
- * TSL function for creating a `subgroupSize` builtin node.
+ * TSL function for creating a `subgroupSize` builtin node. A device dependent variable
+ * that exposes the size of the current invocation's subgroup.
  *
  * @function
  * @returns {ComputeBuiltinNode<uint>}

src/nodes/gpgpu/WorkgroupInfoNode.js

@@ -5,7 +5,7 @@ import Node from '../core/Node.js';
 /** @module WorkgroupInfoNode **/
 
 /**
- * TODO
+ * Represents an element of a 'workgroup' scoped buffer.
  *
  * @augments ArrayElementNode
  */
@@ -56,18 +56,25 @@ class WorkgroupInfoElementNode extends ArrayElementNode {
 }
 
 /**
- * TODO
+ * A node allowing the user to create a 'workgroup' scoped buffer within the
+ * context of a compute shader. Typically, workgroup scoped buffers are
+ * created to hold data that is transfered from a global storage scope into
+ * a local workgroup scope. For invocations within a workgroup, data
+ * access speeds on 'workgroup' scoped buffers can be significantly faster
+ * than similar access operations on globally accessible storage buffers.
+ *
+ * This node can only be used with a WebGPU backend.
  *
  * @augments Node
  */
 class WorkgroupInfoNode extends Node {
 
 	/**
-	 * Constructs a new workgroup info node.
+	 * Constructs a new buffer scoped to type scope.
 	 *
 	 * @param {String} scope - TODO.
-	 * @param {String} bufferType - The buffer type.
-	 * @param {Number} [bufferCount=0] - The buffer count.
+	 * @param {String} bufferType - The data type of a 'workgroup' scoped buffer element.
+	 * @param {Number} [bufferCount=0] - The number of elements in the buffer.
 	 */
 	constructor( scope, bufferType, bufferCount = 0 ) {
 
@@ -171,10 +178,11 @@ export default WorkgroupInfoNode;
 
 /**
  * TSL function for creating a workgroup info node.
+ * Creates a new 'workgroup' scoped array buffer.
  *
  * @function
- * @param {String} type - The buffer type.
- * @param {Number} [count=0] - The buffer count.
+ * @param {String} type - The data type of a 'workgroup' scoped buffer element.
+ * @param {Number} [count=0] - The number of elements in the buffer.
  * @returns {WorkgroupInfoNode}
  */
 export const workgroupArray = ( type, count ) => nodeObject( new WorkgroupInfoNode( 'Workgroup', type, count ) );