playcanvas/build/playcanvas/src/scene/graphics/radix-sort/compute-radix-sort-onesweep.js

Open-source WebGL/WebGPU 3D engine for the web

import { Vec2 } from "../../../core/math/vec2.js";
import { Compute } from "../../../platform/graphics/compute.js";
import { Shader } from "../../../platform/graphics/shader.js";
import { StorageBuffer } from "../../../platform/graphics/storage-buffer.js";
import { BindGroupFormat, BindStorageBufferFormat, BindUniformBufferFormat } from "../../../platform/graphics/bind-group-format.js";
import { UniformBufferFormat, UniformFormat } from "../../../platform/graphics/uniform-buffer-format.js";
import { BUFFERUSAGE_COPY_DST, SHADERLANGUAGE_WGSL, SHADERSTAGE_COMPUTE, UNIFORMTYPE_UINT } from "../../../platform/graphics/constants.js";
import { onesweepGlobalHistSource } from "../../shader-lib/wgsl/chunks/radix-sort/onesweep-global-hist.js";
import { onesweepScanSource } from "../../shader-lib/wgsl/chunks/radix-sort/onesweep-scan.js";
import { onesweepBinningSource } from "../../shader-lib/wgsl/chunks/radix-sort/onesweep-binning.js";
import { ComputeRadixSortBase } from "./compute-radix-sort-base.js";
const D_DIM = 256;
const KEYS_PER_THREAD = 15;
const PART_SIZE = D_DIM * KEYS_PER_THREAD;
const G_HIST_DIM = 128;
const G_HIST_PART_SIZE = 32768;
const RADIX = 256;
const MAX_PASSES = 4;
class ComputeRadixSortOneSweep extends ComputeRadixSortBase {
	_threadBlocks = 0;
	_allocatedThreadBlocks = 0;
	_globalHist = null;
	_passHist = null;
	_index = null;
	_binningDispatchSize = new Vec2(1, 1);
	_globalHistDispatchSize = new Vec2(1, 1);
	_globalHistBindGroupFormat = null;
	_scanBindGroupFormat = null;
	_binningBindGroupFormat = null;
	_globalHistUniformFormat = null;
	_scanUniformFormat = null;
	_binningUniformFormat = null;
	_globalHistShader = null;
	_scanShader = null;
	_binningShader = null;
	_globalHistCompute = null;
	_scanCompute = null;
	_binningComputes = [];
	constructor(device, indirect = false) {
		super(device, indirect);
		const info = this._indirectInfo;
		info[0] = 2;
		info[1] = PART_SIZE;
		info[2] = G_HIST_PART_SIZE;
		this._globalHistUniformFormat = new UniformBufferFormat(device, [
			new UniformFormat("numKeys", UNIFORMTYPE_UINT),
			new UniformFormat("threadBlocks", UNIFORMTYPE_UINT),
			new UniformFormat("numPasses", UNIFORMTYPE_UINT),
			new UniformFormat("_pad", UNIFORMTYPE_UINT)
		]);
		this._scanUniformFormat = new UniformBufferFormat(device, [
			new UniformFormat("threadBlocks", UNIFORMTYPE_UINT),
			new UniformFormat("_pad0", UNIFORMTYPE_UINT),
			new UniformFormat("_pad1", UNIFORMTYPE_UINT),
			new UniformFormat("_pad2", UNIFORMTYPE_UINT)
		]);
		this._binningUniformFormat = new UniformBufferFormat(device, [
			new UniformFormat("numKeys", UNIFORMTYPE_UINT),
			new UniformFormat("threadBlocks", UNIFORMTYPE_UINT),
			new UniformFormat("pass_", UNIFORMTYPE_UINT),
			new UniformFormat("flags", UNIFORMTYPE_UINT)
		]);
		const minSubgroupSize = device.minSubgroupSize || device.maxSubgroupSize || 32;
		const maxSubgroups = Math.max(1, Math.ceil(256 / minSubgroupSize));
		const suffix = indirect ? "Indirect" : "";
		const histGroupEntries = [
			new BindStorageBufferFormat("b_sort", SHADERSTAGE_COMPUTE, true),
			new BindStorageBufferFormat("b_globalHist", SHADERSTAGE_COMPUTE, false),
			new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE)
		];
		const scanGroupEntries = [
			new BindStorageBufferFormat("b_globalHist", SHADERSTAGE_COMPUTE, true),
			new BindStorageBufferFormat("b_passHist", SHADERSTAGE_COMPUTE, false),
			new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE)
		];
		const binGroupEntries = [
			new BindStorageBufferFormat("inputKeys", SHADERSTAGE_COMPUTE, true),
			new BindStorageBufferFormat("outputKeys", SHADERSTAGE_COMPUTE, false),
			new BindStorageBufferFormat("inputValues", SHADERSTAGE_COMPUTE, true),
			new BindStorageBufferFormat("outputValues", SHADERSTAGE_COMPUTE, false),
			new BindStorageBufferFormat("b_passHist", SHADERSTAGE_COMPUTE, false),
			new BindStorageBufferFormat("b_index", SHADERSTAGE_COMPUTE, false),
			new BindUniformBufferFormat("uniforms", SHADERSTAGE_COMPUTE)
		];
		if (indirect) {
			histGroupEntries.push(new BindStorageBufferFormat("b_sortElementCount", SHADERSTAGE_COMPUTE, true));
			scanGroupEntries.push(new BindStorageBufferFormat("b_sortElementCount", SHADERSTAGE_COMPUTE, true));
			binGroupEntries.push(new BindStorageBufferFormat("b_sortElementCount", SHADERSTAGE_COMPUTE, true));
		}
		this._globalHistBindGroupFormat = new BindGroupFormat(device, histGroupEntries);
		this._scanBindGroupFormat = new BindGroupFormat(device, scanGroupEntries);
		this._binningBindGroupFormat = new BindGroupFormat(device, binGroupEntries);
		const histDefines = /* @__PURE__ */ new Map();
		histDefines.set("{G_HIST_DIM}", G_HIST_DIM);
		histDefines.set("{G_HIST_PART_SIZE}", G_HIST_PART_SIZE);
		if (indirect) histDefines.set("USE_INDIRECT_SORT", "");
		this._globalHistShader = new Shader(device, {
			name: `OneSweepGlobalHist${suffix}`,
			shaderLanguage: SHADERLANGUAGE_WGSL,
			cshader: onesweepGlobalHistSource,
			cdefines: histDefines,
			computeBindGroupFormat: this._globalHistBindGroupFormat,
			computeUniformBufferFormats: { uniforms: this._globalHistUniformFormat }
		});
		const scanDefines = /* @__PURE__ */ new Map();
		scanDefines.set("{MAX_SUBGROUPS}", maxSubgroups);
		scanDefines.set("{PART_SIZE}", PART_SIZE);
		if (indirect) scanDefines.set("USE_INDIRECT_SORT", "");
		this._scanShader = new Shader(device, {
			name: `OneSweepScan${suffix}`,
			shaderLanguage: SHADERLANGUAGE_WGSL,
			cshader: onesweepScanSource,
			cdefines: scanDefines,
			computeBindGroupFormat: this._scanBindGroupFormat,
			computeUniformBufferFormats: { uniforms: this._scanUniformFormat }
		});
		const binDefines = /* @__PURE__ */ new Map();
		binDefines.set("{D_DIM}", D_DIM);
		binDefines.set("{KEYS_PER_THREAD}", KEYS_PER_THREAD);
		binDefines.set("{MAX_SUBGROUPS}", maxSubgroups);
		if (indirect) binDefines.set("USE_INDIRECT_SORT", "");
		this._binningShader = new Shader(device, {
			name: `OneSweepBinning${suffix}`,
			shaderLanguage: SHADERLANGUAGE_WGSL,
			cshader: onesweepBinningSource,
			cdefines: binDefines,
			computeBindGroupFormat: this._binningBindGroupFormat,
			computeUniformBufferFormats: { uniforms: this._binningUniformFormat }
		});
		this._globalHistCompute = new Compute(device, this._globalHistShader, this._globalHistShader.name);
		this._scanCompute = new Compute(device, this._scanShader, this._scanShader.name);
	}
	destroy() {
		this._destroyBuffers();
		this._globalHistShader?.destroy();
		this._scanShader?.destroy();
		this._binningShader?.destroy();
		this._globalHistBindGroupFormat?.destroy();
		this._scanBindGroupFormat?.destroy();
		this._binningBindGroupFormat?.destroy();
		this._globalHistShader = null;
		this._scanShader = null;
		this._binningShader = null;
		this._globalHistCompute = null;
		this._scanCompute = null;
		this._binningComputes.length = 0;
		this._globalHistBindGroupFormat = null;
		this._scanBindGroupFormat = null;
		this._binningBindGroupFormat = null;
		this._globalHistUniformFormat = null;
		this._scanUniformFormat = null;
		this._binningUniformFormat = null;
		super.destroy();
	}
	_destroyBuffers() {
		this._destroyPingPongBuffers();
		this._globalHist?.destroy();
		this._passHist?.destroy();
		this._index?.destroy();
		this._globalHist = null;
		this._passHist = null;
		this._index = null;
		this._allocatedThreadBlocks = 0;
		this._threadBlocks = 0;
	}
	get radixBits() {
		return 8;
	}
	_ensureBinningComputes(numPasses) {
		while (this._binningComputes.length < numPasses) {
			this._binningComputes.push(new Compute(this.device, this._binningShader, `OneSweepBinning-${this._binningComputes.length}`));
		}
	}
	_allocateBuffers(elementCount, forceRealloc = false) {
		const effectiveCount = Math.max(elementCount, this.capacity);
		const allocThreadBlocks = Math.max(1, Math.ceil(effectiveCount / PART_SIZE));
		const currentThreadBlocks = Math.max(1, Math.ceil(elementCount / PART_SIZE));
		const needRealloc = forceRealloc || allocThreadBlocks !== this._allocatedThreadBlocks || !this._keys0;
		if (needRealloc) {
			this._destroyBuffers();
			this._allocatedThreadBlocks = allocThreadBlocks;
			this.capacity = effectiveCount;
			const device = this.device;
			this._allocatePingPongElementBuffers(effectiveCount);
			this._globalHist = new StorageBuffer(device, MAX_PASSES * RADIX * 4, BUFFERUSAGE_COPY_DST);
			this._passHist = new StorageBuffer(device, MAX_PASSES * allocThreadBlocks * RADIX * 4, BUFFERUSAGE_COPY_DST);
			this._index = new StorageBuffer(device, MAX_PASSES * 4, BUFFERUSAGE_COPY_DST);
		}
		this._threadBlocks = currentThreadBlocks;
		const maxPerDim = this.device.limits.maxComputeWorkgroupsPerDimension || 65535;
		Compute.calcDispatchSize(currentThreadBlocks, this._binningDispatchSize, maxPerDim);
		const histBlocks = Math.max(1, Math.ceil(elementCount / G_HIST_PART_SIZE));
		Compute.calcDispatchSize(histBlocks, this._globalHistDispatchSize, maxPerDim);
	}
	sort(keysBuffer, elementCount, numBits = 16, initialValues, skipLastPassKeyWrite = false, destructiveKeys = false) {
		const numPasses = numBits / 8;
		const hasInitialValues = !!initialValues;
		this._elementCount = elementCount;
		this._numBits = numBits;
		this._hasInitialValues = hasInitialValues;
		this._skipLastPassKeyWrite = skipLastPassKeyWrite;
		const prevDestructiveKeys = this._destructiveKeys;
		this._destructiveKeys = destructiveKeys;
		this._allocateBuffers(elementCount, destructiveKeys !== prevDestructiveKeys);
		this._ensureBinningComputes(numPasses);
		if (destructiveKeys) {
			this._keys1 = keysBuffer;
		}
		const device = this.device;
		this._globalHist.clear();
		this._passHist.clear(0, MAX_PASSES * this._threadBlocks * RADIX * 4);
		this._index.clear();
		const histCompute = this._globalHistCompute;
		histCompute.setParameter("b_sort", keysBuffer);
		histCompute.setParameter("b_globalHist", this._globalHist);
		histCompute.setParameter("numKeys", elementCount);
		histCompute.setParameter("threadBlocks", this._threadBlocks);
		histCompute.setParameter("numPasses", numPasses);
		histCompute.setParameter("_pad", 0);
		histCompute.setupDispatch(this._globalHistDispatchSize.x, this._globalHistDispatchSize.y, 1);
		device.computeDispatch([histCompute], "OneSweep-GlobalHist");
		const scanCompute = this._scanCompute;
		scanCompute.setParameter("b_globalHist", this._globalHist);
		scanCompute.setParameter("b_passHist", this._passHist);
		scanCompute.setParameter("threadBlocks", this._threadBlocks);
		scanCompute.setParameter("_pad0", 0);
		scanCompute.setParameter("_pad1", 0);
		scanCompute.setParameter("_pad2", 0);
		scanCompute.setupDispatch(numPasses, 1, 1);
		device.computeDispatch([scanCompute], "OneSweep-Scan");
		let currentKeys = keysBuffer;
		let currentValues = initialValues ?? this._values0;
		let nextKeys = this._keys0;
		let nextValues = this._values1;
		for (let pass = 0; pass < numPasses; pass++) {
			const isFirstPass = pass === 0 && !hasInitialValues;
			const isLastPass = pass === numPasses - 1;
			const flags = (isFirstPass ? 1 : 0) | (isLastPass && skipLastPassKeyWrite ? 2 : 0);
			const binCompute = this._binningComputes[pass];
			binCompute.setParameter("inputKeys", currentKeys);
			binCompute.setParameter("outputKeys", nextKeys);
			binCompute.setParameter("inputValues", currentValues);
			binCompute.setParameter("outputValues", nextValues);
			binCompute.setParameter("b_passHist", this._passHist);
			binCompute.setParameter("b_index", this._index);
			binCompute.setParameter("numKeys", elementCount);
			binCompute.setParameter("threadBlocks", this._threadBlocks);
			binCompute.setParameter("pass_", pass);
			binCompute.setParameter("flags", flags);
			binCompute.setupDispatch(this._binningDispatchSize.x, this._binningDispatchSize.y, 1);
			device.computeDispatch([binCompute], `OneSweep-Binning-${pass}`);
			if (!isLastPass) {
				currentKeys = nextKeys;
				nextKeys = currentKeys === this._keys0 ? this._keys1 : this._keys0;
				const t = currentValues;
				currentValues = nextValues;
				nextValues = t;
			}
		}
		return this.sortedIndices;
	}
	sortIndirect(keysBuffer, maxElementCount, numBits, sortSlotBase, sortElementCountBuffer, initialValues, skipLastPassKeyWrite = false, destructiveKeys = false) {
		const numPasses = numBits / 8;
		const hasInitialValues = !!initialValues;
		this._elementCount = maxElementCount;
		this._numBits = numBits;
		this._hasInitialValues = hasInitialValues;
		this._skipLastPassKeyWrite = skipLastPassKeyWrite;
		const prevDestructiveKeys = this._destructiveKeys;
		this._destructiveKeys = destructiveKeys;
		this._allocateBuffers(maxElementCount, destructiveKeys !== prevDestructiveKeys);
		this._ensureBinningComputes(numPasses);
		if (destructiveKeys) {
			this._keys1 = keysBuffer;
		}
		const device = this.device;
		this._globalHist.clear();
		this._passHist.clear(0, MAX_PASSES * this._allocatedThreadBlocks * RADIX * 4);
		this._index.clear();
		const histCompute = this._globalHistCompute;
		histCompute.setParameter("b_sort", keysBuffer);
		histCompute.setParameter("b_globalHist", this._globalHist);
		histCompute.setParameter("b_sortElementCount", sortElementCountBuffer);
		histCompute.setParameter("numKeys", 0);
		histCompute.setParameter("threadBlocks", 0);
		histCompute.setParameter("numPasses", numPasses);
		histCompute.setParameter("_pad", 0);
		histCompute.setupIndirectDispatch(sortSlotBase + 1);
		device.computeDispatch([histCompute], "OneSweep-GlobalHistIndirect");
		const scanCompute = this._scanCompute;
		scanCompute.setParameter("b_globalHist", this._globalHist);
		scanCompute.setParameter("b_passHist", this._passHist);
		scanCompute.setParameter("b_sortElementCount", sortElementCountBuffer);
		scanCompute.setParameter("threadBlocks", 0);
		scanCompute.setParameter("_pad0", 0);
		scanCompute.setParameter("_pad1", 0);
		scanCompute.setParameter("_pad2", 0);
		scanCompute.setupDispatch(numPasses, 1, 1);
		device.computeDispatch([scanCompute], "OneSweep-ScanIndirect");
		let currentKeys = keysBuffer;
		let currentValues = initialValues ?? this._values0;
		let nextKeys = this._keys0;
		let nextValues = this._values1;
		for (let pass = 0; pass < numPasses; pass++) {
			const isFirstPass = pass === 0 && !hasInitialValues;
			const isLastPass = pass === numPasses - 1;
			const flags = (isFirstPass ? 1 : 0) | (isLastPass && skipLastPassKeyWrite ? 2 : 0);
			const binCompute = this._binningComputes[pass];
			binCompute.setParameter("inputKeys", currentKeys);
			binCompute.setParameter("outputKeys", nextKeys);
			binCompute.setParameter("inputValues", currentValues);
			binCompute.setParameter("outputValues", nextValues);
			binCompute.setParameter("b_passHist", this._passHist);
			binCompute.setParameter("b_index", this._index);
			binCompute.setParameter("b_sortElementCount", sortElementCountBuffer);
			binCompute.setParameter("numKeys", 0);
			binCompute.setParameter("threadBlocks", 0);
			binCompute.setParameter("pass_", pass);
			binCompute.setParameter("flags", flags);
			binCompute.setupIndirectDispatch(sortSlotBase);
			device.computeDispatch([binCompute], `OneSweep-BinningIndirect-${pass}`);
			if (!isLastPass) {
				currentKeys = nextKeys;
				nextKeys = currentKeys === this._keys0 ? this._keys1 : this._keys0;
				const t = currentValues;
				currentValues = nextValues;
				nextValues = t;
			}
		}
		return this.sortedIndices;
	}
}
export {
	ComputeRadixSortOneSweep
};