NodeMaterial: Add support for compute() integrated into the material (#30768)

sunag · web-flow · commit ba6fa5aac424 · 2025-03-21T12:31:51.000-03:00
* add `computeSkinning` and remove `skinningReference`

* Node: Show warning for recursive code generate

* fix `attributeName` undefined

* add compute() support for NodeMaterial

* update example using `computeSkinning`

* cleanup
diff --git a/examples/screenshots/webgpu_skinning_points.jpg b/examples/screenshots/webgpu_skinning_points.jpg
diff --git a/examples/webgpu_skinning_points.html b/examples/webgpu_skinning_points.html
@@ -9,7 +9,8 @@
 	<body>
 
 		<div id="info">
-			<a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> webgpu - skinning points
+			<a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> webgpu - skinning points</br>
+			colors and scale of the points are based on the speed of the skinning
 		</div>
 
 		<script type="importmap">
@@ -26,12 +27,11 @@
 		<script type="module">
 
 			import * as THREE from 'three';
-			import { uniform, skinning } from 'three/tsl';
+			import { color, computeSkinning, objectWorldMatrix, instancedArray, instanceIndex, Fn, shapeCircle } from 'three/tsl';
 
 			import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';
 
 			let camera, scene, renderer;
-
 			let mixer, clock;
 
 			init();
@@ -42,8 +42,11 @@
 				camera.position.set( 0, 300, - 85 );
 
 				scene = new THREE.Scene();
+				scene.background = new THREE.Color( 0x111111 );
 				camera.lookAt( 0, 0, - 85 );
 
+				scene.add( new THREE.AmbientLight( 0xffffff, 10 ) );
+
 				clock = new THREE.Clock();
 
 				const loader = new GLTFLoader();
@@ -61,17 +64,47 @@
 
 							child.visible = false;
 
+							const countOfPoints = child.geometry.getAttribute( 'position' ).count;
+
+							const pointPositionArray = instancedArray( countOfPoints, 'vec3' ).setPBO( true );
+							const pointSpeedArray = instancedArray( countOfPoints, 'vec3' ).setPBO( true );
+
+							const pointSpeedAttribute = pointSpeedArray.toAttribute();
+							const skinningPosition = computeSkinning( child );
+
 							const materialPoints = new THREE.PointsNodeMaterial();
-							materialPoints.colorNode = uniform( new THREE.Color() );
-							materialPoints.positionNode = skinning( child );
+							materialPoints.colorNode = pointSpeedAttribute.mul( .6 ).mix( color( 0x0066ff ), color( 0xff9000 ) );
+							materialPoints.opacityNode = shapeCircle();
+							materialPoints.sizeNode = pointSpeedAttribute.length().exp().min( 5 ).mul( 5 ).add( 1 );
+							materialPoints.sizeAttenuation = false;
+
+							materialPoints.positionNode = Fn( () => {
+
+								const pointPosition = pointPositionArray.element( instanceIndex );
+								const pointSpeed = pointSpeedArray.element( instanceIndex );
 
-							const pointCloud = new THREE.Points( child.geometry, materialPoints );
+								const skinningWorldPosition = objectWorldMatrix( child ).mul( skinningPosition );
+
+								const skinningSpeed = skinningWorldPosition.sub( pointPosition );
+
+								pointSpeed.assign( skinningSpeed );
+								pointPosition.assign( skinningWorldPosition );
+
+								return pointPositionArray.toAttribute();
+
+							} )().compute( countOfPoints );
+
+							const pointCloud = new THREE.Sprite( materialPoints );
+							pointCloud.count = countOfPoints;
 							scene.add( pointCloud );
 
 						}
 
 					} );
 
+					object.scale.set( 100, 100, 100 );
+					object.rotation.x = - Math.PI / 2;
+
 					scene.add( object );
 
 				} );
diff --git a/src/Three.TSL.js b/src/Three.TSL.js
@@ -116,6 +116,7 @@ export const color = TSL.color;
 export const colorSpaceToWorking = TSL.colorSpaceToWorking;
 export const colorToDirection = TSL.colorToDirection;
 export const compute = TSL.compute;
