three.js WebGPU BitonicSort：基于 Compute Shader 的 GPU 原地双调排序实现与实战

本文围绕 three.js 官方文档中的 BitonicSort 类展开，深入剖析 examples/jsm/gpgpu/BitonicSort.js 中这套 GPU 双调排序（bitonic sort）的完整实现：flip/disperse 两类交换的原语设计、workgroup 局部地址空间与全局存储缓冲的双空间调度、双缓冲乒乓（ping-pong）与原地对齐机制，以及配套示例 examples/webgpu_compute_sort_bitonic.html 的可复现用法。读完后你将掌握如何用 TSL（three Shading Language）编写多步 compute dispatch 完成一次完整的 GPU 排序，并理解每一步 dispatch 在做什么。

一、API 总览：new BitonicSort( renderer, dataBuffer, options )

官方 API 文档 docs/pages/BitonicSort_BitonicSort.html.md 给出的核心签名如下：

import { BitonicSort } from 'three/addons/gpgpu/BitonicSort.js';

const bitonicSort = new BitonicSort( renderer, dataBuffer, options = {} );

三个构造参数的官方说明与源码级补充：

参数	官方文档说明	源码级补充
renderer	The current scene's renderer（当前场景的渲染器）	保存为 this.renderer，但实际执行排序时是通过 computeStep( renderer ) / compute( renderer ) 每次传入的 WebGPURenderer
dataBuffer	The data buffer to sort（待排序的数据缓冲）	必须是一个 StorageBufferNode；构造函数直接读取 dataBuffer.value.count 得到元素个数 count，读取 dataBuffer.nodeType 决定临时缓冲与 workgroup 局部数组的标量类型（见 构造函数）
options	Options that modify the bitonic sort，默认 {}	目前仅识别 workgroupSize：options.workgroupSize ? Math.min( dispatchSize, options.workgroupSize ) : Math.min( dispatchSize, 64 )，即不传时默认工作组大小为 64（BitonicSort.js#L122）

需要说明一点：官方文档中“Constructs a new light probe helper”这句描述是源码 JSDoc 中被复制误用的文字（构造函数 JSDoc 同样如此），该类实际构造的是一个 GPU 双调排序模块，与光探针无关。

除构造参数外，构造函数还会在实例上预创建整套排序所需的资源与 compute 节点，理解这些成员是读懂后续调度逻辑的关键：

• this.count：元素个数，取自 dataBuffer.value.count；
• this.dispatchSize = count / 2：每次 compute dispatch 的线程数——双调排序每一步恰好处理 count/2 对元素，因此每个线程负责一对；
• this.workgroupSize：如上表，默认 64；
• this.localStorage：workgroupArray( dataBuffer.nodeType, workgroupSize * 2 )，一个 workgroup 作用域的局部缓冲，每个工作组一次缓存 2 * workgroupSize 个元素（BitonicSort.js#L129）；
• this.tempBuffer：instancedArray( count, dataBuffer.nodeType ) 命名的 TempStorage，与数据缓冲同大小的第二块存储，用于全局交换的乒乓写入（BitonicSort.js#L143）；
• this.infoStorage：new Uint32Array( [ 1, 2, 2 ] ) 的 3 元素 uint 缓冲，三个分量分别是【当前算法类型、当前交换跨度 swap span、最大交换跨度】，初值 (1, 2, 2) 对应 SWAP_LOCAL 起始态（BitonicSort.js#L150）；
• this.swapOpCount / this.stepCount：排序所需的总交换操作数与总 dispatch 步数，公式见第五节；
• this.readBufferName：'Data' 或 'Temp'，标记当前读取源缓冲，实现乒乓切换；
• flipGlobalNodes / disperseGlobalNodes：各含 'Data'、'Temp' 两个方向的 global 交换 compute 节点；
• swapLocalFn：一次性完成“整块局部排序”的 compute 节点；
• disperseLocalNodes：局部 disperse 交换节点（同样分 Data/Temp 两个方向）；
• setAlgoFn / resetFn / alignFn：三个单线程（compute( 1 )）或全量拷贝的工具节点。

二、两类交换的索引映射：flip 与 disperse

双调排序的核心动作只有两种，它们都由纯 TSL 函数计算“当前线程该比较哪一对下标”，这两个函数被导出供外部复用（getBitonicFlipIndices 与 getBitonicDisperseIndices）：

// flip：块内首尾对称配对（用于把两个单调序列合并成一个双调序列）
export const getBitonicFlipIndices = 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;

} );

// disperse：块内上半区与下半区按固定距离配对
export const getBitonicDisperseIndices = 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;

} );

