跳转至

排序模块

排序模块托管了 GPU 基数排序器，用于在绘制投影后的 Splat 之前按深度对其进行排序。它拥有将预处理的未排序输出转换为连续的负载索引列表所需的缓冲、绑定组和计算管线，渲染器可以将该列表提供给间接绘制调用。所有内容都位于 src/sort/ 下，并在预处理器和渲染器之间共享。

概览

项目 路径 目的
接口 src/sort/index.ts 渲染器堆栈中使用的 ISorter, SortedSplats, PointCloudSortStuff 契约。
实现 src/sort/radix_sort.ts GPURSSorter 类加上缓冲分配辅助方法。
着色器 src/shaders/radix_sort.wgsl 四个计算入口点 (zero_histograms, calculate_histogram, prefix_histogram, scatter_even / scatter_odd)。

职责

  • 分配和回收基数排序 Pass 中使用的乒乓键和负载缓冲。
  • 暴露绑定组布局，以便预处理 (group 2) 和渲染 (group 1) 可以绑定排序器资源而无需重复布局。
  • 记录固定大小排序 (recordSort) 或由预处理填充的计数器驱动的间接排序。
  • 维护 GeneralInfo 统一缓冲 (keys_size, padded_size, pass 标志) 以及在预处理、排序和渲染之间共享的间接调度缓冲。
  • 提供辅助 API (createSortStuff, recordResetIndirectBuffer)，以便渲染器可以保持每个模型的缓存与可见 Splat 数量同步。

数据流

预处理器 (深度键)
        |
        | 写入 key_a, payload_a, keys_size, dispatch_x
        v
PointCloudSortStuff (key_a, payload_a, sorter_uni, sorter_dis)
        |
        | 被 GPURSSorter Pass 消费
        v
Zero -> Histogram -> Prefix -> Scatter
        |
        v
渲染器 (绑定 sorter_render_bg, 发出间接绘制)

缓冲与 Uniform 快照

缓冲 来源 用法 备注
key_a / key_b Sorter STORAGE 乒乓深度键 (u32)。
payload_a / payload_b Sorter STORAGE 乒乓负载索引；payload_a 兼作渲染器排序索引缓冲。
internal_mem Sorter STORAGE 用于 WGSL 的直方图、分区和回溯元数据。
sorter_uni Sorter STORAGE GeneralInfo 结构体 (keys_size, padded_size, pass 选择器)。由渲染器和预处理更新。
sorter_dis Sorter STORAGE / INDIRECT / COPY_DST 存储用于间接直方图和 Scatter Pass 以及绘制调用的 dispatch_x/y/z。
sorter_bg_pre Sorter Bind group 向预处理暴露 sorter_uni, key_a, payload_a, sorter_dis (@group(2))。
sorter_render_bg Sorter Bind group 向渲染器暴露 sorter_uni 和 payload_a (@group(1))。

PointCloudSortStuff 捆绑了上述资源，以便渲染器可以为每个点云或每个全局缓冲保留一个缓存。

GPU 管线 (4 个阶段)

  1. Zero pass: 清除直方图和共享内存。
  2. Histogram pass: 每个工作组（256 线程 x 15 行）扫描 3840 个键，提取所有四个基数 Pass 的 8 位数字，并原子地累积计数。
  3. Prefix pass: 128 线程的工作组对 256 个条目的直方图执行独占扫描 (exclusive scan) 以产生全局偏移。
  4. Scatter passes: 两个入口点 (scatter_even, scatter_odd) 处理 Pass {0,2} 和 {1,3}，使用乒乓缓冲以避免冒险 (hazards)。每个线程计算局部排名，加上全局偏移，并写入新的键和负载顺序。

关键常量： - 8 位基数 -> 32 位键需要 4 个 Pass。 - 工作组大小：histogram 和 scatter = 256, prefix = 128。 - RS_HISTOGRAM_BLOCK_ROWS = 15 因此每个工作组处理 3840 个键。 - 所有派生常量都在 processShaderTemplate 中烘焙到 WGSL 中，以保持填充和共享内存大小与 CPU 辅助方法同步。

集成点

  • 预处理 将 sorter_bg_pre 绑定为 @group(2) 以写入深度键、负载索引、keys_size 和间接调度计数。键被填充到下一个 256 * 15 的倍数，以便直方图 Pass 可以无边界检查地迭代。
  • 渲染器 为每个点云缓存一个 PointCloudSortStuff。prepareMulti 重置间接缓冲 (recordResetIndirectBuffer)，计算基准偏移量，运行预处理，然后使用预处理生成的调度数据调用一次 recordSortIndirect。渲染 Pass 绑定 sorter_render_bg 以读取排序后的索引和实例计数。
  • 动态点数: 当 ONNX 管线改变可见 Splat 的数量时，预处理覆盖 GeneralInfo.keys_size 和 sorter_dis.dispatch_x。排序器间接路径直接使用这些值，因此不需要 CPU 回读。
  • 测试: GPURSSorter.create 尝试几种子组大小 (16 -> 32 -> 16 -> 8 -> 1) 并运行 testSort（排序 8192 个浮点数），以确保编译的管线在向渲染器公开排序器之前在当前适配器上工作。

使用示例

import { GPURSSorter } from 'src/sort';

const sorter = await GPURSSorter.create(device, device.queue);
const sortStuff = sorter.createSortStuff(device, pointCloud.numPoints);
sortResourcesCache.set(pointCloud, sortStuff);

// 每帧预处理前
sorter.recordResetIndirectBuffer(sortStuff.sorter_dis, sortStuff.sorter_uni, device.queue);

// ... 预处理将深度键和负载写入 sortStuff ...

const encoder = device.createCommandEncoder();
sorter.recordSortIndirect(sortStuff, sortStuff.sorter_dis, encoder);
device.queue.submit([encoder.finish()]);

// 渲染 Pass
pass.setBindGroup(1, sortStuff.sorter_render_bg);
pass.drawIndirect(drawIndirectBuffer, 0);

相关文档