@plastic-software/three 0.180.0 → 0.181.2

package/examples/jsm/gpgpu/BitonicSort.js

@@ -0,0 +1,715 @@
+import { Fn, uvec2, If, instancedArray, instanceIndex, invocationLocalIndex, Loop, workgroupArray, workgroupBarrier, workgroupId, uint, select, min, max } from 'three/tsl';
+
+const StepType = {
+	NONE: 0,
+	// Swap all values within the local range of workgroupSize * 2
+	SWAP_LOCAL: 1,
+	DISPERSE_LOCAL: 2,
+	// Swap values within global data buffer.
+	FLIP_GLOBAL: 3,
+	DISPERSE_GLOBAL: 4,
+};
+
+
+/**
+ * Returns the indices that will be compared in a bitonic flip operation.
+ *
+ * @tsl
+ * @private
+ * @param {Node<uint>} index - The compute thread's invocation id.
+ * @param {Node<uint>} blockHeight - The height of the block within which elements are being swapped.
+ * @returns {Node<uvec2>} The indices of the elements in the data buffer being compared.
+ */
+export const getBitonicFlipIndices = /*@__PURE__*/ Fn( ( [ index, blockHeight ] ) => {
+
+	const blockOffset = ( index.mul( 2 ).div( blockHeight ) ).mul( blockHeight );
+	const halfHeight = blockHeight.div( 2 );
+	const idx = uvec2(
+		index.mod( halfHeight ),
+		blockHeight.sub( index.mod( halfHeight ) ).sub( 1 )
+	);
+	idx.x.addAssign( blockOffset );
+	idx.y.addAssign( blockOffset );
+
+	return idx;
+
+} ).setLayout( {
+	name: 'getBitonicFlipIndices',
+	type: 'uvec2',
+	inputs: [
+		{ name: 'index', type: 'uint' },
+		{ name: 'blockHeight', type: 'uint' }
+	]
+} );
+
+/**
+ * Returns the indices that will be compared in a bitonic sort's disperse operation.
+ *
+ * @tsl
+ * @private
+ * @param {Node<uint>} index - The compute thread's invocation id.
+ * @param {Node<uint>} swapSpan - The maximum span over which elements are being swapped.
+ * @returns {Node<uvec2>} The indices of the elements in the data buffer being compared.
+ */
+export const getBitonicDisperseIndices = /*@__PURE__*/ Fn( ( [ index, swapSpan ] ) => {
+
+	const blockOffset = ( ( index.mul( 2 ) ).div( swapSpan ) ).mul( swapSpan );
+	const halfHeight = swapSpan.div( 2 );
+	const idx = uvec2(
+		index.mod( halfHeight ),
+		( index.mod( halfHeight ) ).add( halfHeight )
+	);
+
+	idx.x.addAssign( blockOffset );
+	idx.y.addAssign( blockOffset );
+
+	return idx;
+
+} ).setLayout( {
+	name: 'getBitonicDisperseIndices',
+	type: 'uvec2',
+	inputs: [
+		{ name: 'index', type: 'uint' },
+		{ name: 'blockHeight', type: 'uint' }
+	]
+} );
+
+export class BitonicSort {
+
+	/**
+	 * Constructs a new light probe helper.
+	 *
+	 * @param {Renderer} renderer - The current scene's renderer.
+	 * @param {StorageBufferNode} dataBuffer - The data buffer to sort.
+	 * @param {Object} [options={}] - Options that modify the bitonic sort.
+	 */
+	constructor( renderer, dataBuffer, options = {} ) {
+
+		/**
+		 * A reference to the renderer.
+		 *
+		 * @type {Renderer}
+		 */
+		this.renderer = renderer;
+
+		/**
+		 * A reference to the StorageBufferNode holding the data that will be sorted  .
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.dataBuffer = dataBuffer;
+
+		/**
+		 * The size of the data.
+		 *
+		 * @type {StorageBufferNode}
+		 */
+		this.count = dataBuffer.value.count;
+
+		/**
+		 *
+		 * The size of each compute dispatch.
+		 * @type {number}
+		 */
+
+		 this.dispatchSize = this.count / 2;
+
+		/**
+		 * The workgroup size of the compute shaders executed during the sort.
+		 *
+		 * @type {StorageBufferNode}
+		*/
+		this.workgroupSize = options.workgroupSize ? Math.min( this.dispatchSize, options.workgroupSize ) : Math.min( this.dispatchSize, 64 );
+
+		/**
+		 * A node representing a workgroup scoped buffer that holds locally sorted elements.
+		 *
+		 * @type {WorkgroupInfoNode}
+		*/
+		this.localStorage = workgroupArray( dataBuffer.nodeType, this.workgroupSize * 2 );
+
+		this._tempArray = new Uint32Array( this.count );
+		for ( let i = 0; i < this.count; i ++ ) {
+
+			this._tempArray[ i ] = 0;
+
+		}
+
+		/**
+		 * A node representing a storage buffer used for transferring the result of the global sort back to the original data buffer.
+		 *
+		 * @type {StorageBufferNode}
+		*/
+		this.tempBuffer = instancedArray( this.count, dataBuffer.nodeType ).setName( 'TempStorage' );
+
+		/**
+		 * A node containing the current algorithm type, the current swap span, and the highest swap span.
+		 *
+		 * @type {StorageBufferNode}
+		*/
+		this.infoStorage = instancedArray( new Uint32Array( [ 1, 2, 2 ] ), 'uint' ).setName( 'BitonicSortInfo' );
+
+
+		/**
+		 * The number of distinct swap operations ('flips' and 'disperses') executed in an in-place
+		 * bitonic sort of the current data buffer.
+		 *
+		 * @type {number}
+		*/
+		this.swapOpCount = this._getSwapOpCount();
+
+		/**
+		 * The number of steps (i.e prepping and/or executing a swap) needed to fully execute an in-place bitonic sort of the current data buffer.
+		 *
+		 * @type {number}
+		*/
+		this.stepCount = this._getStepCount();
+
+		/**
+		 * The number of the buffer being read from.
+		 *
+		 * @type {string}
+		*/
+		this.readBufferName = 'Data';
+
+		/**
+		 * An object containing compute shaders that execute a 'flip' swap within a global address space on elements in the data buffer.
+		 *
+		 * @type {Object<string, ComputeNode>}
+		*/
+		this.flipGlobalNodes = {
+			'Data': this._getFlipGlobal( this.dataBuffer, this.tempBuffer ),
+			'Temp': this._getFlipGlobal( this.tempBuffer, this.dataBuffer )
+		};
+
+		/**
+		 * An object containing compute shaders that execute a 'disperse' swap within a global address space on elements in the data buffer.
+		 *
+		 * @type {Object<string, ComputeNode>}
+		*/
+		this.disperseGlobalNodes = {
+			'Data': this._getDisperseGlobal( this.dataBuffer, this.tempBuffer ),
+			'Temp': this._getDisperseGlobal( this.tempBuffer, this.dataBuffer )
+		};
+
+		/**
+		 * A compute shader that executes a sequence of flip and disperse swaps within a local address space on elements in the data buffer.
+		 *
+		 * @type {ComputeNode}
+		*/
+		this.swapLocalFn = this._getSwapLocal();
+
+		/**
+		 * A compute shader that executes a sequence of disperse swaps within a local address space on elements in the data buffer.
+		 *
+		 * @type {Object<string, ComputeNode>}
+		*/
+		this.disperseLocalNodes = {
+			'Data': this._getDisperseLocal( this.dataBuffer ),
+			'Temp': this._getDisperseLocal( this.tempBuffer ),
+		};
+
+		// Utility functions
+
+		/**
+		 * A compute shader that sets up the algorithm and the swap span for the next swap operation.
+		 *
+		 * @type {ComputeNode}
+		*/
+		this.setAlgoFn = this._getSetAlgoFn();
+
+		/**
+		 * A compute shader that aligns the result of the global swap operation with the current buffer.
+		 *
+		 * @type {ComputeNode}
+		*/
+		this.alignFn = this._getAlignFn();
+
+
+		/**
+		 * A compute shader that resets the algorithm and swap span information.
+		 *
+		 * @type {ComputeNode}
+		*/
+		this.resetFn = this._getResetFn();
+
+
+		/**
+		 * The current compute shader dispatch within the list of dispatches needed to complete the sort.
+		 *
+		 * @type {number}
+		*/
+		this.currentDispatch = 0;
+
+		/**
+		 * The number of global swap operations that must be executed before the sort
+		 * can swap in local address space.
+		 *
+		 * @type {number}
+		*/
+		this.globalOpsRemaining = 0;
+
+		/**
+		 * The total number of global operations needed to sort elements within the current swap span.
+		 *
+		 * @type {number}
+		*/
+		this.globalOpsInSpan = 0;
+
+
+	}
+
+	/**
+	 * Get total number of distinct swaps that occur in a bitonic sort.
+	 *
+	 * @private
+	 * @returns {number} - The total number of distinct swaps in a bitonic sort
+	 */
+	_getSwapOpCount() {
+
+		const n = Math.log2( this.count );
+		return ( n * ( n + 1 ) ) / 2;
+
+	}
+
+	/**
+	 * Get the number of steps it takes to execute a complete bitonic sort.
+	 *
+	 * @private
+	 * @returns {number} The number of steps it takes to execute a complete bitonic sort.
+	 */
+	_getStepCount() {
+
+		const logElements = Math.log2( this.count );
+		const logSwapSpan = Math.log2( this.workgroupSize * 2 );
+
+		const numGlobalFlips = logElements - logSwapSpan;
+
+		// Start with 1 for initial sort over all local elements
+		let numSteps = 1;
+		let numGlobalDisperses = 0;
+
+		for ( let i = 1; i <= numGlobalFlips; i ++ ) {
+
+			// Increment by the global flip that starts each global block
+			numSteps += 1;
+			// Increment by number of global disperses following the global flip
+			numSteps += numGlobalDisperses;
+			// Increment by local disperse that occurs after all global swaps are finished
+			numSteps += 1;
+
+			// Number of global disperse increases as swapSpan increases by factor of 2
+			numGlobalDisperses += 1;
+
+		}
+
+		return numSteps;
+
+	}
+
+	/**
+	 * Compares and swaps two data points in the data buffer within the global address space.
+	 * @param {Node<uint>} idxBefore - The index of the first data element in the data buffer.
+	 * @param {Node<uint>} idxAfter - The index of the second data element in the data buffer.
+	 * @param {StorageBufferNode} dataBuffer - The buffer of data to read from.
+	 * @param {StorageBufferNode} tempBuffer - The buffer of data to write to.
+	 * @private
+	 *
+	 */
+	_globalCompareAndSwapTSL( idxBefore, idxAfter, dataBuffer, tempBuffer ) {
+
+		const data1 = dataBuffer.element( idxBefore );
+		const data2 = dataBuffer.element( idxAfter );
+
+		tempBuffer.element( idxBefore ).assign( min( data1, data2 ) );
+		tempBuffer.element( idxAfter ).assign( max( data1, data2 ) );
+
+	}
+
+	/**
+	 * Compares and swaps two data points in the data buffer within the local address space.
+	 *
+	 * @private
+	 * @param {Node<uint>} idxBefore - The index of the first data element in the data buffer.
+	 * @param {Node<uint>} idxAfter - The index of the second data element in the data buffer
+	 */
+	_localCompareAndSwapTSL( idxBefore, idxAfter ) {
+
+		const { localStorage } = this;
+
+		const data1 = localStorage.element( idxBefore ).toVar();
+		const data2 = localStorage.element( idxAfter ).toVar();
+
+		localStorage.element( idxBefore ).assign( min( data1, data2 ) );
+		localStorage.element( idxAfter ).assign( max( data1, data2 ) );
+
+	}
+
+
+	/**
+	 * Create the compute shader that performs a global disperse swap on the data buffer.
+	 *
+	 * @private
+	 * @param {StorageBufferNode} readBuffer - The data buffer to read from.
+	 * @param {StorageBufferNode} writeBuffer - The data buffer to read from.
+	 * @returns {ComputeNode} - A compute shader that performs a global disperse swap on the data buffer.
+	 */
+	_getDisperseGlobal( readBuffer, writeBuffer ) {
+
+		const { infoStorage } = this;
+
+		const currentSwapSpan = infoStorage.element( 1 );
+
+		const fnDef = Fn( () => {
+
+			const idx = getBitonicDisperseIndices( instanceIndex, currentSwapSpan );
+			this._globalCompareAndSwapTSL( idx.x, idx.y, readBuffer, writeBuffer );
+
+		} )().compute( this.dispatchSize, [ this.workgroupSize ] );
+
+		return fnDef;
+
+	}
+
+	/**
+	 * Create the compute shader that performs a global flip swap on the data buffer.
+	 *
+	 * @private
+	 * @param {StorageBufferNode} readBuffer - The data buffer to read from.
+	 * @param {StorageBufferNode} writeBuffer - The data buffer to read from.
+	 * @returns {ComputeNode} - A compute shader that executes a global flip swap.
+	 */
+	_getFlipGlobal( readBuffer, writeBuffer ) {
+
+		const { infoStorage } = this;
+
+		const currentSwapSpan = infoStorage.element( 1 );
+
+		const fnDef = Fn( () => {
+
+			const idx = getBitonicFlipIndices( instanceIndex, currentSwapSpan );
+			this._globalCompareAndSwapTSL( idx.x, idx.y, readBuffer, writeBuffer );
+
+		} )().compute( this.dispatchSize, [ this.workgroupSize ] );
+
+		return fnDef;
+
+	}
+
+
+	/**
+	 * Create the compute shader that performs a complete local swap on the data buffer.
+	 *
+	 * @private
+	 * @returns {ComputeNode} - A compute shader that executes a full local swap.
+	 */
+	_getSwapLocal() {
+
+		const { localStorage, dataBuffer, workgroupSize } = this;
+
+		const fnDef = Fn( () => {
+
+			// Get ids of indices needed to populate workgroup local buffer.
+			// Use .toVar() to prevent these values from being recalculated multiple times.
+			const localOffset = uint( workgroupSize ).mul( 2 ).mul( workgroupId.x ).toVar();
+
+			const localID1 = invocationLocalIndex.mul( 2 );
+			const localID2 = invocationLocalIndex.mul( 2 ).add( 1 );
+
+			localStorage.element( localID1 ).assign( dataBuffer.element( localOffset.add( localID1 ) ) );
+			localStorage.element( localID2 ).assign( dataBuffer.element( localOffset.add( localID2 ) ) );
+
+			// Ensure that all local data has been populated
+			workgroupBarrier();
+
+			// Perform a chunk of the sort in a single pass that operates entirely in workgroup local space
+			// SWAP_LOCAL will always be first pass, so we start with known block height of 2
+			const flipBlockHeight = uint( 2 );
+
+			Loop( { start: uint( 2 ), end: uint( workgroupSize * 2 ), type: 'uint', condition: '<=', update: '<<= 1' }, () => {
+
+				// Ensure that last dispatch block executed
+				workgroupBarrier();
+
+				const flipIdx = getBitonicFlipIndices( invocationLocalIndex, flipBlockHeight );
+
+				this._localCompareAndSwapTSL( flipIdx.x, flipIdx.y );
+
+				const localBlockHeight = flipBlockHeight.div( 2 );
+
+				Loop( { start: localBlockHeight, end: uint( 1 ), type: 'uint', condition: '>', update: '>>= 1' }, () => {
+
+					// Ensure that last dispatch op executed
+					workgroupBarrier();
+
+					const disperseIdx = getBitonicDisperseIndices( invocationLocalIndex, localBlockHeight );
+					this._localCompareAndSwapTSL( disperseIdx.x, disperseIdx.y );
+
+					localBlockHeight.divAssign( 2 );
+
+				} );
+
+				// flipBlockHeight *= 2;
+				flipBlockHeight.shiftLeftAssign( 1 );
+
+			} );
+
+			// Ensure that all invocations have swapped their own regions of data
+			workgroupBarrier();
+
+			dataBuffer.element( localOffset.add( localID1 ) ).assign( localStorage.element( localID1 ) );
+			dataBuffer.element( localOffset.add( localID2 ) ).assign( localStorage.element( localID2 ) );
+
+		} )().compute( this.dispatchSize, [ this.workgroupSize ] );
+
+		return fnDef;
+
+	}
+
+	/**
+	 * Create the compute shader that performs a local disperse swap on the data buffer.
+	 *
+	 * @private
+	 * @param {StorageBufferNode} readWriteBuffer - The data buffer to read from and write to.
+	 * @returns {ComputeNode} - A compute shader that executes a local disperse swap.
+	 */
+	_getDisperseLocal( readWriteBuffer ) {
+
+		const { localStorage, workgroupSize } = this;
+
+		const fnDef = Fn( () => {
+
+			// Get ids of indices needed to populate workgroup local buffer.
+			// Use .toVar() to prevent these values from being recalculated multiple times.
+			const localOffset = uint( workgroupSize ).mul( 2 ).mul( workgroupId.x ).toVar();
+
+			const localID1 = invocationLocalIndex.mul( 2 );
+			const localID2 = invocationLocalIndex.mul( 2 ).add( 1 );
+
+			localStorage.element( localID1 ).assign( readWriteBuffer.element( localOffset.add( localID1 ) ) );
+			localStorage.element( localID2 ).assign( readWriteBuffer.element( localOffset.add( localID2 ) ) );
+
+			// Ensure that all local data has been populated
+			workgroupBarrier();
+
+			const localBlockHeight = uint( workgroupSize * 2 );
+
+			Loop( { start: localBlockHeight, end: uint( 1 ), type: 'uint', condition: '>', update: '>>= 1' }, () => {
+
+				// Ensure that last dispatch op executed
+				workgroupBarrier();
+
+				const disperseIdx = getBitonicDisperseIndices( invocationLocalIndex, localBlockHeight );
+				this._localCompareAndSwapTSL( disperseIdx.x, disperseIdx.y );
+
+				localBlockHeight.divAssign( 2 );
+
+			} );
+
+			// Ensure that all invocations have swapped their own regions of data
+			workgroupBarrier();
+
+			readWriteBuffer.element( localOffset.add( localID1 ) ).assign( localStorage.element( localID1 ) );
+			readWriteBuffer.element( localOffset.add( localID2 ) ).assign( localStorage.element( localID2 ) );
+
+		} )().compute( this.dispatchSize, [ this.workgroupSize ] );
+
+		return fnDef;
+
+	}
+
+	/**
+	 * Create the compute shader that resets the sort's algorithm information.
+	 *
+	 * @private
+	 * @returns {ComputeNode} - A compute shader that resets the bitonic sort's algorithm information.
+	 */
+	_getResetFn() {
+
+		const fnDef = Fn( () => {
+
+			const { infoStorage } = this;
+
+			const currentAlgo = infoStorage.element( 0 );
+			const currentSwapSpan = infoStorage.element( 1 );
+			const maxSwapSpan = infoStorage.element( 2 );
+
+			currentAlgo.assign( StepType.SWAP_LOCAL );
+			currentSwapSpan.assign( 2 );
+			maxSwapSpan.assign( 2 );
+
+		} )().compute( 1 );
+
+		return fnDef;
+
+	}
+
+	/**
+	 * Create the compute shader that copies the state of the last global swap to the data buffer.
+	 *
+	 * @private
+	 * @returns {ComputeNode} - A compute shader that copies the state of the last global swap to the data buffer.
+	 */
+	_getAlignFn() {
+
+		const { dataBuffer, tempBuffer } = this;
+
+		// TODO: Only do this in certain instances by ping-ponging which buffer gets sorted
+		// And only aligning if numDispatches % 2 === 1
+		const fnDef = Fn( () => {
+
+			dataBuffer.element( instanceIndex ).assign( tempBuffer.element( instanceIndex ) );
+
+		} )().compute( this.count, [ this.workgroupSize ] );
+
+		return fnDef;
+
+	}
+
+	/**
+	 * Create the compute shader that sets the bitonic sort algorithm's information.
+	 *
+	 * @private
+	 * @returns {ComputeNode} - A compute shader that sets the bitonic sort algorithm's information.
+	 */
+	_getSetAlgoFn() {
+
+		const fnDef = Fn( () => {
+
+			const { infoStorage, workgroupSize } = this;
+
+			const currentAlgo = infoStorage.element( 0 );
+			const currentSwapSpan = infoStorage.element( 1 );
+			const maxSwapSpan = infoStorage.element( 2 );
+
+			If( currentAlgo.equal( StepType.SWAP_LOCAL ), () => {
+
+				const nextHighestSwapSpan = uint( workgroupSize * 4 );
+
+				currentAlgo.assign( StepType.FLIP_GLOBAL );
+				currentSwapSpan.assign( nextHighestSwapSpan );
+				maxSwapSpan.assign( nextHighestSwapSpan );
+
+			} ).ElseIf( currentAlgo.equal( StepType.DISPERSE_LOCAL ), () => {
+
+				currentAlgo.assign( StepType.FLIP_GLOBAL );
+
+				const nextHighestSwapSpan = maxSwapSpan.mul( 2 );
+
+				currentSwapSpan.assign( nextHighestSwapSpan );
+				maxSwapSpan.assign( nextHighestSwapSpan );
+
+			} ).Else( () => {
+
+				const nextSwapSpan = currentSwapSpan.div( 2 );
+				currentAlgo.assign(
+					select(
+						nextSwapSpan.lessThanEqual( uint( workgroupSize * 2 ) ),
+						StepType.DISPERSE_LOCAL,
+						StepType.DISPERSE_GLOBAL
+					).uniformFlow()
+				);
+				currentSwapSpan.assign( nextSwapSpan );
+
+			} );
+
+		} )().compute( 1 );
+
+		return fnDef;
+
+	}
+
+	/**
+	 * Executes a step of the bitonic sort operation.
+	 *
+	 * @param {Renderer} renderer - The current scene's renderer.
+	 */
+	computeStep( renderer ) {
+
+		// Swap local only runs once
+		if ( this.currentDispatch === 0 ) {
+
+			renderer.compute( this.swapLocalFn );
+
+			this.globalOpsRemaining = 1;
+			this.globalOpsInSpan = 1;
+
+		} else if ( this.globalOpsRemaining > 0 ) {
+
+			const swapType = this.globalOpsRemaining === this.globalOpsInSpan ? 'Flip' : 'Disperse';
+
+			renderer.compute( swapType === 'Flip' ? this.flipGlobalNodes[ this.readBufferName ] : this.disperseGlobalNodes[ this.readBufferName ] );
+
+			if ( this.readBufferName === 'Data' ) {
+
+				this.readBufferName = 'Temp';
+
+			} else {
+
+				this.readBufferName = 'Data';
+
+			}
+
+			this.globalOpsRemaining -= 1;
+
+		} else {
+
+			// Then run local disperses when we've finished all global swaps
+			renderer.compute( this.disperseLocalNodes[ this.readBufferName ] );
+
+			const nextSpanGlobalOps = this.globalOpsInSpan + 1;
+			this.globalOpsInSpan = nextSpanGlobalOps;
+			this.globalOpsRemaining = nextSpanGlobalOps;
+
+		}
+
+
+		this.currentDispatch += 1;
+
+		if ( this.currentDispatch === this.stepCount ) {
+
+			// If our last swap addressed only addressed the temp buffer, then re-align it with the data buffer
+			// to fulfill the requirement of an in-place sort.
+			if ( this.readBufferName === 'Temp' ) {
+
+				renderer.compute( this.alignFn );
+				this.readBufferName = 'Data';
+
+			}
+
+			// Just reset the algorithm information
+			renderer.compute( this.resetFn );
+
+			this.currentDispatch = 0;
+			this.globalOpsRemaining = 0;
+			this.globalOpsInSpan = 0;
+
+		} else {
+
+			// Otherwise, determine what next swap span is
+			renderer.compute( this.setAlgoFn );
+
+		}
+
+	}
+
+	/**
+	 * Executes a complete bitonic sort on the data buffer.
+	 *
+	 * @param {Renderer} renderer - The current scene's renderer.
+	 */
+	compute( renderer ) {
+
+		this.globalOpsRemaining = 0;
+		this.globalOpsInSpan = 0;
+		this.currentDispatch = 0;
+
+		for ( let i = 0; i < this.stepCount; i ++ ) {
+
+			this.computeStep( renderer );
+
+		}
+
+	}
+
+}