Accueil / npm / three@0.186.1 / src / nodes / gpgpu / WorkgroupInfoNode.js

WorkgroupInfoNode.js

import ArrayElementNode from '../utils/ArrayElementNode.js';
import Node from '../core/Node.js';
import { error, warn } from '../../utils.js';
import StackTrace from '../core/StackTrace.js';

/**
 * Represents an element of a 'workgroup' scoped buffer.
 *
 * @augments ArrayElementNode
 */
class WorkgroupInfoElementNode extends ArrayElementNode {

	/**
	 * Constructs a new workgroup info element node.
	 *
	 * @param {Node} workgroupInfoNode - The workgroup info node.
	 * @param {Node} indexNode - The index node that defines the element access.
	 */
	constructor( workgroupInfoNode, indexNode ) {

		super( workgroupInfoNode, indexNode );

		/**
		 * This flag can be used for type testing.
		 *
		 * @type {boolean}
		 * @readonly
		 * @default true
		 */
		this.isWorkgroupInfoElementNode = true;

	}

	generate( builder, output ) {

		let snippet;

		const isAssignContext = builder.isContextAssign();
		snippet = super.generate( builder );

		if ( isAssignContext !== true ) {

			const type = this.getNodeType( builder );

			snippet = builder.format( snippet, type, output );

		}

		// TODO: Possibly activate clip distance index on index access rather than from clipping context

		return snippet;

	}

}

/**
 * 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 transferred 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.
	 *
	 * @param {string} scope - The address space the buffer is scoped to. Currently only `'Workgroup'` is used.
	 * @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 ) {

		super( bufferType );

		/**
		 * The buffer type.
		 *
		 * @type {string}
		 */
		this.bufferType = bufferType;

		/**
		 * The buffer count.
		 *
		 * @type {number}
		 * @default 0
		 */
		this.bufferCount = bufferCount;

		/**
		 * Whether the node is atomic or not.
		 *
		 * @type {boolean}
		 * @default false
		 */
		this.isAtomic = false;

		/**
		 * This flag can be used for type testing.
		 *
		 * @type {boolean}
		 * @readonly
		 * @default true
		 */
		this.isWorkgroupInfoNode = true;

		/**
		 * The data type of the array buffer.
		 *
		 * @type {string}
		 */
		this.elementType = bufferType;

		/**
		 * The address space the buffer is scoped to. Currently only `'Workgroup'` is used.
		 *
		 * @type {string}
		 */
		this.scope = scope;

		/**
		 * The name of the workgroup scoped buffer.
		 *
		 * @type {string}
		 * @default ''
		 */
		this.name = '';

	}

	/**
	 * Sets the name of this node.
	 *
	 * @param {string} name - The name to set.
	 * @return {WorkgroupInfoNode} A reference to this node.
	 */
	setName( name ) {

		this.name = name;

		return this;

	}

	/**
	 * Sets the name/label of this node.
	 *
	 * @deprecated
	 * @param {string} name - The name to set.
	 * @return {WorkgroupInfoNode} A reference to this node.
	 */
	label( name ) {

		warn( 'TSL: "label()" has been deprecated. Use "setName()" instead.', new StackTrace() ); // @deprecated r179

		return this.setName( name );

	}

	/**
	 * Sets the scope of this node.
	 *
	 * @param {string} scope - The scope to set.
	 * @return {WorkgroupInfoNode} A reference to this node.
	 */
	setScope( scope ) {

		this.scope = scope;

		return this;

	}

	/**
	 * Defines whether the node is atomic or not.
	 *
	 * @param {boolean} value - The atomic flag.
	 * @return {WorkgroupInfoNode} A reference to this node.
	 */
	setAtomic( value ) {

		this.isAtomic = value;

		return this;

	}

	/**
	 * Convenience method for making this node atomic.
	 *
	 * @return {WorkgroupInfoNode} A reference to this node.
	 */
	toAtomic() {

		return this.setAtomic( true );

	}


	/**
	 * The data type of the array buffer.
	 *
	 * @return {string} The element type.
	 */
	getElementType() {

		return this.elementType;

	}

	/**
	 * Overwrites the default implementation since the input type
	 * is inferred from the scope.
	 *
	 * @param {NodeBuilder} builder - The current node builder.
	 * @return {string} The input type.
	 */
	getInputType( /*builder*/ ) {

		return `${this.scope}Array`;

	}

	/**
	 * This method can be used to access elements via an index node.
	 *
	 * @param {IndexNode} indexNode - indexNode.
	 * @return {WorkgroupInfoElementNode} A reference to an element.
	 */
	element( indexNode ) {

		return new WorkgroupInfoElementNode( this, indexNode );

	}

	generate( builder ) {

		if ( builder.shaderStage !== 'compute' ) {

			error( 'TSL: "workgroupArray()" can only be executed within the compute shader stage' );

		}

		const name = ( this.name !== '' ) ? this.name : `${this.scope}Array_${this.id}`;

		return builder.getScopedArray( name, this.scope.toLowerCase(), this.bufferType, this.bufferCount, this.isAtomic );

	}

}

export default WorkgroupInfoNode;

/**
 * TSL function for creating a workgroup info node.
 * Creates a new 'workgroup' scoped array buffer.
 *
 * @tsl
 * @function
 * @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 ) => new WorkgroupInfoNode( 'Workgroup', type, count );