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

three.js TSL AtomicFunctionNode在 WebGPU Compute Shader 中实现 GPU 原子操作【免费下载链接】three.jsJavaScript 3D Library.项目地址: https://gitcode.com/GitHub_Trending/th/three.jsAtomicFunctionNode是 three.js 节点系统Nodes/TSL中专用于 WebGPU 后端的节点类型它把 WebGPU 着色语言WGSL中的 atomic 内置函数封装成声明式的 TSL 调用使并行执行的 GPU 线程对同一个原子变量的修改成为不可分割、有序的操作。读完本文你能掌握原子函数节点的构造参数、九个原子操作 APIatomicLoad/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 )构造一个新的原子函数节点三个参数含义如下参数类型说明methodstring要构造的原子函数签名即 WGSL atomic 函数名如atomicAddpointerNodeNode原子变量或原子缓冲区中的某个元素通常来自storage(...).toAtomic()的结果valueNodeNode用于修改该原子变量的值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-L318TSL 函数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() 如何产出 WGSLgenerate( builder )L91-L134是理解该节点行为的核心它按以下顺序生成着色器代码取类型type this.getNodeType( builder )、inputType this.getInputType( builder )构造实参列表第一个参数是指针引用${ a.build( builder, inputType ) }——对pointerNode的构建结果加上即把存储缓冲元素以引用reference形式传入 atomic 函数若valueNode ! null再追加第二个参数atomicLoad只有一个参数拼接调用${ builder.getMethod( method, type ) }( ${ params.join( , ) } )由 NodeBuilder 根据后端把 method 映射为对应的原子函数名最终形如atomicAdd( hist[i], u32(1) )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标志。仓库中最完整的范例是 CountingSortGPU 计数排序其四个 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)等字段实现无锁的并行 reduceexamples/webgpu_struct_drawindirect.html —— 用atomicStore写入 draw indirect 结构体中的instanceCount字段再复位归零examples/jsm/gpgpu/CountingSort.js —— 前文详述的计数排序atomicAdd在 scatter pass 中直接承担“并发分配目标下标”的角色。可以看出这些场景有一个共同模式多个 workgroup 对共享的 uint 计数/游标/直方图做并发修改而这正是普通存储缓冲写入无法安全完成的。六、使用约束与注意事项综合文档与源码使用该节点时需注意以下前提仅 WebGPU 后端节点依赖 WGSL atomic 内置函数在 WebGL 模式下不可用需要 CPU 回退时参考 CountingSort#computeCPU 的做法整型约束构造函数以uint为基类型L33指针节点应指向无符号整型元素valueNode建议同样显式包裹为uint(...)valueNode 可为 null仅atomicLoad无第二个参数生成代码中相应只传pointer两种消费语义在语句流中直接使用即产生副作用void 路径将其返回值赋给变量如.toVar()时会自动生成const表达式拿到原子函数返回的旧值const 路径两者由节点父级结构自动判定缓冲需先声明原子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。【免费下载链接】three.jsJavaScript 3D Library.项目地址: https://gitcode.com/GitHub_Trending/th/three.js创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考