WebGPURenderer: partial compute() shader support for WebGL backend (#27367)

* compute support

* code bot fixes

* add StorageBufferAttribute

* adapt to storagetBufferAttribute

* dynamically realign the size of storage buffer based on backend

* compute examples takes too much time to init for puppeteer

* Add StorageBufferAttribute.create()

---------

Co-authored-by: aardgoose <[email protected]>

Author: aardgoose

examples/jsm/nodes/accessors/BufferAttributeNode.js

@@ -67,12 +67,14 @@ class BufferAttributeNode extends InputNode {
 
 		const nodeType = this.getNodeType( builder );
 
-		const nodeUniform = builder.getBufferAttributeFromNode( this, nodeType );
-		const propertyName = builder.getPropertyName( nodeUniform );
+		const nodeAttribute = builder.getBufferAttributeFromNode( this, nodeType );
+		const propertyName = builder.getPropertyName( nodeAttribute );
 
 		let output = null;
 
-		if ( builder.shaderStage === 'vertex' ) {
+		if ( builder.shaderStage === 'vertex' || builder.shaderStage === 'compute' ) {
+
+			this.name = propertyName;
 
 			output = propertyName;
 

examples/jsm/nodes/accessors/StorageBufferNode.js

@@ -1,6 +1,8 @@
 import BufferNode from './BufferNode.js';
+import { bufferAttribute } from './BufferAttributeNode.js';
 import { addNodeClass } from '../core/Node.js';
 import { nodeObject } from '../shadernode/ShaderNode.js';
+import { varying } from '../core/VaryingNode.js';
 
 class StorageBufferNode extends BufferNode {
 
@@ -10,6 +12,9 @@ class StorageBufferNode extends BufferNode {
 
 		this.isStorageBufferNode = true;
 
+		this._attribute = null;
+		this._varying = null;
+
 	}
 
 	getInputType( /*builder*/ ) {
@@ -18,6 +23,28 @@ class StorageBufferNode extends BufferNode {
 
 	}
 
+	generate( builder ) {
+
+		if ( builder.isAvailable( 'storageBuffer' ) ) return super.generate( builder );
+
+		const nodeType = this.getNodeType( builder );
+
+		if ( this._attribute === null ) {
+
+			this._attribute = bufferAttribute( this.value );
+			this._varying = varying( this._attribute );
+
+		}
+
+
+		const output = this._varying.build( builder, nodeType );
+
+		builder.registerTransform( output, this._attribute );
+
+		return output;
+
+	}
+
 }
 
 export default StorageBufferNode;

examples/jsm/nodes/utils/ArrayElementNode.js

@@ -24,6 +24,12 @@ class ArrayElementNode extends Node { // @TODO: If extending from TempNode it br
 		const nodeSnippet = this.node.build( builder );
 		const indexSnippet = this.indexNode.build( builder, 'uint' );
 
+		if ( this.node.isStorageBufferNode && ! builder.isAvailable( 'storageBuffer' ) ) {
+
+			return nodeSnippet;
+
+		}
+
 		return `${nodeSnippet}[ ${indexSnippet} ]`;
 
 	}

examples/jsm/renderers/common/Backend.js

@@ -174,6 +174,12 @@ class Backend {
 
 	}
 
+	has( object ) {
+
+		return this.data.has( object );
+
+	}
+
 	delete( object ) {
 
 		this.data.delete( object );

examples/jsm/renderers/common/Bindings.js

@@ -48,7 +48,7 @@ class Bindings extends DataMap {
 
 			const nodeBuilderState = this.nodes.getForCompute( computeNode );
 
-			const bindings = nodeBuilderState.bindings.compute;
+			const bindings = nodeBuilderState.bindings;
 
 			data.bindings = bindings;
 

examples/jsm/renderers/common/Pipelines.js

@@ -42,18 +42,18 @@ class Pipelines extends DataMap {
 
 			// get shader
 
-			const nodeBuilder = this.nodes.getForCompute( computeNode );
+			const nodeBuilderState = this.nodes.getForCompute( computeNode );
 
 			// programmable stage
 
-			let stageCompute = this.programs.compute.get( nodeBuilder.computeShader );
+			let stageCompute = this.programs.compute.get( nodeBuilderState.computeShader );
 
 			if ( stageCompute === undefined ) {
 
 				if ( previousPipeline && previousPipeline.computeProgram.usedTimes === 0 ) this._releaseProgram( previousPipeline.computeProgram );
 
-				stageCompute = new ProgrammableStage( nodeBuilder.computeShader, 'compute' );
-				this.programs.compute.set( nodeBuilder.computeShader, stageCompute );
+				stageCompute = new ProgrammableStage( nodeBuilderState.computeShader, 'compute', nodeBuilderState.transforms, nodeBuilderState.nodeAttributes );
+				this.programs.compute.set( nodeBuilderState.computeShader, stageCompute );
 
 				backend.createProgram( stageCompute );
 

examples/jsm/renderers/common/ProgrammableStage.js

@@ -2,12 +2,14 @@ let _id = 0;
 
 class ProgrammableStage {
 
-	constructor( code, type ) {
+	constructor( code, type, transforms = null, attributes = null ) {
 
 		this.id = _id ++;
 
 		this.code = code;
 		this.stage = type;
+		this.transforms = transforms;
+		this.attributes = attributes;
 
 		this.usedTimes = 0;
 

examples/jsm/renderers/common/StorageBufferAttribute.js

@@ -0,0 +1,21 @@
+import { InstancedBufferAttribute } from 'three';
+
+class StorageBufferAttribute extends InstancedBufferAttribute {
+
+	constructor( array, itemSize ) {
+
+		super( array, itemSize );
+
+		this.isStorageBufferAttribute = true;
+
+	}
+
+	static create( count, itemSize, typeClass = Float32Array ) {
+
+		return new StorageBufferAttribute( new typeClass( count * itemSize ), itemSize );
+
+	}
+
+}
+
+export default StorageBufferAttribute;

examples/jsm/renderers/common/nodes/NodeBuilderState.js

@@ -1,10 +1,11 @@
 class NodeBuilderState {
 
-	constructor( vertexShader, fragmentShader, computeShader, nodeAttributes, bindings, updateNodes, updateBeforeNodes ) {
+	constructor( vertexShader, fragmentShader, computeShader, nodeAttributes, bindings, updateNodes, updateBeforeNodes, transforms = [] ) {
 
 		this.vertexShader = vertexShader;
 		this.fragmentShader = fragmentShader;
 		this.computeShader = computeShader;
+		this.transforms = transforms;
 
 		this.nodeAttributes = nodeAttributes;
 		this.bindings = bindings;

examples/jsm/renderers/common/nodes/Nodes.js

@@ -164,7 +164,7 @@ class Nodes extends DataMap {
 
 			nodeBuilderState = this._createNodeBuilderState( nodeBuilder );
 
-			computeData.nodeBuilderState = nodeBuilder;
+			computeData.nodeBuilderState = nodeBuilderState;
 
 		}
 
@@ -181,7 +181,8 @@ class Nodes extends DataMap {
 			nodeBuilder.getAttributesArray(),
 			nodeBuilder.getBindings(),
 			nodeBuilder.updateNodes,
-			nodeBuilder.updateBeforeNodes
+			nodeBuilder.updateBeforeNodes,
+			nodeBuilder.transforms
 		);
 
 	}