123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572 |
- <html lang="en">
- <head>
- <title>three.js webgpu - storage pbo external element</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</a>
- <br /> This example demonstrates a bitonic sort running step by step in a compute shader.
- <br /> The left canvas swaps values within workgroup local arrays. The right swaps values within storage buffers.
- <br /> Reference implementation by <a href="https://poniesandlight.co.uk/reflect/bitonic_merge_sort/">Tim Gfrerer</a>
- <br />
- <div id="local_swap" style="
- position: absolute;
- top: 150px;
- left: 0;
- padding: 10px;
- background: rgba( 0, 0, 0, 0.5 );
- color: #fff;
- font-family: monospace;
- font-size: 12px;
- line-height: 1.5;
- pointer-events: none;
- text-align: left;
- "></div>
- <div id="global_swap" style="
- position: absolute;
- top: 150px;
- right: 0;
- padding: 10px;
- background: rgba( 0, 0, 0, 0.5 );
- color: #fff;
- font-family: monospace;
- font-size: 12px;
- line-height: 1.5;
- pointer-events: none;
- text-align: left;
- "></div>
- </div>
- <script type="importmap">
- {
- "imports": {
- "three": "../build/three.webgpu.js",
- "three/tsl": "../build/three.webgpu.js",
- "three/addons/": "./jsm/"
- }
- }
- </script>
- <script type="module">
- import * as THREE from 'three';
- import { storageObject, If, vec3, not, uniform, uv, uint, float, Fn, vec2, abs, int, invocationLocalIndex, workgroupArray, uvec2, floor, instanceIndex, workgroupBarrier, atomicAdd, atomicStore, workgroupId } from 'three/tsl';
- import { GUI } from 'three/addons/libs/lil-gui.module.min.js';
- const StepType = {
- NONE: 0,
- // Swap values within workgroup local buffer.
- FLIP_LOCAL: 1,
- DISPERSE_LOCAL: 2,
- // Swap values within global data buffer.
- FLIP_GLOBAL: 3,
- DISPERSE_GLOBAL: 4,
- };
- const timestamps = {
- local_swap: document.getElementById( 'local_swap' ),
- global_swap: document.getElementById( 'global_swap' )
- };
- const localColors = [ 'rgb(203, 64, 203)', 'rgb(0, 215, 215)' ];
- const globalColors = [ 'rgb(1, 150, 1)', 'red' ];
- // Total number of elements and the dimensions of the display grid.
- const size = 16384;
- const gridDim = Math.sqrt( size );
- const getNumSteps = () => {
- const n = Math.log2( size );
- return ( n * ( n + 1 ) ) / 2;
- };
- // Total number of steps in a bitonic sort with 'size' elements.
- const MAX_STEPS = getNumSteps();
- const WORKGROUP_SIZE = [ 64 ];
- const effectController = {
- // Sqr root of 16834
- gridWidth: uniform( gridDim ),
- gridHeight: uniform( gridDim ),
- highlight: uniform( 1 ),
- 'Display Mode': 'Swap Zone Highlight'
- };
- const gui = new GUI();
- gui.add( effectController, 'Display Mode', [ 'Elements', 'Swap Zone Highlight' ] ).onChange( () => {
- if ( effectController[ 'Display Mode' ] === 'Elements' ) {
- effectController.highlight.value = 0;
-
- } else {
- effectController.highlight.value = 1;
-
- }
-
- } );
- // Allow Workgroup Array Swaps
- init();
- // Global Swaps Only
- init( true );
- // When forceGlobalSwap is true, force all valid local swaps to be global swaps.
- async function init( forceGlobalSwap = false ) {
- let currentStep = 0;
- let nextStepGlobal = false;
- const aspect = ( window.innerWidth / 2 ) / window.innerHeight;
- const camera = new THREE.OrthographicCamera( - aspect, aspect, 1, - 1, 0, 2 );
- camera.position.z = 1;
- const scene = new THREE.Scene();
- const nextAlgoBuffer = new THREE.StorageInstancedBufferAttribute( new Uint32Array( 1 ).fill( forceGlobalSwap ? StepType.FLIP_GLOBAL : StepType.FLIP_LOCAL ), 1 );
- const nextAlgoStorage = storageObject( nextAlgoBuffer, 'uint', nextAlgoBuffer.count ).label( 'NextAlgo' );
- const nextBlockHeightBuffer = new THREE.StorageInstancedBufferAttribute( new Uint32Array( 1 ).fill( 2 ), 1 );
- const nextBlockHeightStorage = storageObject( nextBlockHeightBuffer, 'uint', nextBlockHeightBuffer.count ).label( 'NextBlockHeight' );
- const nextBlockHeightRead = storageObject( nextBlockHeightBuffer, 'uint', nextBlockHeightBuffer.count ).label( 'NextBlockHeight' ).toReadOnly();
- const highestBlockHeightBuffer = new THREE.StorageInstancedBufferAttribute( new Uint32Array( 1 ).fill( 2 ), 1 );
- const highestBlockHeightStorage = storageObject( highestBlockHeightBuffer, 'uint', highestBlockHeightBuffer.count ).label( 'HighestBlockHeight' );
- const counterBuffer = new THREE.StorageBufferAttribute( 1, 1 );
- const counterStorage = storageObject( counterBuffer, 'uint', counterBuffer.count ).toAtomic().label( 'Counter' );
- const array = new Uint32Array( Array.from( { length: size }, ( _, i ) => {
- return i;
- } ) );
- const randomizeDataArray = () => {
- let currentIndex = array.length;
- while ( currentIndex !== 0 ) {
- const randomIndex = Math.floor( Math.random() * currentIndex );
- currentIndex -= 1;
- [ array[ currentIndex ], array[ randomIndex ] ] = [
- array[ randomIndex ],
- array[ currentIndex ],
- ];
- }
- };
- randomizeDataArray();
- const currentElementsBuffer = new THREE.StorageInstancedBufferAttribute( array, 1 );
- const currentElementsStorage = storageObject( currentElementsBuffer, 'uint', size ).label( 'Elements' );
- const tempBuffer = new THREE.StorageInstancedBufferAttribute( array, 1 );
- const tempStorage = storageObject( tempBuffer, 'uint', size ).label( 'Temp' );
- const randomizedElementsBuffer = new THREE.StorageInstancedBufferAttribute( size, 1 );
- const randomizedElementsStorage = storageObject( randomizedElementsBuffer, 'uint', size ).label( 'RandomizedElements' );
- const getFlipIndices = ( index, blockHeight ) => {
- const blockOffset = ( index.mul( 2 ).div( blockHeight ) ).mul( blockHeight );
- const halfHeight = blockHeight.div( 2 );
- const idx = uvec2(
- index.modInt( halfHeight ),
- blockHeight.sub( index.modInt( halfHeight ) ).sub( 1 )
- );
- idx.x.addAssign( blockOffset );
- idx.y.addAssign( blockOffset );
- return idx;
- };
- const getDisperseIndices = ( index, blockHeight ) => {
- const blockOffset = ( ( index.mul( 2 ) ).div( blockHeight ) ).mul( blockHeight );
- const halfHeight = blockHeight.div( 2 );
- const idx = uvec2(
- index.modInt( halfHeight ),
- ( index.modInt( halfHeight ) ).add( halfHeight )
- );
- idx.x.addAssign( blockOffset );
- idx.y.addAssign( blockOffset );
- return idx;
- };
- const localStorage = workgroupArray( 'uint', 64 * 2 );
- // Swap the elements in local storage
- const localCompareAndSwap = ( idxBefore, idxAfter ) => {
- If( localStorage.element( idxAfter ).lessThan( localStorage.element( idxBefore ) ), () => {
- atomicAdd( counterStorage.element( 0 ), 1 );
- const temp = localStorage.element( idxBefore ).toVar();
- localStorage.element( idxBefore ).assign( localStorage.element( idxAfter ) );
- localStorage.element( idxAfter ).assign( temp );
- } );
- };
- const globalCompareAndSwap = ( idxBefore, idxAfter ) => {
- // If the later element is less than the current element
- If( currentElementsStorage.element( idxAfter ).lessThan( currentElementsStorage.element( idxBefore ) ), () => {
- // Apply the swapped values to temporary storage.
- atomicAdd( counterStorage.element( 0 ), 1 );
- tempStorage.element( idxBefore ).assign( currentElementsStorage.element( idxAfter ) );
- tempStorage.element( idxAfter ).assign( currentElementsStorage.element( idxBefore ) );
- } ).Else( () => {
- // Otherwise apply the existing values to temporary storage.
- tempStorage.element( idxBefore ).assign( currentElementsStorage.element( idxBefore ) );
- tempStorage.element( idxAfter ).assign( currentElementsStorage.element( idxAfter ) );
- } );
- };
- const computeInitFn = Fn( () => {
- randomizedElementsStorage.element( instanceIndex ).assign( currentElementsStorage.element( instanceIndex ) );
- } );
- const computeBitonicStepFn = Fn( () => {
- const nextBlockHeight = nextBlockHeightStorage.element( 0 ).toVar();
- const nextAlgo = nextAlgoStorage.element( 0 ).toVar();
- // Get ids of indices needed to populate workgroup local buffer.
- // Use .toVar() to prevent these values from being recalculated multiple times.
- const localOffset = uint( WORKGROUP_SIZE[ 0 ] ).mul( 2 ).mul( workgroupId.x ).toVar();
-
- const localID1 = invocationLocalIndex.mul( 2 );
- const localID2 = invocationLocalIndex.mul( 2 ).add( 1 );
- // If we will perform a local swap, then populate the local data
- If( nextAlgo.lessThanEqual( uint( StepType.DISPERSE_LOCAL ) ), () => {
- localStorage.element( localID1 ).assign( currentElementsStorage.element( localOffset.add( localID1 ) ) );
- localStorage.element( localID2 ).assign( currentElementsStorage.element( localOffset.add( localID2 ) ) );
- } );
- workgroupBarrier();
- // TODO: Convert to switch block.
- If( nextAlgo.equal( uint( StepType.FLIP_LOCAL ) ), () => {
- const idx = getFlipIndices( invocationLocalIndex, nextBlockHeight );
- localCompareAndSwap( idx.x, idx.y );
- } ).ElseIf( nextAlgo.equal( uint( StepType.DISPERSE_LOCAL ) ), () => {
- const idx = getDisperseIndices( invocationLocalIndex, nextBlockHeight );
- localCompareAndSwap( idx.x, idx.y );
- } ).ElseIf( nextAlgo.equal( uint( StepType.FLIP_GLOBAL ) ), () => {
- const idx = getFlipIndices( instanceIndex, nextBlockHeight );
- globalCompareAndSwap( idx.x, idx.y );
- } ).ElseIf( nextAlgo.equal( uint( StepType.DISPERSE_GLOBAL ) ), () => {
- const idx = getDisperseIndices( instanceIndex, nextBlockHeight );
- globalCompareAndSwap( idx.x, idx.y );
- } );
- // Ensure that all invocations have swapped their own regions of data
- workgroupBarrier();
- // Populate output data with the results from our swaps
- If( nextAlgo.lessThanEqual( uint( StepType.DISPERSE_LOCAL ) ), () => {
- currentElementsStorage.element( localOffset.add( localID1 ) ).assign( localStorage.element( localID1 ) );
- currentElementsStorage.element( localOffset.add( localID2 ) ).assign( localStorage.element( localID2 ) );
- } );
- // If the previous algorithm was global, we execute an additional compute step to sync the current buffer with the output buffer.
- } );
- const computeSetAlgoFn = Fn( () => {
- const nextBlockHeight = nextBlockHeightStorage.element( 0 ).toVar();
- const nextAlgo = nextAlgoStorage.element( 0 );
- const highestBlockHeight = highestBlockHeightStorage.element( 0 ).toVar();
- nextBlockHeight.divAssign( 2 );
- If( nextBlockHeight.equal( 1 ), () => {
- highestBlockHeight.mulAssign( 2 );
- if ( forceGlobalSwap ) {
- If( highestBlockHeight.equal( size * 2 ), () => {
-
- nextAlgo.assign( StepType.NONE );
- nextBlockHeight.assign( 0 );
- } ).Else( () => {
- nextAlgo.assign( StepType.FLIP_GLOBAL );
- nextBlockHeight.assign( highestBlockHeight );
- } );
- } else {
- If( highestBlockHeight.equal( size * 2 ), () => {
-
- nextAlgo.assign( StepType.NONE );
- nextBlockHeight.assign( 0 );
- } ).ElseIf( highestBlockHeight.greaterThan( WORKGROUP_SIZE[ 0 ] * 2 ), () => {
- nextAlgo.assign( StepType.FLIP_GLOBAL );
- nextBlockHeight.assign( highestBlockHeight );
- } ).Else( () => {
- nextAlgo.assign( forceGlobalSwap ? StepType.FLIP_GLOBAL : StepType.FLIP_LOCAL );
- nextBlockHeight.assign( highestBlockHeight );
- } );
- }
- } ).Else( () => {
- if ( forceGlobalSwap ) {
- nextAlgo.assign( StepType.DISPERSE_GLOBAL );
- } else {
- nextAlgo.assign( nextBlockHeight.greaterThan( WORKGROUP_SIZE[ 0 ] * 2 ).select( StepType.DISPERSE_GLOBAL, StepType.DISPERSE_LOCAL ) );
- }
- } );
- nextBlockHeightStorage.element( 0 ).assign( nextBlockHeight );
- highestBlockHeightStorage.element( 0 ).assign( highestBlockHeight );
- } );
- const computeAlignCurrentFn = Fn( () => {
- currentElementsStorage.element( instanceIndex ).assign( tempStorage.element( instanceIndex ) );
- } );
- const computeResetBuffersFn = Fn( () => {
- currentElementsStorage.element( instanceIndex ).assign( randomizedElementsStorage.element( instanceIndex ) );
-
- } );
- const computeResetAlgoFn = Fn( () => {
- nextAlgoStorage.element( 0 ).assign( forceGlobalSwap ? StepType.FLIP_GLOBAL : StepType.FLIP_LOCAL );
- nextBlockHeightStorage.element( 0 ).assign( 2 );
- highestBlockHeightStorage.element( 0 ).assign( 2 );
- atomicStore( counterStorage.element( 0 ), 0 );
- } );
- // Initialize each value in the elements buffer.
- const computeInit = computeInitFn().compute( size );
- // Swap a pair of elements in the elements buffer.
- const computeBitonicStep = computeBitonicStepFn().compute( size / 2 );
- // Set the conditions for the next swap.
- const computeSetAlgo = computeSetAlgoFn().compute( 1 );
- // Align the current buffer with the temp buffer if the previous sort was executed in a global scope.
- const computeAlignCurrent = computeAlignCurrentFn().compute( size );
- // Reset the buffers and algorithm information after a full bitonic sort has been completed.
- const computeResetBuffers = computeResetBuffersFn().compute( size );
- const computeResetAlgo = computeResetAlgoFn().compute( 1 );
- const material = new THREE.MeshBasicNodeMaterial( { color: 0x00ff00 } );
- const display = Fn( () => {
- const { gridWidth, gridHeight, highlight } = effectController;
- const newUV = uv().mul( vec2( gridWidth, gridHeight ) );
- const pixel = uvec2( uint( floor( newUV.x ) ), uint( floor( newUV.y ) ) );
- const elementIndex = uint( gridWidth ).mul( pixel.y ).add( pixel.x );
- const colorChanger = currentElementsStorage.element( elementIndex );
- const subtracter = float( colorChanger ).div( gridWidth.mul( gridHeight ) );
- const color = vec3( subtracter.oneMinus() ).toVar();
- If( highlight.equal( 1 ).and( not( nextAlgoStorage.element( 0 ).equal( StepType.NONE ) ) ), () => {
- const boolCheck = int( elementIndex.modInt( nextBlockHeightRead.element( 0 ) ).lessThan( nextBlockHeightRead.element( 0 ).div( 2 ) ) );
- color.z.assign( nextAlgoStorage.element( 0 ).lessThanEqual( StepType.DISPERSE_LOCAL ) );
- color.x.mulAssign( boolCheck );
- color.y.mulAssign( abs( boolCheck.sub( 1 ) ) );
- } );
- return color;
- } );
- material.colorNode = display();
- const plane = new THREE.Mesh( new THREE.PlaneGeometry( 1, 1 ), material );
- scene.add( plane );
- const renderer = new THREE.WebGPURenderer( { antialias: false, trackTimestamp: true } );
- renderer.setPixelRatio( window.devicePixelRatio );
- renderer.setSize( window.innerWidth / 2, window.innerHeight );
- const animate = () => {
- renderer.render( scene, camera );
- };
- renderer.setAnimationLoop( animate );
- document.body.appendChild( renderer.domElement );
- renderer.domElement.style.position = 'absolute';
- renderer.domElement.style.top = '0';
- renderer.domElement.style.left = '0';
- renderer.domElement.style.width = '50%';
- renderer.domElement.style.height = '100%';
- if ( forceGlobalSwap ) {
- renderer.domElement.style.left = '50%';
- scene.background = new THREE.Color( 0x212121 );
-
- } else {
- scene.background = new THREE.Color( 0x313131 );
- }
- await renderer.computeAsync( computeInit );
- renderer.info.autoReset = false;
- const stepAnimation = async function () {
- renderer.info.reset();
- if ( currentStep !== MAX_STEPS ) {
- renderer.compute( computeBitonicStep );
- if ( nextStepGlobal ) {
- renderer.compute( computeAlignCurrent );
-
- }
-
- renderer.compute( computeSetAlgo );
- currentStep ++;
- } else {
- renderer.compute( computeResetBuffers );
- renderer.compute( computeResetAlgo );
- currentStep = 0;
-
- }
- const algo = new Uint32Array( await renderer.getArrayBufferAsync( nextAlgoBuffer ) );
- algo > StepType.DISPERSE_LOCAL ? ( nextStepGlobal = true ) : ( nextStepGlobal = false );
- const totalSwaps = new Uint32Array( await renderer.getArrayBufferAsync( counterBuffer ) );
-
- renderer.render( scene, camera );
- timestamps[ forceGlobalSwap ? 'global_swap' : 'local_swap' ].innerHTML = `
- Compute ${forceGlobalSwap ? 'Global' : 'Local'}: ${renderer.info.compute.frameCalls} pass in ${renderer.info.compute.timestamp.toFixed( 6 )}ms<br>
- Total Swaps: ${totalSwaps}<br>
- <div style="display: flex; flex-direction:row; justify-content: center; align-items: center;">
- ${forceGlobalSwap ? 'Global Swaps' : 'Local Swaps'} Compare Region
- <div style="background-color: ${ forceGlobalSwap ? globalColors[ 0 ] : localColors[ 0 ]}; width:12.5px; height: 1em; border-radius: 20%;"></div>
- to Region
- <div style="background-color: ${ forceGlobalSwap ? globalColors[ 1 ] : localColors[ 1 ]}; width:12.5px; height: 1em; border-radius: 20%;"></div>
- </div>`;
-
- if ( currentStep === MAX_STEPS ) {
- setTimeout( stepAnimation, 1000 );
-
- } else {
- setTimeout( stepAnimation, 50 );
-
- }
- };
- stepAnimation();
- window.addEventListener( 'resize', onWindowResize );
- function onWindowResize() {
- renderer.setSize( window.innerWidth / 2, window.innerHeight );
- const aspect = ( window.innerWidth / 2 ) / window.innerHeight;
- const frustumHeight = camera.top - camera.bottom;
- camera.left = - frustumHeight * aspect / 2;
- camera.right = frustumHeight * aspect / 2;
- camera.updateProjectionMatrix();
- renderer.render( scene, camera );
- }
- }
- </script>
- </body>
- </html>
|