Accueil / npm / three@0.186.1 / examples / jsm / gpgpu / BitonicSort.js
BitonicSort.js
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 );
}
}
}