three.js TSL AtomicFunctionNode：在 WebGPU Compute Shader 中实现 GPU 原子操作

AtomicFunctionNode 是 three.js 节点系统（Nodes/TSL）中专用于 WebGPU 后端的节点类型，它把 WebGPU 着色语言（WGSL）中的 atomic 内置函数封装成声明式的 TSL 调用，使并行执行的 GPU 线程对同一个原子变量的修改成为不可分割、有序的操作。读完本文，你能掌握原子函数节点的构造参数、九个原子操作 API（atomicLoad/atomicStore/atomicAdd 等）的用法，并结合仓库内 CountingSort 等实例理解其源码级的代码生成机制。

一、原子操作节点解决什么问题

在 GPU 上，一个 compute shader 会同时调度成千上万个线程。如果多个线程对同一块存储缓冲区的同一个变量做“读-改-写”（例如计数自增），普通读写会相互覆盖、产生竞态（race condition）。原子操作（atomic operation）正是为此设计：

AtomicFunctionNode 表示着色器中任何可作用于原子变量类型的函数。在原子函数中，对原子变量的任何修改都会作为一个不可分割的步骤发生，并且相对于其他修改具有确定的顺序。因此，即使多个原子函数同时在修改同一个原子变量，这些原子操作也互不干扰。

——src/nodes/gpgpu/AtomicFunctionNode.js 的文档注释

其继承链为 EventDispatcher → Node → AtomicFunctionNode，即直接扩展自节点基类 Node。

适用前提：该节点只能配合 WebGPU 后端使用（WebGPURenderer 处于 WebGPU 模式时）。它依赖 WGSL 的 atomic 内置函数，因此要求运行环境支持 compute shader 与原子存储操作；在 WebGL 回退模式下不可用。仓库中 CountingSort 也明确为不支持 compute shader 的平台准备了 computeCPU 纯 JS 回退路径，这正是“仅 WebGPU”限制的典型应对方式。

二、构造函数与属性

new AtomicFunctionNode( method : string, pointerNode : Node, valueNode : Node )

构造一个新的原子函数节点，三个参数含义如下：

参数	类型	说明
method	string	要构造的原子函数签名，即 WGSL atomic 函数名（如 'atomicAdd'）
pointerNode	Node	原子变量，或原子缓冲区中的某个元素（通常来自 storage(...).toAtomic() 的结果）
valueNode	Node	用于修改该原子变量的值；atomicLoad 时为 null

从 构造函数实现 可以看到两点关键设计：

• super( 'uint' ) —— 节点基类型默认为 uint，因为 WGSL 原子操作只允许无符号整型，这是原子函数的固有约束；
• this.parents = true —— 覆盖父类 Node#parents 的默认行为，主动构建父节点列表，用于在 generate() 中判断“当前节点是否需要返回值”（详见第四节）。

类中暴露的属性：

• .method : string —— 原子函数签名；
• .pointerNode : Node —— 原子变量或原子缓冲元素；
• .valueNode : Node —— 修改原子变量的值；
• .parents : boolean —— 默认 true，为父节点建立列表以检测节点是否需要返回值。

九个原子函数族

源码通过静态 getter 定义了九个原子操作签名常量（L136-L188），并在文件底部导出对应的 TSL 便捷函数（L216-L318）：

TSL 函数	method 签名	语义
atomicLoad( pointerNode )	atomicLoad	原子加载：读取原子变量当前值
atomicStore( pointerNode, valueNode )	atomicStore	原子存储：写入新值
atomicAdd( pointerNode, valueNode )	atomicAdd	原子加：自增/累加
atomicSub( pointerNode, valueNode )	atomicSub	原子减
atomicMax( pointerNode, valueNode )	atomicMax	存储当前值与参数中的较大者
atomicMin( pointerNode, valueNode )	atomicMin	存储当前值与参数中的较小者
atomicAnd( pointerNode, valueNode )	atomicAnd	按位与
atomicOr( pointerNode, valueNode )	atomicOr	按位或
atomicXor( pointerNode, valueNode )	atomicXor	按位异或

