@openhps/core/dist/esm/three/nodes/gpgpu/ComputeBuiltinNode.js

Open Hybrid Positioning System - Core component

import Node from '../core/Node.js';
import { nodeObject } from '../tsl/TSLBase.js';

/**
 * `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
 */
class ComputeBuiltinNode extends Node {
  static get type() {
    return 'ComputeBuiltinNode';
  }

  /**
   * Constructs a new compute builtin node.
   *
   * @param {string} builtinName - The built-in name.
   * @param {string} nodeType - The node type.
   */
  constructor(builtinName, nodeType) {
    super(nodeType);

    /**
     * The built-in name.
     *
     * @private
     * @type {string}
     */
    this._builtinName = builtinName;
  }

  /**
   * This method is overwritten since hash is derived from the built-in name.
   *
   * @param {NodeBuilder} builder - The current node builder.
   * @return {string} The hash.
   */
  getHash(builder) {
    return this.getBuiltinName(builder);
  }

  /**
   * This method is overwritten since the node type is simply derived from `nodeType`..
   *
   * @param {NodeBuilder} builder - The current node builder.
   * @return {string} The node type.
   */
  getNodeType( /*builder*/
  ) {
    return this.nodeType;
  }

  /**
   * Sets the builtin name.
   *
   * @param {string} builtinName - The built-in name.
   * @return {ComputeBuiltinNode} A reference to this node.
   */
  setBuiltinName(builtinName) {
    this._builtinName = builtinName;
    return this;
  }

  /**
   * Returns the builtin name.
   *
   * @param {NodeBuilder} builder - The current node builder.
   * @return {string} The builtin name.
   */
  getBuiltinName( /*builder*/
  ) {
    return this._builtinName;
  }

  /**
   * Whether the current node builder has the builtin or not.
   *
   * @param {NodeBuilder} builder - The current node builder.
   * @return {boolean} Whether the builder has the builtin or not.
   */
  hasBuiltin(builder) {
    return 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;

/**
 * TSL function for creating a compute builtin node.
 *
 * @tsl
 * @function
 * @param {string} name - The built-in name.
 * @param {string} nodeType - The node type.
 * @returns {ComputeBuiltinNode}
 */
const computeBuiltin = (name, nodeType) => nodeObject(new ComputeBuiltinNode(name, nodeType));

/**
 * 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);
 * ```
 *
 * @tsl
 * @type {ComputeBuiltinNode<uvec3>}
 */
export const numWorkgroups = /*@__PURE__*/computeBuiltin('numWorkgroups', 'uvec3');

/**
 * 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.mod( 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];
 * ```
 *
 * @tsl
 * @type {ComputeBuiltinNode<uvec3>}
 */
export const workgroupId = /*@__PURE__*/computeBuiltin('workgroupId', 'uvec3');

/**
 * A non-linearized 3-dimensional representation of the current invocation's position within a 3D global grid.
 *
 * @tsl
 * @type {ComputeBuiltinNode<uvec3>}
 */
export const globalId = /*@__PURE__*/computeBuiltin('globalId', 'uvec3');
/**
 * A non-linearized 3-dimensional representation of the current invocation's position within a 3D workgroup grid.
 *
 * @tsl
 * @type {ComputeBuiltinNode<uvec3>}
 */
export const localId = /*@__PURE__*/computeBuiltin('localId', 'uvec3');

/**
 * A device dependent variable that exposes the size of the current invocation's subgroup.
 *
 * @tsl
 * @type {ComputeBuiltinNode<uint>}
 */
export const subgroupSize = /*@__PURE__*/computeBuiltin('subgroupSize', 'uint');