两者都可以归纳为：给定线程编号 index 与当前块高 blockHeight（即 swap span），先定位所在块起点 blockOffset，再在块内取出一对对称下标 (x, y)。由于每次 dispatch 的线程数是 count / 2，每个线程天然负责一对元素——flip 让线程负责“块内首尾镜像位”（i 与 blockHeight - 1 - i），disperse 让线程负责“上下半区等距位”（i 与 i + halfHeight）。两个函数都通过 .setLayout( { name, type, inputs } ) 声明了 TSL 布局元数据，便于在编辑器类工具中被识别。

三、双地址空间策略：local swap 与 global swap

StepType 定义了五种状态（BitonicSort.js#L3-L11）：

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,
};

设计的出发点是把比较交换尽可能放在 GPU 上最快的地方做：

1. 局部交换（SWAP_LOCAL / DISPERSE_LOCAL）

当当前 swap span 小于等于 workgroupSize * 2 时，整个交换块都能装进一个 workgroup 的局部内存。swapLocalFn（BitonicSort.js#L406-L467）的做法是：

1. 每个线程从全局缓冲搬入自己的两个元素到 localStorage（localID1 = invocationLocalIndex * 2，localID2 = localID1 + 1），随后 workgroupBarrier()；
2. 在局部内存上执行嵌套循环：外层 flipBlockHeight 从 2 每次左移一位到 workgroupSize * 2，每轮先做一次 flip 交换，再对 localBlockHeight 从一半开始做递减的 disperse 交换，每轮之间 workgroupBarrier() 同步；
3. 最后再 workgroupBarrier() 并把结果写回全局缓冲。

局部比较交换的 TSL 原语只有一对 min/max：

_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 ) );

}

2. 全局交换（FLIP_GLOBAL / DISPERSE_GLOBAL）

当 swap span 超出单个 workgroup 的容量时，线程只能操作全局存储。_getFlipGlobal / _getDisperseGlobal 生成的节点结构相同：先用 getBitonicFlipIndices( instanceIndex, currentSwapSpan ) 或 disperse 版本算出下标对，再执行跨缓冲的比较交换（BitonicSort.js#L319-L327）：

_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 ) );

}

这里有一个重要的并发设计：读一个缓冲、写另一个缓冲（Data→Temp 或 Temp→Data），而不是原地写。双调排序中一对元素只会被某一个线程访问，但由于交换结果落在“另一个”缓冲里，读侧与写侧互不冲突，无需任何原子操作即可安全并行。代价是数据在两块缓冲间来回流动，这正是下一节乒乓调度的由来。

四、状态机：infoStorage 驱动的多步调度

排序被拆成多个独立的 compute dispatch，每步之间由 CPU 侧 computeStep( renderer ) 推进。GPU 侧的“下一步做什么”不靠 JS 变量传递，而是存在 infoStorage（[currentAlgo, currentSwapSpan, maxSwapSpan]）里，由单线程节点 setAlgoFn 更新（BitonicSort.js#L575-L620）：

• 当前为 SWAP_LOCAL：转入 FLIP_GLOBAL，span 抬升到 workgroupSize * 4（首个必须走全局的块）；
• 当前为 DISPERSE_LOCAL：转入 FLIP_GLOBAL，maxSwapSpan *= 2（进入下一个更大的双调合并阶段）；
• 其他情况（FLIP_GLOBAL / DISPERSE_GLOBAL）：span 减半，若 span/2 <= workgroupSize * 2 则降级为 DISPERSE_LOCAL（回到局部地址空间），否则保持 DISPERSE_GLOBAL。

这个转移规则恰好刻画了双调排序的自底向上结构：从 2 元素块开始逐级翻倍合并，合并过程中的“收敛”阶段（span 递减）尽量留在 workgroup 局部空间完成，只有跨度超出局部容量时才升级到全局交换。

CPU 侧的调度器 computeStep 逻辑如下：

