TSL: Introduce uniformFlow() (#31531)

* init branch

* fixes

* move to new approach, adjust existing conditional node flow to take advantage of more existing flow code functionality

* remove forceWebGL checks

* cleanup

* revert requested examples

* fix typo

* try build again

* fix spacing in select statement

* change

Author: cmhhelgeson

examples/webgpu_compute_cloth.html

@@ -561,4 +561,4 @@
 
 		</script>
 	</body>
-</html>
+</html>

examples/webgpu_compute_points.html

@@ -30,7 +30,7 @@
 			import Stats from 'stats-gl';
 
 
-			import { Fn, uniform, instancedArray, float, vec2, vec3, color, instanceIndex } from 'three/tsl';
+			import { Fn, uniform, instancedArray, float, vec2, color, instanceIndex } from 'three/tsl';
 
 			import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
 
@@ -76,7 +76,7 @@
 					const pointerSize = 0.1;
 					const distanceFromPointer = pointer.sub( position ).length();
 
-					particle.assign( distanceFromPointer.lessThanEqual( pointerSize ).select( vec3(), position ) );
+					particle.assign( distanceFromPointer.lessThanEqual( pointerSize ).select( vec2(), position ) );
 
 				} );
 

examples/webgpu_compute_sort_bitonic.html

@@ -383,7 +383,7 @@
 
 						} else {
 
-							nextAlgo.assign( nextBlockHeight.greaterThan( WORKGROUP_SIZE[ 0 ] * 2 ).select( StepType.DISPERSE_GLOBAL, StepType.DISPERSE_LOCAL ) );
+							nextAlgo.assign( nextBlockHeight.greaterThan( WORKGROUP_SIZE[ 0 ] * 2 ).select( StepType.DISPERSE_GLOBAL, StepType.DISPERSE_LOCAL ).uniformFlow() );
 
 						}
 

src/Three.TSL.js

@@ -547,6 +547,7 @@ export const uniform = TSL.uniform;
 export const uniformArray = TSL.uniformArray;
 export const uniformCubeTexture = TSL.uniformCubeTexture;
 export const uniformGroup = TSL.uniformGroup;
+export const uniformFlow = TSL.uniformFlow;
 export const uniformTexture = TSL.uniformTexture;
 export const unpremultiplyAlpha = TSL.unpremultiplyAlpha;
 export const userData = TSL.userData;

src/nodes/core/ContextNode.js

@@ -131,6 +131,16 @@ export default ContextNode;
  */
 export const context = /*@__PURE__*/ nodeProxy( ContextNode ).setParameterLength( 1, 2 );
 
+/**
+ * TSL function for defining a uniformFlow context value for a given node.
+ *
+ * @tsl
+ * @function
+ * @param {Node} node - The node whose dependencies should all execute within a uniform control-flow path.
+ * @returns {ContextNode}
+ */
+export const uniformFlow = ( node ) => context( node, { uniformFlow: true } );
+
 /**
  * TSL function for defining a name for the context value for a given node.
  *
@@ -162,4 +172,5 @@ export function label( node, name ) {
 
 addMethodChaining( 'context', context );
 addMethodChaining( 'label', label );
+addMethodChaining( 'uniformFlow', uniformFlow );
 addMethodChaining( 'setName', setName );

src/nodes/core/NodeBuilder.js

@@ -826,6 +826,22 @@ class NodeBuilder {
 
 	}
 
+	/**
+	 * Returns the native snippet for a ternary operation. E.g. GLSL would output
+	 * a ternary op as `cond ? x : y` whereas WGSL would output it as `select(y, x, cond)`
+	 *
+	 * @abstract
+	 * @param {string} condSnippet - The condition determining which expression gets resolved.
+	 * @param {string} ifSnippet - The expression to resolve to if the condition is true.
+	 * @param {string} elseSnippet - The expression to resolve to if the condition is false.
+	 * @return {string} The resolved method name.
+	 */
+	getTernary( /* condSnippet, ifSnippet, elseSnippet*/ ) {
+
+		return null;
+
+	}
+
 	/**
 	 * Returns a node for the given hash, see {@link NodeBuilder#setHashNode}.
 	 *

src/nodes/math/ConditionalNode.js

@@ -112,10 +112,12 @@ class ConditionalNode extends Node {
 
 		//
 
+		const isUniformFlow = builder.context.uniformFlow;
+
 		const properties = builder.getNodeProperties( this );
 		properties.condNode = condNode;
-		properties.ifNode = ifNode.context( { nodeBlock: ifNode } );
-		properties.elseNode = elseNode ? elseNode.context( { nodeBlock: elseNode } ) : null;
+		properties.ifNode = isUniformFlow ? ifNode : ifNode.context( { nodeBlock: ifNode } );
+		properties.elseNode = elseNode ? ( isUniformFlow ? elseNode : elseNode.context( { nodeBlock: elseNode } ) ) : null;
 
 	}
 
@@ -140,6 +142,20 @@ class ConditionalNode extends Node {
 		nodeData.nodeProperty = nodeProperty;
 
 		const nodeSnippet = condNode.build( builder, 'bool' );
+		const isUniformFlow = builder.context.uniformFlow;
+
+		if ( isUniformFlow && elseNode !== null ) {
+
+			const ifSnippet = ifNode.build( builder, type );
+			const elseSnippet = elseNode.build( builder, type );
+
+			const mathSnippet = builder.getTernary( nodeSnippet, ifSnippet, elseSnippet );
+
+			// TODO: If node property already exists return something else
+
+			return builder.format( mathSnippet, type, output );
+
+		}
 
 		builder.addFlowCode( `\n${ builder.tab }if ( ${ nodeSnippet } ) {\n\n` ).addFlowTab();
 

src/renderers/webgl-fallback/nodes/GLSLNodeBuilder.js

@@ -134,6 +134,20 @@ class GLSLNodeBuilder extends NodeBuilder {
 
 	}
 
+	/**
+	 * Returns the native snippet for a ternary operation.
+	 *
+	 * @param {string} condSnippet - The condition determining which expression gets resolved.
+	 * @param {string} ifSnippet - The expression to resolve to if the condition is true.
+	 * @param {string} elseSnippet - The expression to resolve to if the condition is false.
+	 * @return {string} The resolved method name.
+	 */
+	getTernary( condSnippet, ifSnippet, elseSnippet ) {
+
+		return `${condSnippet} ? ${ifSnippet} : ${elseSnippet}`;
+
+	}
+
 	/**
 	 * Returns the output struct name. Not relevant for GLSL.
 	 *

src/renderers/webgpu/nodes/WGSLNodeBuilder.js

@@ -1899,6 +1899,21 @@ ${ flowData.code }
 
 	}
 
+	/**
+	 * Returns the native snippet for a ternary operation.
+	 *
+	 * @param {string} condSnippet - The condition determining which expression gets resolved.
+	 * @param {string} ifSnippet - The expression to resolve to if the condition is true.
+	 * @param {string} elseSnippet - The expression to resolve to if the condition is false.
+	 * @return {string} The resolved method name.
+	 */
+	getTernary( condSnippet, ifSnippet, elseSnippet ) {
+
+		return `select( ${elseSnippet}, ${ifSnippet}, ${condSnippet} )`;
+
+	}
+
+
 	/**
 	 * Returns the WGSL type of the given node data type.
 	 *