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

SubgroupFunctionNode.js

import { error } from '../../utils.js';
import TempNode from '../core/TempNode.js';
import { nodeProxyIntent } from '../tsl/TSLCore.js';


/**
 * This class represents a set of built in WGSL shader functions that sync
 * synchronously execute an operation across a subgroup, or 'warp', of compute
 * or fragment shader invocations within a workgroup. Typically, these functions
 * will synchronously execute an operation using data from all active invocations
 * within the subgroup, then broadcast that result to all active invocations. In
 * other graphics APIs, subgroup functions are also referred to as wave intrinsics
 * (DirectX/HLSL) or warp intrinsics (CUDA).
 *
 * @augments TempNode
 */
class SubgroupFunctionNode extends TempNode {

	static get type() {

		return 'SubgroupFunctionNode';

	}

	/**
	 * Constructs a new function node.
	 *
	 * @param {string} method - The subgroup/wave intrinsic method to construct.
	 * @param {Node} [aNode=null] - The method's first argument.
	 * @param {Node} [bNode=null] - The method's second argument.
	 */
	constructor( method, aNode = null, bNode = null ) {

		super();

		/**
		 * The subgroup/wave intrinsic method to construct.
		 *
		 * @type {string}
		 */
		this.method = method;

		/**
		 * The method's first argument.
		 *
		 * @type {Node}
		 */
		this.aNode = aNode;

		/**
		 * The method's second argument.
		 *
		 * @type {Node}
		 */
		this.bNode = bNode;

	}

	getInputType( builder ) {

		const aType = this.aNode ? this.aNode.getNodeType( builder ) : null;
		const bType = this.bNode ? this.bNode.getNodeType( builder ) : null;

		const aLen = builder.isMatrix( aType ) ? 0 : builder.getTypeLength( aType );
		const bLen = builder.isMatrix( bType ) ? 0 : builder.getTypeLength( bType );

		if ( aLen > bLen ) {

			return aType;

		} else {

			return bType;

		}

	}

	generateNodeType( builder ) {

		const method = this.method;

		if ( method === SubgroupFunctionNode.SUBGROUP_ELECT ) {

			return 'bool';

		} else if ( method === SubgroupFunctionNode.SUBGROUP_BALLOT ) {

			return 'uvec4';

		} else {

			return this.getInputType( builder );

		}

	}

	generate( builder, output ) {

		const method = this.method;

		if ( builder.shaderStage === 'vertex' ) {

			error( `TSL: "${this.method}" is not supported in the vertex shader stage.` );

		}

		const type = this.getNodeType( builder );
		const inputType = this.getInputType( builder );

		const a = this.aNode;
		const b = this.bNode;

		const params = [];

		if (
			method === SubgroupFunctionNode.SUBGROUP_BROADCAST ||
			method === SubgroupFunctionNode.SUBGROUP_SHUFFLE ||
			method === SubgroupFunctionNode.QUAD_BROADCAST
		) {

			const bType = b.getNodeType( builder );

			params.push(
				a.build( builder, type ),
				b.build( builder, bType === 'float' ? 'int' : type )
			);

		} else if (
			method === SubgroupFunctionNode.SUBGROUP_SHUFFLE_XOR ||
			method === SubgroupFunctionNode.SUBGROUP_SHUFFLE_DOWN ||
			method === SubgroupFunctionNode.SUBGROUP_SHUFFLE_UP
		) {

			params.push(
				a.build( builder, type ),
				b.build( builder, 'uint' )
			);

		} else {

			if ( a !== null ) params.push( a.build( builder, inputType ) );
			if ( b !== null ) params.push( b.build( builder, inputType ) );

		}

		const paramsString = params.length === 0 ? '()' : `( ${params.join( ', ' )} )`;

		return builder.format( `${ builder.getMethod( method, type ) }${paramsString}`, type, output );



	}

	serialize( data ) {

		super.serialize( data );

		data.method = this.method;

	}

	deserialize( data ) {

		super.deserialize( data );

		this.method = data.method;

	}

	// 0 inputs
	static get SUBGROUP_ELECT() {

		return 'subgroupElect';

	}

	// 1 input
	static get SUBGROUP_BALLOT() {

		return 'subgroupBallot';

	}

	static get SUBGROUP_ADD() {

		return 'subgroupAdd';

	}

	static get SUBGROUP_INCLUSIVE_ADD() {

		return 'subgroupInclusiveAdd';

	}

	static get SUBGROUP_EXCLUSIVE_ADD() {

		return 'subgroupExclusiveAdd';

	}

	static get SUBGROUP_MUL() {

		return 'subgroupMul';

	}