computeStep( renderer ) {

	// 第 0 步：一次性完成所有 2*workgroupSize 元素块的局部排序
	if ( this.currentDispatch === 0 ) {

		renderer.compute( this.swapLocalFn );
		this.globalOpsRemaining = 1;
		this.globalOpsInSpan = 1;

	} else if ( this.globalOpsRemaining > 0 ) {

		// 全局阶段：span 内首个操作是 Flip，其余全是 Disperse
		const swapType = this.globalOpsRemaining === this.globalOpsInSpan ? 'Flip' : 'Disperse';
		renderer.compute( swapType === 'Flip'
			? this.flipGlobalNodes[ this.readBufferName ]
			: this.disperseGlobalNodes[ this.readBufferName ] );

		// 乒乓：读 Data 就写 Temp，读 Temp 就写 Data
		this.readBufferName = this.readBufferName === 'Data' ? 'Temp' : 'Data';
		this.globalOpsRemaining -= 1;

	} else {

		// 全局交换全部完成后，执行本地 disperse 收敛
		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 ) {

		// 若最终结果落在 Temp 缓冲，则整体拷回 Data 保证“原地”语义
		if ( this.readBufferName === 'Temp' ) {
			renderer.compute( this.alignFn );
			this.readBufferName = 'Data';
		}

		renderer.compute( this.resetFn );
		this.currentDispatch = 0;
		this.globalOpsRemaining = 0;
		this.globalOpsInSpan = 0;

	} else {

		renderer.compute( this.setAlgoFn ); // 单线程更新 infoStorage
	}

}

要点归纳：

• 第 0 步永远是 swapLocalFn：它把每个 2 * workgroupSize 大小的块排好序，作为后续所有双调合并的“已排好序的砖块”；
• 全局阶段呈“一次 Flip + 若干次 Disperse”：globalOpsInSpan 记录当前 span 内的全局操作总数，首个操作必为 Flip，其后是 Disperse，且随 span 翻倍每次 +1；
• 乒乓与对齐：每次全局交换后 readBufferName 翻转；若整个排序结束时结果停在 Temp 缓冲，则执行 alignFn（全量 dataBuffer[i] = tempBuffer[i]）把数据搬回原缓冲，从而保证排序对外呈现为原地（in-place）。源码中留有 TODO 注释：可以改为仅在 numDispatches % 2 === 1 时对齐，进一步省掉一次全量拷贝；
• 可重复使用：一次 compute 结束后内部计数器归零、infoStorage 被 resetFn 重置为 (SWAP_LOCAL, 2, 2)，因此同一实例可对同一缓冲反复排序。

对外最简单的入口是一次跑完整个排序（compute）：

compute( renderer ) {

	this.globalOpsRemaining = 0;
	this.globalOpsInSpan = 0;
	this.currentDispatch = 0;

	for ( let i = 0; i < this.stepCount; i ++ ) {

		this.computeStep( renderer );

	}

}

五、步数与交换次数的数学

构造函数用两个纯数学公式预算工作量（BitonicSort.js#L268-L308）：

_getSwapOpCount() {

	const n = Math.log2( this.count );
	return ( n * ( n + 1 ) ) / 2;

}

swapOpCount = n(n+1)/2（n = log2(count)）就是经典双调排序的比较交换轮数。而 stepCount 进一步考虑了“局部空间能吞掉多少交换”：

const logElements = Math.log2( this.count );
const logSwapSpan = Math.log2( this.workgroupSize * 2 );
const numGlobalFlips = logElements - logSwapSpan;

let numSteps = 1;               // 第一步整块局部排序
let numGlobalDisperses = 0;
for ( let i = 1; i <= numGlobalFlips; i ++ ) {

	numSteps += 1;                // 每个全局块以一次 Flip 开头
	numSteps += numGlobalDisperses; // 紧随其后的全局 Disperse 数量
	numSteps += 1;                // 全局交换完毕后的本地 Disperse 收敛
	numGlobalDisperses += 1;      // span 翻倍后，全局 Disperse 数量 +1

}

以官方示例的规模代入可以具体化：count = 16384 = 2^14、workgroupSize = 64，则 logElements = 14、logSwapSpan = 7、numGlobalFlips = 7，得到 swapOpCount = 14*15/2 = 105 次比较交换、stepCount = 36 次 compute dispatch。也就是说整个排序在 CPU 看来只是 36 次 renderer.compute( ... ) 调用。

六、实战用法：webgpu_compute_sort_bitonic.html

配套示例 examples/webgpu_compute_sort_bitonic.html 把 16384 个 uint 元素排成 128×128 的可视化网格，左半屏运行 BitonicSort 模块（混合 local/global 交换），右半屏运行一段仅用全局交换的参考实现（复用导出的 getBitonicFlipIndices / getBitonicDisperseIndices），并用 infoStorage 驱动颜色高亮显示当前正在比较的区域——这正是第二节截图所展示的界面。

最小可复现的关键代码路径如下。首先通过 importmap 引入 WebGPU 构建与 TSL：

<script type="importmap">
	{
		"imports": {
			"three/webgpu": "../build/three.webgpu.js",
			"three/tsl": "../build/three.tsl.js",
			"three/addons/": "./jsm/"
		}
	}
