webgpu_compute_birds.html 16 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508
  1. <!DOCTYPE html>
  2. <html lang="en">
  3. <head>
  4. <title>three.js webgpu - compute - flocking</title>
  5. <meta charset="utf-8">
  6. <meta name="viewport" content="width=device-width, user-scalable=no, minimum-scale=1.0, maximum-scale=1.0">
  7. <link type="text/css" rel="stylesheet" href="main.css">
  8. <style>
  9. body {
  10. background-color: #fff;
  11. color: #444;
  12. }
  13. a {
  14. color:#08f;
  15. }
  16. </style>
  17. </head>
  18. <body>
  19. <div id="info">
  20. <a href="https://threejs.org" target="_blank" rel="noopener">three.js</a> - webgpu compute birds<br/>
  21. Move mouse to disturb birds.
  22. </div>
  23. <script type="importmap">
  24. {
  25. "imports": {
  26. "three": "../build/three.webgpu.js",
  27. "three/tsl": "../build/three.webgpu.js",
  28. "three/addons/": "./jsm/"
  29. }
  30. }
  31. </script>
  32. <script type="module">
  33. import * as THREE from 'three';
  34. import { uniform, varying, vec4, add, sub, max, dot, sin, mat3, uint, negate, cameraProjectionMatrix, cameraViewMatrix, positionLocal, modelWorldMatrix, sqrt, attribute, property, float, storage, storageObject, Fn, If, cos, Loop, Continue, normalize, instanceIndex, length } from 'three/tsl';
  35. import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
  36. import Stats from 'three/addons/libs/stats.module.js';
  37. import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
  38. let container, stats;
  39. let camera, scene, renderer;
  40. let last = performance.now();
  41. let pointer, raycaster;
  42. let computeVelocity, computePosition, effectController;
  43. const BIRDS = 16384;
  44. const SPEED_LIMIT = 9.0;
  45. const BOUNDS = 800, BOUNDS_HALF = BOUNDS / 2;
  46. // Custom Geometry - using 3 triangles each. No normals currently.
  47. class BirdGeometry extends THREE.BufferGeometry {
  48. constructor() {
  49. super();
  50. const trianglesPerBird = 3;
  51. const triangles = BIRDS * trianglesPerBird;
  52. const points = triangles * 3;
  53. const vertices = new THREE.BufferAttribute( new Float32Array( points * 3 ), 3 );
  54. const references = new THREE.BufferAttribute( new Uint32Array( points ), 1 );
  55. const birdVertex = new THREE.BufferAttribute( new Uint32Array( points ), 1 );
  56. this.setAttribute( 'position', vertices );
  57. this.setAttribute( 'reference', references );
  58. this.setAttribute( 'birdVertex', birdVertex );
  59. let v = 0;
  60. function verts_push() {
  61. for ( let i = 0; i < arguments.length; i ++ ) {
  62. vertices.array[ v ++ ] = arguments[ i ];
  63. }
  64. }
  65. const wingsSpan = 20;
  66. for ( let f = 0; f < BIRDS; f ++ ) {
  67. // Body
  68. verts_push(
  69. 0, 0, - 20,
  70. 0, - 8, 10,
  71. 0, 0, 30
  72. );
  73. // Wings
  74. verts_push(
  75. 0, 0, - 15,
  76. - wingsSpan, 0, 5,
  77. 0, 0, 15
  78. );
  79. verts_push(
  80. 0, 0, 15,
  81. wingsSpan, 0, 5,
  82. 0, 0, - 15
  83. );
  84. }
  85. for ( let v = 0; v < triangles * 3; v ++ ) {
  86. const triangleIndex = ~ ~ ( v / 3 );
  87. const birdIndex = ~ ~ ( triangleIndex / trianglesPerBird );
  88. references.array[ v ] = birdIndex;
  89. birdVertex.array[ v ] = v % 9;
  90. }
  91. this.scale( 0.2, 0.2, 0.2 );
  92. }
  93. }
  94. function init() {
  95. container = document.createElement( 'div' );
  96. document.body.appendChild( container );
  97. camera = new THREE.PerspectiveCamera( 50, window.innerWidth / window.innerHeight, 1, 5000 );
  98. camera.position.z = 1000;
  99. scene = new THREE.Scene();
  100. scene.fog = new THREE.Fog( 0xffffff, 700, 3000 );
  101. // Pointer
  102. pointer = new THREE.Vector2();
  103. raycaster = new THREE.Raycaster();
  104. // Sky
  105. const geometry = new THREE.IcosahedronGeometry( 1, 6 );
  106. const material = new THREE.MeshBasicNodeMaterial( {
  107. // Use vertex positions to create atmosphere colors
  108. colorNode: varying(
  109. vec4(
  110. sub( 0.25, positionLocal.y ),
  111. sub( - 0.25, positionLocal.y ),
  112. add( 1.5, positionLocal.y ),
  113. 1.0
  114. )
  115. ),
  116. side: THREE.BackSide
  117. } );
  118. const mesh = new THREE.Mesh( geometry, material );
  119. mesh.rotation.z = 0.75;
  120. mesh.scale.setScalar( 1200 );
  121. scene.add( mesh );
  122. //
  123. renderer = new THREE.WebGPURenderer( { antialiasing: true } );
  124. renderer.setPixelRatio( window.devicePixelRatio );
  125. renderer.setSize( window.innerWidth, window.innerHeight );
  126. renderer.setAnimationLoop( animate );
  127. renderer.toneMapping = THREE.NeutralToneMapping;
  128. container.appendChild( renderer.domElement );
  129. const controls = new OrbitControls( camera, renderer.domElement );
  130. controls.connect( /* renderer.domElement */ );
  131. // Initialize position, velocity, and phase values
  132. const positionArray = new Float32Array( BIRDS * 3 );
  133. const velocityArray = new Float32Array( BIRDS * 3 );
  134. const phaseArray = new Float32Array( BIRDS );
  135. for ( let i = 0; i < BIRDS; i ++ ) {
  136. const posX = Math.random() * BOUNDS - BOUNDS_HALF;
  137. const posY = Math.random() * BOUNDS - BOUNDS_HALF;
  138. const posZ = Math.random() * BOUNDS - BOUNDS_HALF;
  139. positionArray[ i * 3 + 0 ] = posX;
  140. positionArray[ i * 3 + 1 ] = posY;
  141. positionArray[ i * 3 + 2 ] = posZ;
  142. const velX = Math.random() - 0.5;
  143. const velY = Math.random() - 0.5;
  144. const velZ = Math.random() - 0.5;
  145. velocityArray[ i * 3 + 0 ] = velX * 10;
  146. velocityArray[ i * 3 + 1 ] = velY * 10;
  147. velocityArray[ i * 3 + 2 ] = velZ * 10;
  148. phaseArray[ i ] = 1;
  149. }
  150. // Create storage buffer attributes.
  151. const positionBufferAttribute = new THREE.StorageBufferAttribute( positionArray, 3 );
  152. const velocityBufferAttribute = new THREE.StorageBufferAttribute( velocityArray, 3 );
  153. const phaseBufferAttribute = new THREE.StorageBufferAttribute( phaseArray, 1 );
  154. // Labels applied to storage nodes and uniform nodes are reflected within the shader output,
  155. // and are useful for debugging purposes.
  156. // Access storage buffer attribute data from within shaders with a StorageNode.
  157. const positionStorage = storage( positionBufferAttribute, 'vec3', positionBufferAttribute.count ).label( 'positionStorage' );
  158. const velocityStorage = storage( velocityBufferAttribute, 'vec3', velocityBufferAttribute.count ).label( 'velocityStorage' );
  159. const phaseStorage = storage( phaseBufferAttribute, 'float', phaseBufferAttribute.count ).label( 'phaseStorage' );
  160. // Create read-only storage nodes. Storage nodes can only be accessed outside of compute shaders in a read-only state.
  161. const positionRead = storageObject( positionBufferAttribute, 'vec3', positionBufferAttribute.count ).toReadOnly();
  162. const velocityRead = storageObject( velocityBufferAttribute, 'vec3', velocityBufferAttribute.count ).toReadOnly();
  163. const phaseRead = storageObject( phaseBufferAttribute, 'float', phaseBufferAttribute.count ).toReadOnly();
  164. // Define Uniforms. Uniforms only need to be defined once rather than per shader.
  165. effectController = {
  166. separation: uniform( 15.0 ).label( 'separation' ),
  167. alignment: uniform( 20.0 ).label( 'alignment' ),
  168. cohesion: uniform( 20.0 ).label( 'cohesion' ),
  169. freedom: uniform( 0.75 ).label( 'freedom' ),
  170. now: uniform( 0.0 ),
  171. deltaTime: uniform( 0.0 ).label( 'deltaTime' ),
  172. rayOrigin: uniform( new THREE.Vector3() ).label( 'rayOrigin' ),
  173. rayDirection: uniform( new THREE.Vector3() ).label( 'rayDirection' )
  174. };
  175. // Create geometry
  176. const birdGeometry = new BirdGeometry();
  177. const birdMaterial = new THREE.NodeMaterial();
  178. // Animate bird mesh within vertex shader, then apply position offset to vertices.
  179. const birdVertexTSL = Fn( () => {
  180. const reference = attribute( 'reference' );
  181. const birdVertex = attribute( 'birdVertex' );
  182. const position = positionLocal.toVar();
  183. const newPhase = phaseRead.element( reference ).toVar();
  184. const newVelocity = normalize( velocityRead.element( reference ) ).toVar();
  185. If( birdVertex.equal( 4 ).or( birdVertex.equal( 7 ) ), () => {
  186. // flap wings
  187. position.y = sin( newPhase ).mul( 5.0 );
  188. } );
  189. const newPosition = modelWorldMatrix.mul( position );
  190. newVelocity.z.mulAssign( - 1.0 );
  191. const xz = length( newVelocity.xz );
  192. const xyz = float( 1.0 );
  193. const x = sqrt( ( newVelocity.y.mul( newVelocity.y ) ).oneMinus() );
  194. const cosry = newVelocity.x.div( xz ).toVar();
  195. const sinry = newVelocity.z.div( xz ).toVar();
  196. const cosrz = x.div( xyz );
  197. const sinrz = newVelocity.y.div( xyz ).toVar();
  198. // Nodes must be negated with negate(). Using '-', their values will resolve to NaN.
  199. const maty = mat3(
  200. cosry, 0, negate( sinry ),
  201. 0, 1, 0,
  202. sinry, 0, cosry
  203. );
  204. const matz = mat3(
  205. cosrz, sinrz, 0,
  206. negate( sinrz ), cosrz, 0,
  207. 0, 0, 1
  208. );
  209. const finalVert = maty.mul( matz ).mul( newPosition );
  210. finalVert.addAssign( positionRead.element( reference ) );
  211. return cameraProjectionMatrix.mul( cameraViewMatrix ).mul( finalVert );
  212. } );
  213. birdMaterial.vertexNode = birdVertexTSL();
  214. birdMaterial.side = THREE.DoubleSide;
  215. const birdMesh = new THREE.Mesh( birdGeometry, birdMaterial );
  216. birdMesh.rotation.y = Math.PI / 2;
  217. birdMesh.matrixAutoUpdate = false;
  218. birdMesh.updateMatrix();
  219. // Define GPU Compute shaders.
  220. // Shaders are computationally identical to their GLSL counterparts outside of texture destructuring.
  221. computeVelocity = Fn( () => {
  222. // Define consts
  223. const PI = float( 3.141592653589793 );
  224. const PI_2 = PI.mul( 2.0 );
  225. const limit = property( 'float', 'limit' ).assign( SPEED_LIMIT );
  226. // Destructure uniforms
  227. const { alignment, separation, cohesion, deltaTime, rayOrigin, rayDirection } = effectController;
  228. const zoneRadius = separation.add( alignment ).add( cohesion );
  229. const separationThresh = separation.div( zoneRadius );
  230. const alignmentThresh = ( separation.add( alignment ) ).div( zoneRadius );
  231. const zoneRadiusSq = zoneRadius.mul( zoneRadius );
  232. const position = positionStorage.element( instanceIndex );
  233. const velocity = velocityStorage.element( instanceIndex );
  234. // Add influence of pointer position to velocity.
  235. const directionToRay = rayOrigin.sub( position );
  236. const projectionLength = dot( directionToRay, rayDirection );
  237. const closestPoint = rayOrigin.sub( rayDirection.mul( projectionLength ) );
  238. const directionToClosestPoint = closestPoint.sub( position );
  239. const distanceToClosestPoint = length( directionToClosestPoint );
  240. const distanceToClosestPointSq = distanceToClosestPoint.mul( distanceToClosestPoint );
  241. const rayRadius = float( 150.0 );
  242. const rayRadiusSq = rayRadius.mul( rayRadius );
  243. If( distanceToClosestPointSq.lessThan( rayRadiusSq ), () => {
  244. // Scale bird velocity inversely with distance from prey radius center.
  245. const velocityAdjust = ( distanceToClosestPointSq.div( rayRadiusSq ).sub( 1.0 ) ).mul( deltaTime ).mul( 100.0 );
  246. velocity.addAssign( normalize( directionToClosestPoint ).mul( velocityAdjust ) );
  247. limit.addAssign( 5.0 );
  248. } );
  249. // Attract flocks to center
  250. const dirToCenter = position.toVar();
  251. dirToCenter.y.mulAssign( 2.5 );
  252. velocity.subAssign( normalize( dirToCenter ).mul( deltaTime ).mul( 5.0 ) );
  253. Loop( { start: uint( 0 ), end: uint( BIRDS ), type: 'uint', condition: '<' }, ( { i } ) => {
  254. const birdPosition = positionStorage.element( i );
  255. const dirToBird = birdPosition.sub( position );
  256. const distToBird = length( dirToBird );
  257. // Don't apply any changes to velocity if the distance to this bird is negligable.
  258. If( distToBird.lessThan( 0.0001 ), () => {
  259. Continue();
  260. } );
  261. const distToBirdSq = distToBird.mul( distToBird );
  262. // Don't apply any changes to velocity if changes if the bird is outsize the zone's radius.
  263. If( distToBirdSq.greaterThan( zoneRadiusSq ), () => {
  264. Continue();
  265. } );
  266. // Determine which threshold the bird is flying within and adjust its velocity accordingly
  267. const percent = distToBirdSq.div( zoneRadiusSq );
  268. If( percent.lessThan( separationThresh ), () => {
  269. // Separation - Move apart for comfort
  270. const velocityAdjust = ( separationThresh.div( percent ).sub( 1.0 ) ).mul( deltaTime );
  271. velocity.subAssign( normalize( dirToBird ).mul( velocityAdjust ) );
  272. } ).ElseIf( percent.lessThan( alignmentThresh ), () => {
  273. // Alignment - fly the same direction
  274. const threshDelta = alignmentThresh.sub( separationThresh );
  275. const adjustedPercent = ( percent.sub( separationThresh ) ).div( threshDelta );
  276. const birdVelocity = velocityStorage.element( i );
  277. const cosRange = cos( adjustedPercent.mul( PI_2 ) );
  278. const cosRangeAdjust = float( 0.5 ).sub( cosRange.mul( 0.5 ) ).add( 0.5 );
  279. const velocityAdjust = cosRangeAdjust.mul( deltaTime );
  280. velocity.addAssign( normalize( birdVelocity ).mul( velocityAdjust ) );
  281. } ).Else( () => {
  282. // Attraction / Cohesion - move closer
  283. const threshDelta = alignmentThresh.oneMinus();
  284. const adjustedPercent = threshDelta.equal( 0.0 ).select( 1.0, ( percent.sub( alignmentThresh ) ).div( threshDelta ) );
  285. const cosRange = cos( adjustedPercent.mul( PI_2 ) );
  286. const adj1 = cosRange.mul( - 0.5 );
  287. const adj2 = adj1.add( 0.5 );
  288. const adj3 = float( 0.5 ).sub( adj2 );
  289. const velocityAdjust = adj3.mul( deltaTime );
  290. velocity.addAssign( normalize( dirToBird ).mul( velocityAdjust ) );
  291. } );
  292. } );
  293. If( length( velocity ).greaterThan( limit ), () => {
  294. velocity.assign( normalize( velocity ).mul( limit ) );
  295. } );
  296. } )().compute( BIRDS );
  297. computePosition = Fn( () => {
  298. const { deltaTime } = effectController;
  299. positionStorage.element( instanceIndex ).addAssign( velocityStorage.element( instanceIndex ).mul( deltaTime ).mul( 15.0 ) );
  300. const velocity = velocityStorage.element( instanceIndex );
  301. const phase = phaseStorage.element( instanceIndex );
  302. const modValue = phase.add( deltaTime ).add( length( velocity.xz ).mul( deltaTime ).mul( 3.0 ) ).add( max( velocity.y, 0.0 ).mul( deltaTime ).mul( 6.0 ) );
  303. phaseStorage.element( instanceIndex ).assign( modValue.mod( 62.83 ) );
  304. } )().compute( BIRDS );
  305. scene.add( birdMesh );
  306. stats = new Stats();
  307. container.appendChild( stats.dom );
  308. container.style.touchAction = 'none';
  309. container.addEventListener( 'pointermove', onPointerMove );
  310. window.addEventListener( 'resize', onWindowResize );
  311. const gui = new GUI();
  312. gui.add( effectController.separation, 'value', 0.0, 100.0, 1.0 ).name( 'Separation' );
  313. gui.add( effectController.alignment, 'value', 0.0, 100, 0.001 ).name( 'Alignment ' );
  314. gui.add( effectController.cohesion, 'value', 0.0, 100, 0.025 ).name( 'Cohesion' );
  315. gui.close();
  316. }
  317. function onWindowResize() {
  318. camera.aspect = window.innerWidth / window.innerHeight;
  319. camera.updateProjectionMatrix();
  320. renderer.setSize( window.innerWidth, window.innerHeight );
  321. }
  322. function onPointerMove( event ) {
  323. if ( event.isPrimary === false ) return;
  324. pointer.x = ( event.clientX / window.innerWidth ) * 2.0 - 1.0;
  325. pointer.y = 1.0 - ( event.clientY / window.innerHeight ) * 2.0;
  326. }
  327. function animate() {
  328. render();
  329. stats.update();
  330. }
  331. function render() {
  332. const now = performance.now();
  333. let deltaTime = ( now - last ) / 1000;
  334. if ( deltaTime > 1 ) deltaTime = 1; // safety cap on large deltas
  335. last = now;
  336. raycaster.setFromCamera( pointer, camera );
  337. effectController.now.value = now;
  338. effectController.deltaTime.value = deltaTime;
  339. effectController.rayOrigin.value.copy( raycaster.ray.origin );
  340. effectController.rayDirection.value.copy( raycaster.ray.direction );
  341. renderer.compute( computeVelocity );
  342. renderer.compute( computePosition );
  343. renderer.render( scene, camera );
  344. // Move pointer away so we only affect birds when moving the mouse
  345. pointer.y = 10;
  346. }
  347. init();
  348. </script>
  349. </body>
  350. </html>