	static get SUBGROUP_INCLUSIVE_MUL() {

		return 'subgroupInclusiveMul';

	}

	static get SUBGROUP_EXCLUSIVE_MUL() {

		return 'subgroupExclusiveMul';

	}

	static get SUBGROUP_AND() {

		return 'subgroupAnd';

	}

	static get SUBGROUP_OR() {

		return 'subgroupOr';

	}

	static get SUBGROUP_XOR() {

		return 'subgroupXor';

	}

	static get SUBGROUP_MIN() {

		return 'subgroupMin';

	}

	static get SUBGROUP_MAX() {

		return 'subgroupMax';

	}

	static get SUBGROUP_ALL() {

		return 'subgroupAll';

	}

	static get SUBGROUP_ANY() {

		return 'subgroupAny';

	}

	static get SUBGROUP_BROADCAST_FIRST() {

		return 'subgroupBroadcastFirst';

	}

	static get QUAD_SWAP_X() {

		return 'quadSwapX';

	}

	static get QUAD_SWAP_Y() {

		return 'quadSwapY';

	}

	static get QUAD_SWAP_DIAGONAL() {

		return 'quadSwapDiagonal';

	}

	// 2 inputs
	static get SUBGROUP_BROADCAST() {

		return 'subgroupBroadcast';

	}

	static get SUBGROUP_SHUFFLE() {

		return 'subgroupShuffle';

	}

	static get SUBGROUP_SHUFFLE_XOR() {

		return 'subgroupShuffleXor';

	}

	static get SUBGROUP_SHUFFLE_UP() {

		return 'subgroupShuffleUp';

	}

	static get SUBGROUP_SHUFFLE_DOWN() {

		return 'subgroupShuffleDown';

	}

	static get QUAD_BROADCAST() {

		return 'quadBroadcast';

	}

}

export default SubgroupFunctionNode;

/**
 * Returns true if this invocation has the lowest subgroup_invocation_id
 * among active invocations in the subgroup.
 *
 * @tsl
 * @method
 * @return {bool} The result of the computation.
 */
export const subgroupElect = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_ELECT ).setParameterLength( 0 );

/**
 * Returns a set of bitfields where the bit corresponding to subgroup_invocation_id
 * is 1 if pred is true for that active invocation and 0 otherwise.
 *
 * @tsl
 * @method
 * @param {bool} pred - A boolean that sets the bit corresponding to the invocations subgroup invocation id.
 * @return {vec4<u32>}- A bitfield corresponding to the pred value of each subgroup invocation.
 */
export const subgroupBallot = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_BALLOT ).setParameterLength( 1 );

/**
 * A reduction that adds e among all active invocations and returns that result.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the reduction by the current invocation.
 * @return {number} The accumulated result of the reduction operation.
 */
export const subgroupAdd = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_ADD ).setParameterLength( 1 );

/**
 * An inclusive scan returning the sum of e for all active invocations with subgroup_invocation_id less than or equal to this invocation.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the inclusive scan by the current invocation.
 * @return {number} The accumulated result of the inclusive scan operation.
 */
export const subgroupInclusiveAdd = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_INCLUSIVE_ADD ).setParameterLength( 1 );

/**
 * An exclusive scan that returns the sum of e for all active invocations with subgroup_invocation_id less than this invocation.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the exclusive scan by the current invocation.
 * @return {number} The accumulated result of the exclusive scan operation.
 */
export const subgroupExclusiveAdd = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_EXCLUSIVE_ADD ).setParameterLength( 1 );

/**
 * A reduction that multiplies e among all active invocations and returns that result.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the reduction by the current invocation.
 * @return {number} The accumulated result of the reduction operation.
 */
export const subgroupMul = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_MUL ).setParameterLength( 1 );

/**
 * An inclusive scan returning the product of e for all active invocations with subgroup_invocation_id less than or equal to this invocation.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the inclusive scan by the current invocation.
 * @return {number} The accumulated result of the inclusive scan operation.
 */
export const subgroupInclusiveMul = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_INCLUSIVE_MUL ).setParameterLength( 1 );

/**
 * An exclusive scan that returns the product of e for all active invocations with subgroup_invocation_id less than this invocation.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the exclusive scan by the current invocation.
 * @return {number} The accumulated result of the exclusive scan operation.
 */
export const subgroupExclusiveMul = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_EXCLUSIVE_MUL ).setParameterLength( 1 );

/**
 * A reduction that performs a bitwise and of e among all active invocations and returns that result.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the reduction by the current invocation.
 * @return {number} The result of the reduction operation.
 */
export const subgroupAnd = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_AND ).setParameterLength( 1 );

