WebGPU-Framework/demos/sort/demo.js

import Shader 			from "../../framework/WebGpu.js"

import Measure			from "../../framework/Measure.js"


const adapter 		= await navigator.gpu.requestAdapter();

if (!adapter ) {

	throw new Error("Failed to get GPU adapter");

}

const particleCount 	= 131072;

const workgroupSize 	= 256;

const numWorkgroups 	= Math.ceil( particleCount / workgroupSize );

const device 			= await adapter.requestDevice();


var passIndex 			= 0;

var jArray 				= new Array();

var kArray 				= new Array();


for ( let k = 512; k <= particleCount; k <<= 1 ) {

	for ( let j = k >> 1; j > 0; j >>= 1 ) {

		jArray.push( j );

		kArray.push( k );	

	}

}

const indices 				= new Uint32Array( particleCount );

const threadPassIndices 	= new Uint32Array( particleCount );

for (var i = 0; i < particleCount; i++) {

	indices[particleCount-i-1] 	= i;

	threadPassIndices[0]		= 0;

}

const localSortShader = new Shader( device, "./shaders/localSort.wgsl");

await localSortShader.addStage("main", GPUShaderStage.COMPUTE );

localSortShader.setVariable( "compare", indices );

localSortShader.setVariable( "totalCount", particleCount );


const bitonicSortGridHashShader = new Shader( device, "./shaders/bitonicSortUIntMultiPass.wgsl");

await bitonicSortGridHashShader.addStage("main", GPUShaderStage.COMPUTE );

bitonicSortGridHashShader.setVariable( "totalCount", particleCount );

bitonicSortGridHashShader.setBuffer( "compare", localSortShader.getBuffer("compare") );

bitonicSortGridHashShader.setVariable( "jArray", jArray );

bitonicSortGridHashShader.setVariable( "kArray", kArray );


const measure = new Measure();

measure.writeToPage = true;

measure.element		= document.querySelector(".result");

for (var i = 3; i < 59; i++) {

	bitonicSortGridHashShader.setVariable("threadPassIndices", threadPassIndices);
	
	measure.start( "sort"+i );

	//await localSortShader.execute( numWorkgroups );

const commandEncoder = device.createCommandEncoder();


// === Local sort ===
{
	const passEncoder = commandEncoder.beginComputePass();

	passEncoder.setPipeline( localSortShader.pipeline );

	for ( const [ index, bindGroup ] of localSortShader.bindGroups.entries() ) {

		passEncoder.setBindGroup( index, bindGroup );

	}

	passEncoder.dispatchWorkgroups( numWorkgroups );

	passEncoder.end();
}

// === Bitonic global merge ===
{
	const passEncoder = commandEncoder.beginComputePass();

	passEncoder.setPipeline( bitonicSortGridHashShader.pipeline );

	for ( const [ index, bindGroup ] of bitonicSortGridHashShader.bindGroups.entries() ) {

		passEncoder.setBindGroup( index, bindGroup );

	}

	var numberOfPasses = 0;

	for ( let k = 512; k <= particleCount; k <<= 1 ) {

		for ( let j = k >> 1; j > 0; j >>= 1 ) {

			numberOfPasses++;

			passEncoder.dispatchWorkgroups( numWorkgroups );

		}

	}

	passEncoder.end();
	
}


const commandBuffer = commandEncoder.finish();

device.queue.submit( [ commandBuffer ] );

	await device.queue.onSubmittedWorkDone();

	measure.end( "sort"+i );

}


var result = await bitonicSortGridHashShader.debugBuffer("compare");

await bitonicSortGridHashShader.debugBuffer("threadPassIndices");

var error = false;

for (var i = 0; i < particleCount; i++ ) {

	if( result[i] != i ) {

		error = true;

	} 

}

console.log("There is error?", error, particleCount);