除 atomicLoad 外，其余函数均接收 valueNode；注意各函数内部统一通过 atomicFunc( method, pointerNode, valueNode ) 创建节点并调用 .toStack()（L216-L220），即把节点追加进程序化执行流（stack），使原子调用成为 compute 函数体中的一条顺序语句，而不是一个可参与表达式求值的临时值。这与 docs/TSL.md “Compute” 一节中列出的 TSL API 表格完全对应。

三、代码生成机制：generate() 如何产出 WGSL

generate( builder )（L91-L134）是理解该节点行为的核心，它按以下顺序生成着色器代码：

1. 取类型：type = this.getNodeType( builder )、inputType = this.getInputType( builder )；
2. 构造实参列表：第一个参数是指针引用 &${ a.build( builder, inputType ) }——对 pointerNode 的构建结果加上 &，即把存储缓冲元素以引用（reference）形式传入 atomic 函数；若 valueNode !== null 再追加第二个参数（atomicLoad 只有一个参数）；
3. 拼接调用：`${ builder.getMethod( method, type ) }( ${ params.join( ', ' ) } )`，由 NodeBuilder 根据后端把 method 映射为对应的原子函数名，最终形如 atomicAdd( &hist[i], u32(1) )；
4. void 还是表达式：
   • 若父节点列表存在且唯一父节点 isStackNode === true（即该调用直接挂在 Fn 的语句流下、没有人消费它的返回值），则通过 builder.addLineFlowCode( methodSnippet, this ) 生成一条独立语句——原子操作只产生副作用，不产生值；
   • 否则，缓存/创建一个 expression( methodSnippet, type ).toConst() 并返回其构建结果，把原子函数返回的旧值作为 const 表达式供后续节点引用。

第 4 步正是 .parents = true 存在的意义：原子操作既是语句又是可取值表达式（WGSL 的 atomic 交换类函数返回旧值），节点通过父节点结构动态决定走哪条路径。

类型推断的两个覆盖方法

文档中列出的两个方法覆盖（源码 L66-L89）：

• .getInputType( builder ) : string —— 覆盖默认实现，直接返回 this.pointerNode.getNodeType( builder )，即输入类型 = 指针节点所指向元素的类型（例如 u32）；
• .generateNodeType( builder ) : string —— 因节点类型由输入类型推断，直接返回 this.getInputType( builder )。

这意味着整个原子节点的类型链完全由 pointerNode 锚定：存储缓冲声明为 'uint' 时，指针元素、valueNode 的构建类型、生成的 WGSL 表达式类型三者保持一致。

四、实战：与 storage().toAtomic() 配合的完整用法

AtomicFunctionNode 的 pointerNode 通常来自 StorageBufferNode 的 toAtomic()（L292-L296，底层置 isAtomic 标志）。仓库中最完整的范例是 CountingSort（GPU 计数排序），其四个 compute pass 几乎覆盖了全部原子 API：

// 1. 原子缓冲的创建：histogram / offset 两个 uint 缓冲声明为 atomic
this.histogramAtomic = storage( histogramAttribute, 'uint', binCount ).toAtomic();
this.offsetAtomic    = storage( offsetAttribute,    'uint', binCount ).toAtomic();

// 2. Reset pass —— atomicStore 清零
Fn( () => {
    atomicStore( this.histogramAtomic.element( instanceIndex ), uint( 0 ) );
    atomicStore( this.offsetAtomic.element( instanceIndex ), uint( 0 ) );
} )().compute( binCount, [ workgroupSize ] ).setName( 'CountingSortReset' );

// 3. Histogram pass —— atomicAdd 对直方图桶计数
Fn( () => {
    const bin = binNode().toVar( 'bin' );
    this.binWrite.element( instanceIndex ).assign( bin );
    atomicAdd( this.histogramAtomic.element( bin ), uint( 1 ) );
} )().compute( count, [ workgroupSize ] ).setName( 'CountingSortHistogram' );

// 4. Prefix pass —— atomicLoad + atomicStore 计算前缀和
const binCountValue = atomicLoad( this.histogramAtomic.element( bin ) ).toVar( 'count' );
atomicStore( this.offsetAtomic.element( bin ), sum );

