WebGPURenderer: Introduce compute geometry (mrdoob#28072)

* introduce compute geometry * update * update names * update speed * add TODO * read-only
duzeng · Apr 5, 2024 · f21c099 · f21c099
1 parent 9870c85
commit f21c099
Show file tree

Hide file tree

Showing 5 changed files with 161 additions and 1 deletion.
diff --git a/examples/files.json b/examples/files.json
@@ -324,6 +324,7 @@
 		"webgpu_clearcoat",
 		"webgpu_clipping",
 		"webgpu_compute_audio",
+		"webgpu_compute_geometry",
 		"webgpu_compute_particles",
 		"webgpu_compute_particles_rain",
 		"webgpu_compute_particles_snow",

diff --git a/examples/jsm/nodes/accessors/StorageBufferNode.js b/examples/jsm/nodes/accessors/StorageBufferNode.js
@@ -18,6 +18,15 @@ class StorageBufferNode extends BufferNode {
 		this._attribute = null;
 		this._varying = null;
 
+		if ( value.isStorageBufferAttribute !== true && value.isStorageInstancedBufferAttribute !== true ) {
+
+			// TOOD: Improve it, possibly adding a new property to the BufferAttribute to identify it as a storage buffer read-only attribute in Renderer
+
+			if ( value.isInstancedBufferAttribute ) value.isStorageInstancedBufferAttribute = true;
+			else value.isStorageBufferAttribute = true;
+
+		}
+
 	}
 
 	getInputType( /*builder*/ ) {

diff --git a/examples/jsm/renderers/webgpu/utils/WebGPUAttributeUtils.js b/examples/jsm/renderers/webgpu/utils/WebGPUAttributeUtils.js
@@ -64,7 +64,6 @@ class WebGPUAttributeUtils {
 
 			if ( ( bufferAttribute.isStorageBufferAttribute || bufferAttribute.isStorageInstancedBufferAttribute ) && bufferAttribute.itemSize === 3 ) {
 
-				bufferAttribute.itemSize = 4;
 				array = new array.constructor( bufferAttribute.count * 4 );
 
 				for ( let i = 0; i < bufferAttribute.count; i ++ ) {
@@ -73,6 +72,10 @@ class WebGPUAttributeUtils {
 
 				}
 
+				// Update BufferAttribute
+				bufferAttribute.itemSize = 4;
+				bufferAttribute.array = array;
+
 			}
 
 			const size = array.byteLength + ( ( 4 - ( array.byteLength % 4 ) ) % 4 ); // ensure 4 byte alignment, see #20441

diff --git a/examples/screenshots/webgpu_compute_geometry.jpg b/examples/screenshots/webgpu_compute_geometry.jpg
diff --git a/examples/webgpu_compute_geometry.html b/examples/webgpu_compute_geometry.html
@@ -0,0 +1,147 @@
+<!DOCTYPE html>
+<html lang="en">
+	<head>
+		<title>three.js webgpu - compute geometry</title>
+		<meta charset="utf-8">
+		<meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
+		<link type="text/css" rel="stylesheet" href="main.css">
+	</head>
+	<body>
+
+		<div id="info">
+			<a href="https://threejs.org" target="_blank" rel="noopener">three.js webgpu</a> - compute geometry
+		</div>
+
+		<script type="importmap">
+			{
+				"imports": {
+					"three": "../build/three.module.js",
+					"three/addons/": "./jsm/",
+					"three/nodes": "./jsm/nodes/Nodes.js"
+				}
+			}
+		</script>
+
+		<script type="module">
+
+			import * as THREE from 'three';
+			import { vec3, cos, sin, mat3, storage, tslFn, instanceIndex, timerLocal } from 'three/nodes';
+
+			import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';
+
+			import WebGPURenderer from 'three/addons/renderers/webgpu/WebGPURenderer.js';
+
+			import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
+
+			import StorageBufferAttribute from 'three/addons/renderers/common/StorageBufferAttribute.js';
+
+			let camera, scene, renderer;
+			let computeUpdate;
+
+			init();
+
+			function init() {
+
+				camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 0.1, 10 );
+				camera.position.set( 0, 0, 1 );
+
+				scene = new THREE.Scene();
+				scene.background = new THREE.Color( 0x333333 );
+
+				new GLTFLoader().load( 'models/gltf/LeePerrySmith/LeePerrySmith.glb', function ( gltf ) {
+
+					const mesh = gltf.scene.children[ 0 ];
+					mesh.scale.setScalar( .1 );
+					mesh.material = new THREE.MeshNormalMaterial();
+					scene.add( mesh );
+
+					//
+
+					const positionBaseAttribute = mesh.geometry.attributes.position;
+					const normalBaseAttribute = mesh.geometry.attributes.normal;
+
+					// replace geometry attributes for storage buffer attributes
+
+					const positionStorageBufferAttribute = new StorageBufferAttribute( positionBaseAttribute.count, 4 );
+					const normalStorageBufferAttribute = new StorageBufferAttribute( normalBaseAttribute.count, 4 );
+
+					mesh.geometry.setAttribute( 'position', positionStorageBufferAttribute );
+					mesh.geometry.setAttribute( 'normal', normalStorageBufferAttribute );
+
+					// compute shader
+
+					const computeFn = tslFn( () => {
+
+						const positionAttribute = storage( positionBaseAttribute, 'vec3', positionBaseAttribute.count );
+						const normalAttribute = storage( normalBaseAttribute, 'vec3', normalBaseAttribute.count );
+
+						const positionStorageAttribute = storage( positionStorageBufferAttribute, 'vec4', positionStorageBufferAttribute.count );
+						const normalStorageAttribute = storage( normalStorageBufferAttribute, 'vec4', normalStorageBufferAttribute.count );
+
+						const time = timerLocal( 1 );
+						const scale = 0.3;
+
+						//
+
+						const position = vec3( positionAttribute.element( instanceIndex ) );
+						const normal = vec3( normalAttribute.element( instanceIndex ) );
+
+						const theta = sin( time.add( position.y ) ).mul( scale );
+
+						const c = cos( theta );
+						const s = sin( theta );
+
+						const m = mat3(
+							c, 0, s,
+							0, 1, 0,
+							s.negate(), 0, c
+						);
+
+						const transformed = position.mul( m );
+						const transformedNormal = normal.mul( m );
+
+						positionStorageAttribute.element( instanceIndex ).assign( transformed );
+						normalStorageAttribute.element( instanceIndex ).assign( transformedNormal );
+
+					} );
+
+					computeUpdate = computeFn().compute( positionBaseAttribute.count );
+
+				} );
+
+				// renderer
+
+				renderer = new WebGPURenderer( { antialias: true } );
+				renderer.setPixelRatio( window.devicePixelRatio );
+				renderer.setSize( window.innerWidth, window.innerHeight );
+				renderer.setAnimationLoop( animate );
+				document.body.appendChild( renderer.domElement );
+
+				const controls = new OrbitControls( camera, renderer.domElement );
+				controls.minDistance = .7;
+				controls.maxDistance = 2;
+
+				window.addEventListener( 'resize', onWindowResize );
+
+			}
+
+			function onWindowResize() {
+
+				camera.aspect = window.innerWidth / window.innerHeight;
+				camera.updateProjectionMatrix();
+
+				renderer.setSize( window.innerWidth, window.innerHeight );
+
+			}
+
+			async function animate() {
+
+				if ( computeUpdate ) await renderer.computeAsync( computeUpdate );
+
+				renderer.render( scene, camera );
+
+			}
+
+		</script>
+	</body>
+</html>