/**
 * A reduction that performs a bitwise or of e among all active invocations and returns that result.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the reduction by the current invocation.
 * @return {number} The result of the reduction operation.
 */
export const subgroupOr = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_OR ).setParameterLength( 1 );

/**
 * A reduction that performs a bitwise xor of e among all active invocations and returns that result.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the reduction by the current invocation.
 * @return {number} The result of the reduction operation.
 */
export const subgroupXor = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_XOR ).setParameterLength( 1 );

/**
 * A reduction that performs a min of e among all active invocations and returns that result.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the reduction by the current invocation.
 * @return {number} The result of the reduction operation.
 */
export const subgroupMin = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_MIN ).setParameterLength( 1 );

/**
 * A reduction that performs a max of e among all active invocations and returns that result.
 *
 * @tsl
 * @method
 * @param {number} e - The value provided to the reduction by the current invocation.
 * @return {number} The result of the reduction operation.
 */
export const subgroupMax = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_MAX ).setParameterLength( 1 );

/**
 * Returns true if e is true for all active invocations in the subgroup.
 *
 * @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( 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( 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.
 * @return {number} The broadcast value.
 */
export const subgroupBroadcastFirst = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_BROADCAST_FIRST ).setParameterLength( 1 );

/**
 * Swaps e between invocations in the quad in the X direction.
 *
 * @tsl
 * @method
 * @param {number} e - The value to swap from the current invocation.
 * @return {number} The value received from the swap operation.
 */
export const quadSwapX = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.QUAD_SWAP_X ).setParameterLength( 1 );

/**
 * Swaps e between invocations in the quad in the Y direction.
 *
 * @tsl
 * @method
 * @param {number} e - The value to swap from the current invocation.
 * @return {number} The value received from the swap operation.
 */
export const quadSwapY = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.QUAD_SWAP_Y ).setParameterLength( 1 );

/**
 * Swaps e between invocations in the quad diagonally.
 *
 * @tsl
 * @method
 * @param {number} e - The value to swap from the current invocation.
 * @return {number} The value received from the swap operation.
 */
export const quadSwapDiagonal = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.QUAD_SWAP_DIAGONAL ).setParameterLength( 1 );

/**
 * Broadcasts e from the invocation whose subgroup_invocation_id matches id, to all active invocations.
 *
 * @tsl
 * @method
 * @param {number} e - The value to broadcast from subgroup invocation 'id'.
 * @param {number} id - The subgroup invocation to broadcast from.
 * @return {number} The broadcast value.
 */
export const subgroupBroadcast = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_BROADCAST ).setParameterLength( 2 );

/**
 * Returns v from the active invocation whose subgroup_invocation_id matches id
 *
 * @tsl
 * @method
 * @param {number} v - The value to return from subgroup invocation id^mask.
 * @param {number} id - The subgroup invocation which returns the value v.
 * @return {number} The broadcast value.
 */
export const subgroupShuffle = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_SHUFFLE ).setParameterLength( 2 );

/**
 * Returns v from the active invocation whose subgroup_invocation_id matches subgroup_invocation_id ^ mask.
 *
 * @tsl
 * @method
 * @param {number} v - The value to return from subgroup invocation id^mask.
 * @param {number} mask - A bitmask that determines the target invocation via a XOR operation.
 * @return {number} The broadcast value.
 */
export const subgroupShuffleXor = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_SHUFFLE_XOR ).setParameterLength( 2 );

/**
 * Returns v from the active invocation whose subgroup_invocation_id matches subgroup_invocation_id - delta
 *
 * @tsl
 * @method
 * @param {number} v - The value to return from subgroup invocation id^mask.
 * @param {number} delta - A value that offsets the current in.
 * @return {number} The broadcast value.
 */
export const subgroupShuffleUp = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_SHUFFLE_UP ).setParameterLength( 2 );

/**
 * Returns v from the active invocation whose subgroup_invocation_id matches subgroup_invocation_id + delta
 *
 * @tsl
 * @method
 * @param {number} v - The value to return from subgroup invocation id^mask.
 * @param {number} delta - A value that offsets the current subgroup invocation.
 * @return {number} The broadcast value.
 */
export const subgroupShuffleDown = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.SUBGROUP_SHUFFLE_DOWN ).setParameterLength( 2 );

/**
 * Broadcasts e from the quad invocation with id equal to id.
 *
 * @tsl
 * @method
 * @param {number} e - The value to broadcast.
 * @return {number} The broadcast value.
 */
export const quadBroadcast = /*@__PURE__*/ nodeProxyIntent( SubgroupFunctionNode, SubgroupFunctionNode.QUAD_BROADCAST ).setParameterLength( 1 );