// 5. Scatter pass —— 利用 atomicAdd 返回的旧值作为写游标
const targetIndex = atomicAdd( this.offsetAtomic.element( bin ), uint( 1 ) ).toVar( 'targetIndex' );
this.orderWrite.element( targetIndex ).assign( instanceIndex );

（代码摘自 examples/jsm/gpgpu/CountingSort.js#L145-L188）

这段示例把 generate() 的两条路径都演示了一遍：histogram 与 reset pass 中的 atomicAdd/atomicStore 直接位于 Fn 语句流下，走 void 语句路径；而 scatter pass 中的 atomicAdd(...) 旧值被 .toVar('targetIndex') 消费，走 const 表达式路径——原子返回的旧值恰好就是每个 bin 的下一个空闲槽位，这是原子操作“既改值又取值”语义在 GPU 并行写入中的经典应用。

五、仓库中其他真实使用场景

搜索整个仓库，AtomicFunctionNode 通过 src/nodes/TSL.js 与 src/nodes/Nodes.js 对外导出（分别供 TSL 函数与节点类两种用法），在示例中的代表用例包括：

• examples/webgpu_compute_rasterizer.html —— GPU 光栅化器：用 atomicStore 重置屏幕三角形/实例计数缓冲，用 atomicAdd 给 work queue 分配任务索引，把多 workgroup 的并行任务分发建立在原子计数之上；
• examples/webgpu_compute_particles_fluid.html —— 流体粒子：粒子按空间格点散列，各线程用 atomicAdd 把定点编码后的动量/质量累加进所属格点（cell.get('x') 等字段），实现无锁的并行 reduce；
• examples/webgpu_struct_drawindirect.html —— 用 atomicStore 写入 draw indirect 结构体中的 instanceCount 字段，再复位归零；
• examples/jsm/gpgpu/CountingSort.js —— 前文详述的计数排序，atomicAdd 在 scatter pass 中直接承担“并发分配目标下标”的角色。

可以看出这些场景有一个共同模式：多个 workgroup 对共享的 uint 计数/游标/直方图做并发修改，而这正是普通存储缓冲写入无法安全完成的。

六、使用约束与注意事项

综合文档与源码，使用该节点时需注意以下前提：

1. 仅 WebGPU 后端：节点依赖 WGSL atomic 内置函数，在 WebGL 模式下不可用；需要 CPU 回退时参考 CountingSort#computeCPU 的做法；
2. 整型约束：构造函数以 'uint' 为基类型（L33），指针节点应指向无符号整型元素，valueNode 建议同样显式包裹为 uint(...)；
3. valueNode 可为 null：仅 atomicLoad 无第二个参数，生成代码中相应只传 &pointer；
4. 两种消费语义：在语句流中直接使用即产生副作用（void 路径）；将其返回值赋给变量（如 .toVar()）时会自动生成 const 表达式拿到原子函数返回的旧值（const 路径），两者由节点父级结构自动判定；
5. 缓冲需先声明原子：pointerNode 必须来自标记过 toAtomic() 的存储节点，否则底层存储不是 atomic 存储，无法通过原子函数引用。

七、小结

AtomicFunctionNode 以“method 签名 + 指针节点 + 值节点”三要素把 WGSL 的九个原子函数统一封装进 three.js 的 TSL 体系：getInputType/generateNodeType 保证类型沿指针节点推断，generate() 通过 & 引用参数与 void/const 双路径完成 WGSL 代码生成，parents = true 则让同一个节点既能当纯副作用语句、又能把旧值继续流入表达式。配合 storage(...).toAtomic() 与 Fn(...).compute(...)，它可以覆盖从 GPU 排序、并行累加到任务队列分配的常见 compute shader 并发需求。

相关源码入口：src/nodes/gpgpu/AtomicFunctionNode.js、src/nodes/accessors/StorageBufferNode.js、API 参考 docs/pages/AtomicFunctionNode.html.md 与 docs/TSL.md。