three.js TSL 计算着色器中的 WorkgroupInfoElementNodeworkgroup 作用域缓冲区元素访问原理与实战【免费下载链接】three.jsJavaScript 3D Library.项目地址: https://gitcode.com/GitHub_Trending/th/three.jsthree.js 的 TSLThree Shading Language为 WebGPU 计算着色器提供了完整、类型安全的着色器节点体系而WorkgroupInfoElementNode正是其中负责对workgroup作用域缓冲区即工作组本地数组进行取下标、读元素、写元素的数组访问节点。本文以 WorkgroupInfoElementNode 文档 为主线结合其底层实现与webgpu_compute_reduce等官方示例讲解它的继承关系、构造函数、元素访问语义、赋值上下文行为以及它在 GPU 归约等并行算法中的真实用法帮助你写出可运行的高性能 WebGPU 计算着色器。什么是 workgroup 作用域缓冲区在 WebGPU / WGSL 的地址空间体系中数据可以存放在storage全局存储、uniform、private、function、workgroup等不同作用域中。其中workgroup作用域的缓冲区是归属于某个工作组workgroup的本地数组它由一个工作组内的所有调用invocation共享生命周期与工作组一致访问速度显著快于全局storage缓冲区因为它在物理上更接近处理单元典型用途是把全局存储中的数据先搬运进工作组本地数组再做多轮协作式处理如归约、排序、扫描期间配合workgroupBarrier()同步。在 three.js 的 TSL 中创建这种缓冲区的节点就是WorkgroupInfoNode而其每个元素的读写访问则统一由WorkgroupInfoElementNode表达。注意这类节点只能配合WebGPU 后端使用WebGL 后端无法表达 workgroup 地址空间。节点继承关系与类职责根据 WorkgroupInfoElementNode 文档 顶部的继承链EventDispatcher → Node → ArrayElementNode → WorkgroupInfoElementNode含义如下EventDispatcherthree.js 中一切具有事件能力的基类NodeTSL 着色器节点的通用基类承担类型推断、构建、序列化等职责ArrayElementNodeTSL 中表达对数组类节点数据结构的元素访问的通用基类定义于 src/nodes/utils/ArrayElementNode.js它持有被访问的数组节点node与下标节点indexNodeWorkgroupInfoElementNode继承ArrayElementNode把通用数组元素访问特化为 workgroup 作用域缓冲区的元素访问并处理赋值上下文assign context下的生成逻辑。WorkgroupInfoElementNode与WorkgroupInfoNode定义在同一个源文件中见 src/nodes/gpgpu/WorkgroupInfoNode.js。构造函数与参数说明new WorkgroupInfoElementNode( workgroupInfoNode, indexNode )参数定义如下参数类型说明workgroupInfoNodeNodeWorkgroupInfoNode被访问的 workgroup 作用域缓冲区节点即数组的宿主对象indexNodeNode通常是IndexNode定义元素访问的下标节点决定访问数组中的哪个元素构造函数内部把两个参数直接透传给父类ArrayElementNode并设置类型测试标志见 src/nodes/gpgpu/WorkgroupInfoNode.jsconstructor( workgroupInfoNode, indexNode ) { super( workgroupInfoNode, indexNode ); this.isWorkgroupInfoElementNode true; }在实际的 TSL 编程中你通常不会直接new这个类而是通过宿主节点WorkgroupInfoNode提供的element()方法获得实例。见 src/nodes/gpgpu/WorkgroupInfoNode.jselement( indexNode ) { return new WorkgroupInfoElementNode( this, indexNode ); }也就是说workgroupInfoNode.element( indexNode )返回一个WorkgroupInfoElementNode它在最终生成的 WGSL 中表现为对 workgroup 数组的下标访问表达式。类型测试标志isWorkgroupInfoElementNode文档中列出的唯一实例属性是只读的类型测试标志属性类型说明.isWorkgroupInfoElementNodeboolean只读用于类型测试的标记默认值为true它由构造函数中this.isWorkgroupInfoElementNode true;写入代码与 many other TSL 节点一致用于让开发者在拿到一个Node后快速判断它是否为 workgroup 元素访问节点从而安全地调用相关方法。注意它同时继承了父类的isArrayElementNode true见 ArrayElementNode因此该节点既可被识别为通用数组元素节点也可被识别为 workgroup 专用元素节点。赋值上下文generate() 的语义差异WorkgroupInfoElementNode覆写了父类的generate()以区分取值与赋值两种使用场景见 src/nodes/gpgpu/WorkgroupInfoNode.jsgenerate( builder, output ) { let snippet; const isAssignContext builder.isContextAssign(); snippet super.generate( builder ); if ( isAssignContext ! true ) { const type this.getNodeType( builder ); snippet builder.format( snippet, type, output ); } // TODO: Possibly activate clip distance index on index access rather than from clipping context return snippet; }这段代码揭示了关键语义先通过builder.isContextAssign()判断当前构建上下文是否是赋值上下文例如tSum.element(i).assign(0)时赋值的左侧非赋值上下文即作为右值读取元素时会对生成的代码片段按节点类型做format保证返回的 WGSL 片段类型正确赋值上下文中直接返回super.generate(builder)生成的裸左值片段array[ index ]以便后续拼接赋值避免多余的格式化。其父类ArrayElementNode.generate()生成的代码片段就是数组下标访问表达式见 src/nodes/utils/ArrayElementNode.jsgenerate( builder ) { const indexType this.indexNode.getNodeType( builder ); const nodeSnippet this.node.build( builder ); const indexSnippet this.indexNode.build( builder, ... ); return ${ nodeSnippet }[ ${ indexSnippet } ]; }因此一个WorkgroupInfoElementNode最终在 WGSL 中体现为形如workgroupArray_1[ localId ]的表达式并且该节点既能作为右值参与读取运算也能作为左值承载.assign()、.addAssign()等写操作。WGSL 生成链路getScopedArray当WorkgroupInfoNode.generate()被构建器调用时会请求后端注册这个作用域数组见 src/nodes/gpgpu/WorkgroupInfoNode.jsgenerate( builder ) { const name ( this.name ! ) ? this.name : ${this.scope}Array_${this.id}; return builder.getScopedArray( name, this.scope.toLowerCase(), this.bufferType, this.bufferCount ); }若用户通过setName()指定了名称则用该名称否则自动生成WorkgroupArray_id后端根据 scope 小写化后的值workgroup、元素类型与数量注册该数组。WebGPU 后端的WGSLNodeBuilder实现了getScopedArray()并只允许在 compute 阶段输出这些作用域数组见 src/renderers/webgpu/nodes/WGSLNodeBuilder.jsgetScopedArray( name, scope, bufferType, bufferCount ) { if ( this.scopedArrays.has( name ) false ) { this.scopedArrays.set( name, { name, scope, bufferType, bufferCount } ); } return name; } getScopedArrays( shaderStage ) { if ( shaderStage ! compute ) { return; // 顶点/片元阶段不支持 workgroup 数组 } // ... 生成 WGSL 作用域数组声明 }这也从后端层面印证了文档仅 WebGPU 计算着色器可用的约束作用域数组的 WGSL 声明只在 compute stage 生成顶点或片元阶段调用getScopedArrays()会直接返回undefined。实战用法从创建到读写元素下面以官方示例 examples/webgpu_compute_reduce.html 中的归约核函数为模板演示WorkgroupInfoNodeworkgroupArray与WorkgroupInfoElementNodeelement()的完整配合。首先创建 workgroup 作用域数组其中workgroupSize为工作组线程数import { Fn, uint, instanceIndex, invocationLocalIndex, Loop, workgroupArray, workgroupBarrier, storage } from three/tsl; // 声明一个类型为 uint、元素个数为 workgroupSize 的 workgroup 缓冲区 const tSum workgroupArray( uint, workgroupSize );注意workgroupArray( type, count )只是WorkgroupInfoNode的 TSL 工厂函数定义于 src/nodes/gpgpu/WorkgroupInfoNode.jsexport const workgroupArray ( type, count ) new WorkgroupInfoNode( Workgroup, type, count );然后通过.element( indexNode )获得WorkgroupInfoElementNode既可用于赋值左值也可用于读取与累加右值const fnDef Fn( () { // 每个线程初始化自己负责的那个桶 tSum.element( invocationLocalIndex ).assign( uint( 0 ) ); const k instanceIndex.toVar( k ); // 跨步累加每个线程把全局 storage 中的多个元素累加进本地桶 Loop( k.lessThan( uint( numElements ) ), () { tSum.element( invocationLocalIndex ).addAssign( inputBuffer.element( k ) ); k.addAssign( uint( dispatchSize ) ); } ); // 同步工作组内所有线程确保桶数据全部写入完成 workgroupBarrier(); // 树形归约读相邻桶并累加 k.assign( pow2Ceil( uint( workgroupSize ) ).div( 2 ) ); Loop( k.greaterThan( 0 ), () { If( invocationLocalIndex.lessThan( k ).and( invocationLocalIndex.add( k ).lessThan( workgroupSize ) ), () { // 右侧读 tSum.element(...)左侧 addAssign 也是元素节点写操作 tSum.element( invocationLocalIndex ).addAssign( tSum.element( invocationLocalIndex.add( k ) ) ); } ); workgroupBarrier(); k.divAssign( 2 ); } ); // 由线程 0 把结果写回全局 storage If( invocationLocalIndex.equal( uint( 0 ) ), () { inputBuffer.element( workgroupId.x ).assign( tSum.element( uint( 0 ) ) ); } ); } )().compute( dispatchSize, [ workgroupSize ] );上述代码的要点tSum.element( i ).assign( v )、.addAssign( v )走的是赋值上下文节点作为左值生成workgroupArray[ i ] .../workgroupArray[ i ] ...tSum.element( i )作为右值参与算术时读取对应桶的数值下标节点可以是invocationLocalIndex每个线程的本地索引、invocationLocalIndex.add( k )相对偏移或常量uint( 0 )等任意返回整数类型的 TSL 表达式多次元素访问配合workgroupBarrier()即经典的工作组内归约模式。更多应用BitonicSort 与 blockReduceWorkgroupInfoElementNode不只出现在归约示例中双调排序bitonic sort也是它的重要应用场景。在 examples/jsm/gpgpu/BitonicSort.js 中排序器为每个比较阶段维护一个本地排序缓冲区this.localStorage workgroupArray( dataBuffer.nodeType, this.workgroupSize * 2 );随后通过this.localStorage.element( localId )完成对元素数组的读写与交换见同文件中基于element()的读写逻辑。这说明只要核函数内需要工作组内共享的可随机读写数组workgroupArray与WorkgroupInfoElementNode就是标准组合。同样地在 examples/webgpu_compute_reduce.html 中还利用workgroupArray( uint, workgroupSize / minSubgroupSize )实现按子组subgroup维度划分的多级归约进一步体现了该节点在复杂并行算法中作为高性能本地暂存区的普适性。关键实现速查关注点内容源码位置类定义WorkgroupInfoElementNode继承ArrayElementNodesrc/nodes/gpgpu/WorkgroupInfoNode.js构造参数workgroupInfoNode、indexNode透传父类src/nodes/gpgpu/WorkgroupInfoNode.js类型标志isWorkgroupInfoElementNode truesrc/nodes/gpgpu/WorkgroupInfoNode.js工厂创建element()返回元素节点src/nodes/gpgpu/WorkgroupInfoNode.js工厂函数workgroupArray(type, count)src/nodes/gpgpu/WorkgroupInfoNode.js基类行为数组下标代码生成src/nodes/utils/ArrayElementNode.jsWGSL 后端compute 阶段注册并声明作用域数组src/renderers/webgpu/nodes/WGSLNodeBuilder.js使用示例工作组内归约examples/webgpu_compute_reduce.html使用示例GPU 双调排序examples/jsm/gpgpu/BitonicSort.js小结WorkgroupInfoElementNode虽然本身只是一个薄薄的数组元素访问节点但它是 TSL 计算着色器中访问工作组共享内存的标准入口。理解它需要把握三个层次结构上它继承ArrayElementNode把通用数组下标访问限定到 workgroup 作用域由WorkgroupInfoNode.element(indexNode)工厂创建语义上generate()根据builder.isContextAssign()区分读写右值读取会做类型格式化赋值上下文则直接生成数组下标左值后端上它依赖 WebGPU 后端的getScopedArray()在 compute 阶段声明workgroup数组因此仅能在 WebGPU 后端下工作。掌握了这三个层次你就可以在归约、扫描、排序等需要线程协作的 GPU 算法中自如地使用workgroupArray(...).element(i)充分发挥 workgroup 本地存储的带宽优势。【免费下载链接】three.jsJavaScript 3D Library.项目地址: https://gitcode.com/GitHub_Trending/th/three.js创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考