+export const computeSkinning = TSL.computeSkinning;
 export const cond = TSL.cond;
 export const Const = TSL.Const;
 export const context = TSL.context;
@@ -433,7 +434,6 @@ export const sign = TSL.sign;
 export const sin = TSL.sin;
 export const sinc = TSL.sinc;
 export const skinning = TSL.skinning;
-export const skinningReference = TSL.skinningReference;
 export const smoothstep = TSL.smoothstep;
 export const smoothstepElement = TSL.smoothstepElement;
 export const specularColor = TSL.specularColor;
diff --git a/src/materials/nodes/NodeMaterial.js b/src/materials/nodes/NodeMaterial.js
@@ -11,7 +11,7 @@ import { instancedMesh } from '../../nodes/accessors/InstancedMeshNode.js';
 import { batch } from '../../nodes/accessors/BatchNode.js';
 import { materialReference } from '../../nodes/accessors/MaterialReferenceNode.js';
 import { positionLocal, positionView } from '../../nodes/accessors/Position.js';
-import { skinningReference } from '../../nodes/accessors/SkinningNode.js';
+import { skinning } from '../../nodes/accessors/SkinningNode.js';
 import { morphReference } from '../../nodes/accessors/MorphNode.js';
 import { mix } from '../../nodes/math/MathNode.js';
 import { float, vec3, vec4 } from '../../nodes/tsl/TSLBase.js';
