webgpu_compute_water.html 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639
  1. <!DOCTYPE html>
  2. <html lang="en">
  3. <head>
  4. <title>three.js webgpu - compute water</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="example.css">
  8. </head>
  9. <body>
  10. <div id="info">
  11. <a href="https://threejs.org/" target="_blank" rel="noopener" class="logo-link"></a>
  12. <div class="title-wrapper">
  13. <a href="https://threejs.org/" target="_blank" rel="noopener">three.js</a>
  14. <span>Compute Water</span>
  15. </div>
  16. <small>
  17. Click and move mouse to disturb water.
  18. </small>
  19. </div>
  20. <script type="importmap">
  21. {
  22. "imports": {
  23. "three": "../build/three.webgpu.js",
  24. "three/webgpu": "../build/three.webgpu.js",
  25. "three/tsl": "../build/three.tsl.js",
  26. "three/addons/": "./jsm/"
  27. }
  28. }
  29. </script>
  30. <script type="module">
  31. import * as THREE from 'three/webgpu';
  32. import { instanceIndex, struct, If, uint, int, floor, float, length, clamp, vec2, cos, vec3, vertexIndex, Fn, uniform, instancedArray, min, max, positionLocal, transformNormalToView, select, globalId } from 'three/tsl';
  33. import { Inspector } from 'three/addons/inspector/Inspector.js';
  34. import { SimplexNoise } from 'three/addons/math/SimplexNoise.js';
  35. import { GLTFLoader } from 'three/addons/loaders/GLTFLoader.js';
  36. import { HDRLoader } from 'three/addons/loaders/HDRLoader.js';
  37. import { DRACOLoader } from 'three/addons/loaders/DRACOLoader.js';
  38. import { OrbitControls } from 'three/addons/controls/OrbitControls.js';
  39. import WebGPU from 'three/addons/capabilities/WebGPU.js';
  40. // Dimensions of simulation grid.
  41. const WIDTH = 128;
  42. // Water size in system units.
  43. const BOUNDS = 6;
  44. const BOUNDS_HALF = BOUNDS * 0.5;
  45. const limit = BOUNDS_HALF - 0.2;
  46. const waterMaxHeight = 0.1;
  47. let container;
  48. let camera, scene, renderer, controls;
  49. let mouseDown = false;
  50. let firstClick = true;
  51. let updateOriginMouseDown = false;
  52. const mouseCoords = new THREE.Vector2();
  53. const raycaster = new THREE.Raycaster();
  54. let frame = 0;
  55. const effectController = {
  56. mousePos: uniform( new THREE.Vector2() ).setName( 'mousePos' ),
  57. mouseSpeed: uniform( new THREE.Vector2() ).setName( 'mouseSpeed' ),
  58. mouseDeep: uniform( .5 ).setName( 'mouseDeep' ),
  59. mouseSize: uniform( 0.12 ).setName( 'mouseSize' ),
  60. viscosity: uniform( 0.96 ).setName( 'viscosity' ),
  61. ducksEnabled: true,
  62. wireframe: false,
  63. speed: 5,
  64. };
  65. let sun;
  66. let waterMesh;
  67. let poolBorder;
  68. let meshRay;
  69. let computeHeightAtoB, computeHeightBtoA, computeDucks;
  70. let pingPong = 0;
  71. const readFromA = uniform( 1 );
  72. let duckModel = null;
  73. const NUM_DUCKS = 100;
  74. const simplex = new SimplexNoise();
  75. // TODO: Fix example with WebGL backend
  76. if ( WebGPU.isAvailable() === false ) {
  77. document.body.appendChild( WebGPU.getErrorMessage() );
  78. throw new Error( 'No WebGPU support' );
  79. }
  80. init();
  81. function noise( x, y ) {
  82. let multR = waterMaxHeight;
  83. let mult = 0.025;
  84. let r = 0;
  85. for ( let i = 0; i < 15; i ++ ) {
  86. r += multR * simplex.noise( x * mult, y * mult );
  87. multR *= 0.53 + 0.025 * i;
  88. mult *= 1.25;
  89. }
  90. return r;
  91. }
  92. async function init() {
  93. container = document.createElement( 'div' );
  94. document.body.appendChild( container );
  95. camera = new THREE.PerspectiveCamera( 75, window.innerWidth / window.innerHeight, 1, 3000 );
  96. camera.position.set( 0, 2.00, 4 );
  97. camera.lookAt( 0, 0, 0 );
  98. scene = new THREE.Scene();
  99. sun = new THREE.DirectionalLight( 0xFFFFFF, 4.0 );
  100. sun.position.set( - 1, 2.6, 1.4 );
  101. scene.add( sun );
  102. //
  103. // Initialize height storage buffers
  104. const heightArray = new Float32Array( WIDTH * WIDTH );
  105. const prevHeightArray = new Float32Array( WIDTH * WIDTH );
  106. let p = 0;
  107. for ( let j = 0; j < WIDTH; j ++ ) {
  108. for ( let i = 0; i < WIDTH; i ++ ) {
  109. const x = i * 128 / WIDTH;
  110. const y = j * 128 / WIDTH;
  111. const height = noise( x, y );
  112. heightArray[ p ] = height;
  113. prevHeightArray[ p ] = height;
  114. p ++;
  115. }
  116. }
  117. // Ping-pong height storage buffers
  118. const heightStorageA = instancedArray( heightArray ).setName( 'HeightA' );
  119. const heightStorageB = instancedArray( new Float32Array( heightArray ) ).setName( 'HeightB' );
  120. const prevHeightStorage = instancedArray( prevHeightArray ).setName( 'PrevHeight' );
  121. // Get Indices of Neighbor Values of an Index in the Simulation Grid
  122. const getNeighborIndicesTSL = ( index ) => {
  123. const width = uint( WIDTH );
  124. // Get 2-D compute coordinate from one-dimensional instanceIndex. The calculation will
  125. // still work even if you dispatch your compute shader 2-dimensionally, since within a compute
  126. // context, instanceIndex is a 1-dimensional value derived from the workgroup dimensions.
  127. // Cast to int to prevent unintended index overflow upon subtraction.
  128. const x = int( index.mod( WIDTH ) );
  129. const y = int( index.div( WIDTH ) );
  130. // The original shader accesses height via texture uvs. However, unlike with textures, we can't
  131. // access areas that are out of bounds. Accordingly, we emulate the Clamp to Edge Wrapping
  132. // behavior of accessing a DataTexture with out of bounds uvs.
  133. const leftX = max( 0, x.sub( 1 ) );
  134. const rightX = min( x.add( 1 ), width.sub( 1 ) );
  135. const bottomY = max( 0, y.sub( 1 ) );
  136. const topY = min( y.add( 1 ), width.sub( 1 ) );
  137. const westIndex = y.mul( width ).add( leftX );
  138. const eastIndex = y.mul( width ).add( rightX );
  139. const southIndex = bottomY.mul( width ).add( x );
  140. const northIndex = topY.mul( width ).add( x );
  141. return { northIndex, southIndex, eastIndex, westIndex };
  142. };
  143. // Get simulation index neighbor values
  144. const getNeighborValuesTSL = ( index, store ) => {
  145. const { northIndex, southIndex, eastIndex, westIndex } = getNeighborIndicesTSL( index );
  146. const north = store.element( northIndex );
  147. const south = store.element( southIndex );
  148. const east = store.element( eastIndex );
  149. const west = store.element( westIndex );
  150. return { north, south, east, west };
  151. };
  152. // Create compute shader for height simulation with explicit read/write buffers
  153. const createComputeHeight = ( readBuffer, writeBuffer ) => Fn( () => {
  154. const { viscosity, mousePos, mouseSize, mouseDeep, mouseSpeed } = effectController;
  155. const height = readBuffer.element( instanceIndex ).toVar();
  156. const prevHeight = prevHeightStorage.element( instanceIndex ).toVar();
  157. const { north, south, east, west } = getNeighborValuesTSL( instanceIndex, readBuffer );
  158. const neighborHeight = north.add( south ).add( east ).add( west );
  159. neighborHeight.mulAssign( 0.5 );
  160. neighborHeight.subAssign( prevHeight );
  161. const newHeight = neighborHeight.mul( viscosity );
  162. // Get x and y position of the coordinate in the water plane
  163. const x = float( globalId.x ).mul( 1 / WIDTH );
  164. const y = float( globalId.y ).mul( 1 / WIDTH );
  165. // Mouse influence
  166. const centerVec = vec2( 0.5 );
  167. // Get length of position in range [ -BOUNDS / 2, BOUNDS / 2 ], offset by mousePos, then scale.
  168. const mousePhase = clamp( length( ( vec2( x, y ).sub( centerVec ) ).mul( BOUNDS ).sub( mousePos ) ).mul( Math.PI ).div( mouseSize ), 0.0, Math.PI );
  169. // "Indent" water down by scaled distance from center of mouse impact
  170. newHeight.addAssign( cos( mousePhase ).add( 1.0 ).mul( mouseDeep ).mul( mouseSpeed.length() ) );
  171. prevHeightStorage.element( instanceIndex ).assign( height );
  172. writeBuffer.element( instanceIndex ).assign( newHeight );
  173. } )().compute( WIDTH * WIDTH, [ 16, 16 ] );
  174. // Create both ping-pong compute shaders
  175. computeHeightAtoB = createComputeHeight( heightStorageA, heightStorageB ).setName( 'Update Height A→B' );
  176. computeHeightBtoA = createComputeHeight( heightStorageB, heightStorageA ).setName( 'Update Height B→A' );
  177. // Water Geometry corresponds with buffered compute grid.
  178. const waterGeometry = new THREE.PlaneGeometry( BOUNDS, BOUNDS, WIDTH - 1, WIDTH - 1 );
  179. const waterMaterial = new THREE.MeshStandardNodeMaterial( {
  180. color: 0x9bd2ec,
  181. metalness: 0.9,
  182. roughness: 0,
  183. transparent: true,
  184. opacity: 0.8,
  185. side: THREE.DoubleSide
  186. } );
  187. // Helper to get height from the current read buffer
  188. const getCurrentHeight = ( index ) => {
  189. return select( readFromA, heightStorageA.element( index ), heightStorageB.element( index ) );
  190. };
  191. // Helper to get normals from the current read buffer
  192. const getCurrentNormals = ( index ) => {
  193. const { northIndex, southIndex, eastIndex, westIndex } = getNeighborIndicesTSL( index );
  194. const north = getCurrentHeight( northIndex );
  195. const south = getCurrentHeight( southIndex );
  196. const east = getCurrentHeight( eastIndex );
  197. const west = getCurrentHeight( westIndex );
  198. const normalX = ( west.sub( east ) ).mul( WIDTH / BOUNDS );
  199. const normalY = ( south.sub( north ) ).mul( WIDTH / BOUNDS );
  200. return { normalX, normalY };
  201. };
  202. waterMaterial.normalNode = Fn( () => {
  203. // To correct the lighting as our mesh undulates, we have to reassign the normals in the normal shader.
  204. const { normalX, normalY } = getCurrentNormals( vertexIndex );
  205. return transformNormalToView( vec3( normalX, normalY.negate(), 1.0 ) ).toVertexStage();
  206. } )();
  207. waterMaterial.positionNode = Fn( () => {
  208. return vec3( positionLocal.x, positionLocal.y, getCurrentHeight( vertexIndex ) );
  209. } )();
  210. waterMesh = new THREE.Mesh( waterGeometry, waterMaterial );
  211. waterMesh.rotation.x = - Math.PI * 0.5;
  212. waterMesh.matrixAutoUpdate = false;
  213. waterMesh.updateMatrix();
  214. scene.add( waterMesh );
  215. // Pool border
  216. const borderGeom = new THREE.TorusGeometry( 4.2, 0.1, 12, 4 );
  217. borderGeom.rotateX( Math.PI * 0.5 );
  218. borderGeom.rotateY( Math.PI * 0.25 );
  219. poolBorder = new THREE.Mesh( borderGeom, new THREE.MeshStandardMaterial( { color: 0x908877, roughness: 0.2 } ) );
  220. scene.add( poolBorder );
  221. // THREE.Mesh just for mouse raycasting
  222. const geometryRay = new THREE.PlaneGeometry( BOUNDS, BOUNDS, 1, 1 );
  223. meshRay = new THREE.Mesh( geometryRay, new THREE.MeshBasicMaterial( { color: 0xFFFFFF, visible: false } ) );
  224. meshRay.rotation.x = - Math.PI / 2;
  225. meshRay.matrixAutoUpdate = false;
  226. meshRay.updateMatrix();
  227. scene.add( meshRay );
  228. // Initialize sphere mesh instance position and velocity.
  229. // position<vec3> + velocity<vec2> + unused<vec3> = 8 floats per sphere.
  230. // for structs arrays must be enclosed in multiple of 4
  231. const duckStride = 8;
  232. const duckInstanceDataArray = new Float32Array( NUM_DUCKS * duckStride );
  233. // Only hold velocity in x and z directions.
  234. // The sphere is wedded to the surface of the water, and will only move vertically with the water.
  235. for ( let i = 0; i < NUM_DUCKS; i ++ ) {
  236. duckInstanceDataArray[ i * duckStride + 0 ] = ( Math.random() - 0.5 ) * BOUNDS * 0.7;
  237. duckInstanceDataArray[ i * duckStride + 1 ] = 0;
  238. duckInstanceDataArray[ i * duckStride + 2 ] = ( Math.random() - 0.5 ) * BOUNDS * 0.7;
  239. }
  240. const DuckStruct = struct( {
  241. position: 'vec3',
  242. velocity: 'vec2'
  243. } );
  244. // Duck instance data storage
  245. const duckInstanceDataStorage = instancedArray( duckInstanceDataArray, DuckStruct ).setName( 'DuckInstanceData' );
  246. computeDucks = Fn( () => {
  247. const yOffset = float( - 0.04 );
  248. const verticalResponseFactor = float( 0.98 );
  249. const waterPushFactor = float( 0.015 );
  250. const linearDamping = float( 0.92 );
  251. const bounceDamping = float( - 0.4 );
  252. // Get 2-D compute coordinate from one-dimensional instanceIndex.
  253. const instancePosition = duckInstanceDataStorage.element( instanceIndex ).get( 'position' ).toVar();
  254. const velocity = duckInstanceDataStorage.element( instanceIndex ).get( 'velocity' ).toVar();
  255. const gridCoordX = instancePosition.x.div( BOUNDS ).add( 0.5 ).mul( WIDTH );
  256. const gridCoordZ = instancePosition.z.div( BOUNDS ).add( 0.5 ).mul( WIDTH );
  257. // Cast to int to prevent unintended index overflow upon subtraction.
  258. const xCoord = uint( clamp( floor( gridCoordX ), 0, WIDTH - 1 ) );
  259. const zCoord = uint( clamp( floor( gridCoordZ ), 0, WIDTH - 1 ) );
  260. const heightInstanceIndex = zCoord.mul( WIDTH ).add( xCoord );
  261. // Get height of water at the duck's position (use current read buffer)
  262. const waterHeight = getCurrentHeight( heightInstanceIndex );
  263. const { normalX, normalY } = getCurrentNormals( heightInstanceIndex );
  264. // Calculate the target Y position based on the water height and the duck's vertical offset
  265. const targetY = waterHeight.add( yOffset );
  266. const deltaY = targetY.sub( instancePosition.y );
  267. instancePosition.y.addAssign( deltaY.mul( verticalResponseFactor ) ); // Gradually update position
  268. // Get the normal of the water surface at the duck's position
  269. const pushX = normalX.mul( waterPushFactor );
  270. const pushZ = normalY.mul( waterPushFactor );
  271. // Apply the water push to the duck's velocity
  272. velocity.x.mulAssign( linearDamping );
  273. velocity.y.mulAssign( linearDamping );
  274. velocity.x.addAssign( pushX );
  275. velocity.y.addAssign( pushZ );
  276. // update position based on velocity
  277. instancePosition.x.addAssign( velocity.x );
  278. instancePosition.z.addAssign( velocity.y );
  279. // Clamp position to the pool bounds
  280. If( instancePosition.x.lessThan( - limit ), () => {
  281. instancePosition.x = - limit;
  282. velocity.x.mulAssign( bounceDamping );
  283. } ).ElseIf( instancePosition.x.greaterThan( limit ), () => {
  284. instancePosition.x = limit;
  285. velocity.x.mulAssign( bounceDamping );
  286. } );
  287. If( instancePosition.z.lessThan( - limit ), () => {
  288. instancePosition.z = - limit;
  289. velocity.y.mulAssign( bounceDamping ); // Invert and damp vz (velocity.y)
  290. } ).ElseIf( instancePosition.z.greaterThan( limit ), () => {
  291. instancePosition.z = limit;
  292. velocity.y.mulAssign( bounceDamping );
  293. } );
  294. // assignment of new values to the instance data storage
  295. duckInstanceDataStorage.element( instanceIndex ).get( 'position' ).assign( instancePosition );
  296. duckInstanceDataStorage.element( instanceIndex ).get( 'velocity' ).assign( velocity );
  297. } )().compute( NUM_DUCKS ).setName( 'Update Ducks' );
  298. // Models / Textures
  299. const hdrLoader = new HDRLoader().setPath( './textures/equirectangular/' );
  300. const glbloader = new GLTFLoader().setPath( 'models/gltf/' );
  301. glbloader.setDRACOLoader( new DRACOLoader().setDecoderPath( 'jsm/libs/draco/gltf/' ) );
  302. const [ env, model ] = await Promise.all( [ hdrLoader.loadAsync( 'blouberg_sunrise_2_1k.hdr' ), glbloader.loadAsync( 'duck.glb' ) ] );
  303. env.mapping = THREE.EquirectangularReflectionMapping;
  304. scene.environment = env;
  305. scene.background = env;
  306. scene.backgroundBlurriness = 0.3;
  307. scene.environmentIntensity = 1.25;
  308. duckModel = model.scene.children[ 0 ];
  309. duckModel.material.positionNode = Fn( () => {
  310. const instancePosition = duckInstanceDataStorage.element( instanceIndex ).get( 'position' );
  311. const newPosition = positionLocal.add( instancePosition );
  312. return newPosition;
  313. } )();
  314. const duckMesh = new THREE.InstancedMesh( duckModel.geometry, duckModel.material, NUM_DUCKS );
  315. scene.add( duckMesh );
  316. renderer = new THREE.WebGPURenderer( { antialias: true } );
  317. renderer.setPixelRatio( window.devicePixelRatio );
  318. renderer.setSize( window.innerWidth, window.innerHeight );
  319. renderer.toneMapping = THREE.ACESFilmicToneMapping;
  320. renderer.toneMappingExposure = 0.5;
  321. renderer.setAnimationLoop( render );
  322. container.appendChild( renderer.domElement );
  323. renderer.inspector = new Inspector();
  324. document.body.appendChild( renderer.inspector.domElement );
  325. controls = new OrbitControls( camera, container );
  326. //
  327. container.style.touchAction = 'none';
  328. container.addEventListener( 'pointermove', onPointerMove );
  329. container.addEventListener( 'pointerdown', onPointerDown );
  330. container.addEventListener( 'pointerup', onPointerUp );
  331. window.addEventListener( 'resize', onWindowResize );
  332. // GUI
  333. const gui = renderer.inspector.createParameters( 'Settings' );
  334. gui.add( effectController.mouseSize, 'value', 0.1, .3 ).name( 'Mouse Size' );
  335. gui.add( effectController.mouseDeep, 'value', 0.1, 1 ).name( 'Mouse Deep' );
  336. gui.add( effectController.viscosity, 'value', 0.9, 0.96, 0.001 ).name( 'viscosity' );
  337. gui.add( effectController, 'speed', 1, 6, 1 );
  338. gui.add( effectController, 'ducksEnabled' ).onChange( () => {
  339. duckMesh.visible = effectController.ducksEnabled;
  340. } );
  341. gui.add( effectController, 'wireframe' ).onChange( () => {
  342. waterMesh.material.wireframe = ! waterMesh.material.wireframe;
  343. poolBorder.material.wireframe = ! poolBorder.material.wireframe;
  344. duckModel.material.wireframe = ! duckModel.material.wireframe;
  345. waterMesh.material.needsUpdate = true;
  346. poolBorder.material.needsUpdate = true;
  347. } );
  348. }
  349. function onWindowResize() {
  350. camera.aspect = window.innerWidth / window.innerHeight;
  351. camera.updateProjectionMatrix();
  352. renderer.setSize( window.innerWidth, window.innerHeight );
  353. }
  354. function setMouseCoords( x, y ) {
  355. mouseCoords.set( ( x / renderer.domElement.clientWidth ) * 2 - 1, - ( y / renderer.domElement.clientHeight ) * 2 + 1 );
  356. }
  357. function onPointerDown() {
  358. mouseDown = true;
  359. firstClick = true;
  360. updateOriginMouseDown = true;
  361. }
  362. function onPointerUp() {
  363. mouseDown = false;
  364. firstClick = false;
  365. updateOriginMouseDown = false;
  366. controls.enabled = true;
  367. }
  368. function onPointerMove( event ) {
  369. if ( event.isPrimary === false ) return;
  370. setMouseCoords( event.clientX, event.clientY );
  371. }
  372. function raycast() {
  373. if ( mouseDown && ( firstClick || ! controls.enabled ) ) {
  374. raycaster.setFromCamera( mouseCoords, camera );
  375. const intersects = raycaster.intersectObject( meshRay );
  376. if ( intersects.length > 0 ) {
  377. const point = intersects[ 0 ].point;
  378. if ( updateOriginMouseDown ) {
  379. effectController.mousePos.value.set( point.x, point.z );
  380. updateOriginMouseDown = false;
  381. }
  382. effectController.mouseSpeed.value.set(
  383. ( point.x - effectController.mousePos.value.x ),
  384. ( point.z - effectController.mousePos.value.y )
  385. );
  386. effectController.mousePos.value.set( point.x, point.z );
  387. if ( firstClick ) {
  388. controls.enabled = false;
  389. }
  390. } else {
  391. updateOriginMouseDown = true;
  392. effectController.mouseSpeed.value.set( 0, 0 );
  393. }
  394. firstClick = false;
  395. } else {
  396. updateOriginMouseDown = true;
  397. effectController.mouseSpeed.value.set( 0, 0 );
  398. }
  399. }
  400. function render() {
  401. raycast();
  402. frame ++;
  403. if ( frame >= 7 - effectController.speed ) {
  404. // Ping-pong: alternate which buffer we read from and write to
  405. if ( pingPong === 0 ) {
  406. renderer.compute( computeHeightAtoB, [ 8, 8, 1 ] );
  407. readFromA.value = 0; // Material now reads from B (just written)
  408. } else {
  409. renderer.compute( computeHeightBtoA, [ 8, 8, 1 ] );
  410. readFromA.value = 1; // Material now reads from A (just written)
  411. }
  412. pingPong = 1 - pingPong;
  413. if ( effectController.ducksEnabled ) {
  414. renderer.compute( computeDucks );
  415. }
  416. frame = 0;
  417. }
  418. renderer.render( scene, camera );
  419. }
  420. </script>
  421. </body>
  422. </html>
粤ICP备19079148号