</script>

import { BitonicSort } from 'three/addons/gpgpu/BitonicSort.js';
import { storage } from 'three/tsl';

// 1. 准备数据：16384 个随机打乱的 uint（2 的幂）
const array = new Uint32Array( 16384 );
// ... randomizeDataArray( array ) ...

// 2. 包装为 StorageBufferNode；setPBO( true ) 便于 WebGPU 之外的回退路径使用
const currentElementsBuffer = new THREE.StorageInstancedBufferAttribute( array, 1 );
const currentElementsStorage = storage( currentElementsBuffer, 'uint', 16384 ).setPBO( true ).setName( 'Elements' );

// 3. 实例化排序器（workgroupSize 缺省即 64，此处显式写出）
const bitonicSort = new BitonicSort( renderer, currentElementsStorage, { workgroupSize: 64 } );

// 一次性完成排序：
bitonicSort.compute( renderer );

// 或逐帧单步推进，用于可视化（示例每 100ms 推进一步，每步读取
// bitonicSort.infoStorage 来渲染“当前交换区间”的高亮）：
bitonicSort.computeStep( renderer );

示例中还有几个值得注意的细节：

• 前置能力检查：WebGPU.isAvailable()（来自 examples/jsm/capabilities/WebGPU.js）不可用时直接抛出 No WebGPU support——BitonicSort 完全依赖 compute shader，纯 WebGL 环境无法运行；
• 单步可视化与 stepCount 配合：示例用 currentStep < bitonicSortModule.stepCount 作为推进条件，跑完 36 步后暂停 1 秒并把数据重置回随机初始态循环播放；
• 右侧对照实现：不使用 BitonicSort 类，而是自己用 Fn().compute( size / 2 )、compute( 1 ) 手工编排 FLIP_GLOBAL / DISPERSE_GLOBAL 的状态推进，展示了该算法不依赖封装类时的裸 TSL 写法。

七、使用限制与方案选择

从源码结构看，使用 BitonicSort 需要注意以下前提：

1. 元素个数应为 2 的幂。_getSwapOpCount 与 _getStepCount 直接使用 Math.log2( this.count )，双调排序本身也假定 N = 2^n；示例因此取 size = 16384。非 2 的幂会使步数计算退化为非整数、调度失去意义，这一点在文档中未显式声明，属于由源码结构可以推断的硬约束；
2. 需要 WebGPU compute 支持。排序全程由 renderer.compute( node ) 驱动；dataBuffer 必须是 StorageBufferNode，count 从 dataBuffer.value.count 读取，因此数据要先落到 GPU 存储缓冲里再参与排序；
3. 排序语义是无条件 min/max 交换。_globalCompareAndSwapTSL 与 _localCompareAndSwapTSL 一律写回 (min, max)，不区分升/降方向标志——该实现只产生升序结果，且对相等元素不做稳定性保证；
4. 规模选择：双调排序的比较交换轮数为 n(n+1)/2，适合 GPU 上数万到百万级的精确排序；同目录下的 examples/jsm/gpgpu/CountingSort.js 在注释中给出了互补建议——当只需要按键“近似分箱”（如按深度排序粒子）且规模达到数十万到数百万时，固定 4 趟（reset/histogram/prefix/scatter）的计数排序比精确比较排序更快，并可配合 computeCPU 在无 compute 支持的后端降级。

八、相关文件速览

文件	作用
examples/jsm/gpgpu/BitonicSort.js	BitonicSort 类、getBitonicFlipIndices、getBitonicDisperseIndices 的完整 TSL 实现
examples/webgpu_compute_sort_bitonic.html	官方分步可视化示例（local 与 global 双对照）
examples/jsm/gpgpu/CountingSort.js	同目录的 GPU 计数排序，适合大规模近似排序场景
docs/pages/BitonicSort_BitonicSort.html.md	本文对应的官方 API 文档页
examples/screenshots/webgpu_compute_sort_bitonic.jpg	示例运行截图

总结来看，BitonicSort 是 three.js WebGPU 管线中一个把“算法调度”与“GPU 并行”分层得很干净的设计：算法状态机放在 GPU 可写的 infoStorage 中、CPU 只按 stepCount 逐拍推进 dispatch；交换原语用 workgroup 局部内存吃下小跨度、用双缓冲乒乓吃下大跨度，最终以一次可选的 alignFn 对齐兑现原地排序语义。理解了 flip/disperse 索引映射与 local/global 两级调度之后，你就可以把它接入自己的 compute 工作流，或按同样的模式改造出其他需要多步 GPU 排序的场景。