@@ -699,7 +699,7 @@ class NodeMaterial extends Material {
 
 		if ( object.isSkinnedMesh === true ) {
 
-			skinningReference( object ).append();
+			skinning( object ).append();
 
 		}
 
diff --git a/src/nodes/accessors/SkinningNode.js b/src/nodes/accessors/SkinningNode.js
@@ -10,6 +10,9 @@ import { tangentLocal } from './Tangent.js';
 import { uniform } from '../core/UniformNode.js';
 import { buffer } from './BufferNode.js';
 import { getDataFromObject } from '../core/NodeUtils.js';
+import { storage } from './StorageBufferNode.js';
+import { InstancedBufferAttribute } from '../../core/InstancedBufferAttribute.js';
+import { instanceIndex } from '../core/IndexNode.js';
 
 const _frameId = new WeakMap();
 
@@ -31,9 +34,8 @@ class SkinningNode extends Node {
 	 * Constructs a new skinning node.
 	 *
 	 * @param {SkinnedMesh} skinnedMesh - The skinned mesh.
-	 * @param {boolean} [useReference=false] - Whether to use reference nodes for internal skinned mesh related data or not.
 	 */
-	constructor( skinnedMesh, useReference = false ) {
+	constructor( skinnedMesh ) {
 
 		super( 'void' );
 
@@ -44,14 +46,6 @@ class SkinningNode extends Node {
 		 */
 		this.skinnedMesh = skinnedMesh;
 
-		/**
-		 * Whether to use reference nodes for internal skinned mesh related data or not.
-		 * TODO: Explain the purpose of the property.
-		 *
-		 * @type {boolean}
-		 */
-		this.useReference = useReference;
-
 		/**
 		 * The update type overwritten since skinning nodes are updated per object.
 		 *
@@ -75,42 +69,40 @@ class SkinningNode extends Node {
 		 */
 		this.skinWeightNode = attribute( 'skinWeight', 'vec4' );
 
-		let bindMatrixNode, bindMatrixInverseNode, boneMatricesNode;
-
-		if ( useReference ) {
-
-			bindMatrixNode = reference( 'bindMatrix', 'mat4' );
-			bindMatrixInverseNode = reference( 'bindMatrixInverse', 'mat4' );
-			boneMatricesNode = referenceBuffer( 'skeleton.boneMatrices', 'mat4', skinnedMesh.skeleton.bones.length );
-
-		} else {
-
-			bindMatrixNode = uniform( skinnedMesh.bindMatrix, 'mat4' );
-			bindMatrixInverseNode = uniform( skinnedMesh.bindMatrixInverse, 'mat4' );
-			boneMatricesNode = buffer( skinnedMesh.skeleton.boneMatrices, 'mat4', skinnedMesh.skeleton.bones.length );
-
-		}
-
 		/**
 		 * The bind matrix node.
 		 *
 		 * @type {Node<mat4>}
 		 */
-		this.bindMatrixNode = bindMatrixNode;
+		this.bindMatrixNode = reference( 'bindMatrix', 'mat4' );
 
 		/**
 		 * The bind matrix inverse node.
 		 *
 		 * @type {Node<mat4>}
 		 */
-		this.bindMatrixInverseNode = bindMatrixInverseNode;
+		this.bindMatrixInverseNode = reference( 'bindMatrixInverse', 'mat4' );
 
 		/**
 		 * The bind matrices as a uniform buffer node.
 		 *
 		 * @type {Node}
 		 */
-		this.boneMatricesNode = boneMatricesNode;
+		this.boneMatricesNode = referenceBuffer( 'skeleton.boneMatrices', 'mat4', skinnedMesh.skeleton.bones.length );
+
+		/**
+		 * The current vertex position in local space.
+		 *
+		 * @type {Node<vec3>}
+		 */
+		this.positionNode = positionLocal;
+
+		/**
+		 * The result of vertex position in local space.
+		 *
+		 * @type {Node<vec3>}
+		 */
+		this.toPositionNode = positionLocal;
 
 		/**
 		 * The previous bind matrices as a uniform buffer node.
@@ -127,10 +119,10 @@ class SkinningNode extends Node {
 	 * Transforms the given vertex position via skinning.
 	 *
 	 * @param {Node} [boneMatrices=this.boneMatricesNode] - The bone matrices
-	 * @param {Node<vec3>} [position=positionLocal] - The vertex position in local space.
+	 * @param {Node<vec3>} [position=this.positionNode] - The vertex position in local space.
 	 * @return {Node<vec3>} The transformed vertex position.
 	 */
-	getSkinnedPosition( boneMatrices = this.boneMatricesNode, position = positionLocal ) {
+	getSkinnedPosition( boneMatrices = this.boneMatricesNode, position = this.positionNode ) {
 
 		const { skinIndexNode, skinWeightNode, bindMatrixNode, bindMatrixInverseNode } = this;
 
@@ -225,6 +217,7 @@ class SkinningNode extends Node {
 	 * Setups the skinning node by assigning the transformed vertex data to predefined node variables.
 	 *
 	 * @param {NodeBuilder} builder - The current node builder.
+	 * @return {Node<vec3>} The transformed vertex position.
 	 */
 	setup( builder ) {
 
@@ -236,8 +229,9 @@ class SkinningNode extends Node {
 
 		const skinPosition = this.getSkinnedPosition();
 
+		if ( this.toPositionNode ) this.toPositionNode.assign( skinPosition );
 
-		positionLocal.assign( skinPosition );
+		//
 
 		if ( builder.hasGeometryAttribute( 'normal' ) ) {
 
@@ -253,6 +247,8 @@ class SkinningNode extends Node {
 
 		}
 
+		return skinPosition;
+
 	}
 
 	/**
@@ -266,7 +262,7 @@ class SkinningNode extends Node {
 
 		if ( output !== 'void' ) {
 
-			return positionLocal.build( builder, output );
+			return super.generate( builder, output );
 
 		}
 
@@ -279,8 +275,7 @@ class SkinningNode extends Node {
 	 */
 	update( frame ) {
 
-		const object = this.useReference ? frame.object : this.skinnedMesh;
-		const skeleton = object.skeleton;
+		const skeleton = frame.object && frame.object.skeleton ? frame.object.skeleton : this.skinnedMesh.skeleton;
 
 		if ( _frameId.get( skeleton ) === frame.frameId ) return;
 
@@ -307,11 +302,25 @@ export default SkinningNode;
 export const skinning = ( skinnedMesh ) => nodeObject( new SkinningNode( skinnedMesh ) );
 
 /**
- * TSL function for creating a skinning node with reference usage.
+ * TSL function for computing skinning.
  *
  * @tsl
  * @function
  * @param {SkinnedMesh} skinnedMesh - The skinned mesh.
+ * @param {Node<vec3>} [toPosition=null] - The target position.
  * @returns {SkinningNode}
  */
-export const skinningReference = ( skinnedMesh ) => nodeObject( new SkinningNode( skinnedMesh, true ) );
+export const computeSkinning = ( skinnedMesh, toPosition = null ) => {
+
+	const node = new SkinningNode( skinnedMesh );
+	node.positionNode = storage( new InstancedBufferAttribute( skinnedMesh.geometry.getAttribute( 'position' ).array, 3 ), 'vec3' ).setPBO( true ).toReadOnly().element( instanceIndex ).toVar();
+	node.skinIndexNode = storage( new InstancedBufferAttribute( new Uint32Array( skinnedMesh.geometry.getAttribute( 'skinIndex' ).array ), 4 ), 'uvec4' ).setPBO( true ).toReadOnly().element( instanceIndex ).toVar();
+	node.skinWeightNode = storage( new InstancedBufferAttribute( skinnedMesh.geometry.getAttribute( 'skinWeight' ).array, 4 ), 'vec4' ).setPBO( true ).toReadOnly().element( instanceIndex ).toVar();
+	node.bindMatrixNode = uniform( skinnedMesh.bindMatrix, 'mat4' );
+	node.bindMatrixInverseNode = uniform( skinnedMesh.bindMatrixInverse, 'mat4' );
+	node.boneMatricesNode = buffer( skinnedMesh.skeleton.boneMatrices, 'mat4', skinnedMesh.skeleton.bones.length );
+	node.toPositionNode = toPosition;
+
+	return nodeObject( node );
+
+};
diff --git a/src/nodes/core/Node.js b/src/nodes/core/Node.js
@@ -674,9 +674,21 @@ class Node extends EventDispatcher {
 
 				if ( result === undefined ) {
 
-					result = this.generate( builder ) || '';
+					if ( nodeData.generated === undefined ) {
 
-					nodeData.snippet = result;
+						nodeData.generated = true;
+
+						result = this.generate( builder ) || '';
+
+						nodeData.snippet = result;
+
+					} else {
+
+						console.warn( 'THREE.Node: Recursion detected.', this );
+
+						result = '';
+
+					}
 
 				} else if ( nodeData.flowCodes !== undefined && builder.context.nodeBlock !== undefined ) {
 
diff --git a/src/nodes/gpgpu/ComputeNode.js b/src/nodes/gpgpu/ComputeNode.js
@@ -163,17 +163,35 @@ class ComputeNode extends Node {
 
 	}
 
-	generate( builder ) {
+	setup( builder ) {
+
+		const result = this.computeNode.setup( builder );
+
+		const properties = builder.getNodeProperties( this );
+		properties.outputComputeNode = result.outputNode;
+
+		result.outputNode = null;
+
+		return result;
+
+	}
+
+	generate( builder, output ) {
 
 		const { shaderStage } = builder;
 
 		if ( shaderStage === 'compute' ) {
 
-			const snippet = this.computeNode.build( builder, 'void' );
+			this.computeNode.build( builder, 'void' );
+
+		} else {
+
+			const properties = builder.getNodeProperties( this );
+			const outputComputeNode = properties.outputComputeNode;
 
-			if ( snippet !== '' ) {
+			if ( outputComputeNode ) {
 
-				builder.addLineFlowCode( snippet, this );
+				return outputComputeNode.build( builder, output );
 
 			}
 
diff --git a/src/renderers/webgl-fallback/nodes/GLSLNodeBuilder.js b/src/renderers/webgl-fallback/nodes/GLSLNodeBuilder.js
@@ -775,7 +775,7 @@ ${ flowData.code }
 
 					const flat = type.includes( 'int' ) || type.includes( 'uv' ) || type.includes( 'iv' ) ? 'flat ' : '';
 
-					snippet += `${flat} out ${type} ${varying.name};\n`;
+					snippet += `${flat}out ${type} ${varying.name};\n`;
 
 				} else {
 
@@ -1073,10 +1073,9 @@ ${ flowData.code }
 		for ( let i = 0; i < transforms.length; i ++ ) {
 
 			const transform = transforms[ i ];
-
 			const attributeName = this.getPropertyName( transform.attributeNode );
 
-			snippet += `${ transform.varyingName } = ${ attributeName };\n\t`;
+			if ( attributeName ) snippet += `${ transform.varyingName } = ${ attributeName };\n\t`;
 
 		}
 
