跳转至

Preprocessing 模块架构

本文档解释了 GaussianPreprocessor 如何将 WebGPU 资源拼接在一起,打包 Uniform,并驱动 preprocess.wgsl,以便每帧可以将多个点云投影到单个全局 2D Splat 缓冲区中。它重点关注 GaussianRenderer.prepareMulti 今天执行的 B 阶段(多模型)路径。

管道概览

┌──────────────┐  ┌──────────────────────────────┐  ┌────────────────┐
│ PointCloud   │──▶│  GaussianPreprocessor        │──▶│ GPURSSorter    │
│ GPU 缓冲区    │  │  dispatchModel(...)          │  │ (基数排序)       │
└─────┬────────┘  │ ├ 相机 + 设置 Uniform          │  └────────┬───────┘
      │           │ ├ 投影、剔除、评估 SH           │           │
      │           │ └ 写入全局 Splat 切片          │           │
      │           └──────────────┬───────────────┘     Renderer
      │                          │
      │           全局 Splat 缓冲区 + 排序计数器
      │           (间接绘制)

特点:

  • 每个点云一次计算分发 (compute dispatch),每个工作组 256 个线程,每个线程处理一个高斯。
  • 所有模型共享相同的 splat2D 缓冲区和排序资源;每次分发写入 [baseOffset, baseOffset + count) 区间。
  • 运行时存在两个管道:SH(视图相关)和原始 RGB(用于 DynamicPointCloud)。着色器路径通过 USE_RAW_COLOR 管道常量选择。
  • 排序计数器(keys_size,间接 dispatch_x)在着色器内原子更新,因此基数排序可以通过间接分发运行,无需 CPU 干预。

绑定组与管道布局

initialize() 构建具有四个组的管道布局,以便我们在保持 WebGPU 绑定组限制的同时支持多模型输出:

@group(0) 相机 Uniform (272 B)
  view, view⁻¹, proj (Y翻转), proj⁻¹, viewport, 焦距

@group(1) 点云数据
  binding 0: 高斯缓冲区 (只读存储)
  binding 1: SH / 原始颜色缓冲区 (只读存储)
  binding 2: splat2D 输出 (存储) —— 绑定到全局缓冲区
  binding 3: 点云绘制 Uniform (Uniform 缓冲区)

@group(2) 排序缓冲区 (来自 GPURSSorter.createPreprocessBindGroupLayout)
  binding 0: 排序信息和原子计数器 (存储)
  binding 1: 深度键 (存储)
  binding 2: 负载索引 (存储)
  binding 3: 间接分发缓冲区 (存储)

@group(3) 设置 + 模型参数
  binding 0: 渲染设置 Uniform (80 B)
  binding 1: 模型参数 Uniform (128 B, 由 PointCloud 拥有)

计算管道由 preprocess.wgsl 编译,SH 阶数注入 WGSL 源代码,USE_RAW_COLOR 通过管道常量提供。

分发流程 (dispatchModel)

对于帧中的每个点云,渲染器调用 GaussianPreprocessor.dispatchModel({...}, encoder)。该方法在记录计算通道之前执行以下步骤:

  1. 打包相机数据 —— packCameraUniforms 用视图/投影矩阵(包括 WebGPU Y 轴翻转)、其逆矩阵、视口和从 PerspectiveCamera.projection.focal() 派生的焦距填充 272 字节的暂存缓冲区。缓冲区通过 UniformBuffer.setData + flush 刷新。
  2. 打包渲染设置 —— packSettingsUniforms 将剔除框、gaussianScalingmaxSHDegree、环境/Mip 标志、kernelSizewalltimesceneExtend 和场景中心写入 80 字节的渲染设置 Uniform,然后刷新它。
  3. 更新每个模型的参数 —— pointCloud.updateModelParamsWithOffset(modelMatrix, baseOffset) 将变换和切片偏移存储在点云自己的 128 字节 Uniform 中。如果点云暴露 setPrecisionForShader(动态 ONNX 路径),则调用它,以便量化元数据(数据类型、缩放、零点)在缓冲区刷新之前落入第 96-119 字节。
  4. 处理动态计数 —— 当提供 countBuffer 时,预处理器首先刷新 pointCloud.modelParamsUniforms,然后将 countBuffer 中的四个字节复制到模型参数缓冲区的字节偏移量 68 处(num_points 字段)。这使得 ONNX 生成器可以驱动间接绘制,无需 CPU 往返。
  5. 绑定资源 —— 使用新更新的缓冲区组装第 0-3 组。第 1 组的绑定 2 指向全局 splat2D 缓冲区,而不是点云本地缓冲区。
  6. 分发 —— 工作组数 = ceil(pointCloud.numPoints / 256)。计算通道将 Splat 写入基偏移区域,更新排序键/负载,并原子递增排序器计数器。

处理完所有模型后,渲染器触发单个 GPURSSorter.recordSortIndirect(...),随后是一个间接绘制。间接绘制的 instanceCount 通过将 sorter_uni.keys_size 复制到绘制缓冲区来填充。

Uniform 布局

相机 Uniform (272 字节)

0-63    : 视图矩阵 (mat4x4<f32>)
64-127  : 视图逆矩阵
128-191 : 投影矩阵 (带 Y 翻转)
192-255 : 投影逆矩阵
256-263 : 视口 (宽, 高)
264-271 : 焦距 (fx, fy)

渲染设置 Uniform (80 字节)

0-15    : 剔除框 min (vec4)
16-31   : 剔除框 max (vec4)
32      : gaussianScaling (f32)
36      : maxSHDegree (u32)
40      : showEnvMap 标志 (u32)
44      : mipSplatting 标志 (u32)
48      : kernelSize (f32)
52      : walltime (f32)
56      : sceneExtend (f32)
60      : 填充
64-79   : 场景中心 (vec3 + 填充)

模型参数 Uniform (128 字节)

PointCloud 内部管理。它存储模型矩阵、baseOffset (u32@64)、num_points (u32@68)、高斯缩放、最大 SH 阶数、内核/不透明度/截止缩放、渲染模式和精度元数据(数据类型 + 缩放/零点)。预处理依赖于此缓冲区在分发前已更新。

着色器内部 (preprocess.wgsl)

工作组拓扑

@compute @workgroup_size(256, 1, 1)
fn preprocess(@builtin(global_invocation_id) gid: vec3<u32>) {
  let idx = gid.x;
  if (idx >= arrayLength(&gaussians)) {
    return;
  }
  // ... 处理一个高斯
}

每个工作组 256 个线程在 warp 占用率、寄存器压力和内存合并之间取得了平衡,同时匹配排序器的 keys_per_workgroup = 256 * 15 常量,用于间接分发大小调整。

缓冲区布局

struct Gaussian {
  pos_opacity: array<u32, 2>; // 打包的 f16 xyz + 不透明度
  cov        : array<u32, 3>; // 打包的 f16 协方差 (6 个值)
}

struct Splat {
  v0: u32; v1: u32;          // 打包为 f16 的主/副轴
  pos: u32;                  // 屏幕位置 (f16)
  color0: u32; color1: u32;  // 打包的 RGBA (f16)
}

SH/RAW 缓冲区是一个 array<array<u32, 24>>,每个高斯提供 48 个 f16 值(16 个 SH 系数 × 3 个通道),着色器按需解包。

投影与协方差

  1. 解包 对称的 3×3 协方差矩阵(从三个 u32 值,即六个 f16 中),并按高斯缩放因子的平方进行缩放。
  2. 计算 透视投影的雅可比矩阵,使用焦距和相机空间深度。
  3. 提取 世界到相机的旋转(视图矩阵左上角 3×3 块的转置)。
  4. 通过 Σ₂D = (W·J)ᵀ · Σ₃D · (W·J) 投影 以生成屏幕空间中的 2×2 协方差。
let Vrk = mat3x3<f32>( /* 解包 */ ) * scaling * scaling;
let J = mat3x3<f32>( /* 焦距, 深度项 */ );
let W = transpose(mat3x3<f32>(camera.view[0].xyz, camera.view[1].xyz, camera.view[2].xyz));
let T = W * J;
let cov2d = transpose(T) * Vrk * T;

特征值/向量通过解析方法推导,kernel_size 添加到对角线上以进行抗锯齿,最小特征值被限制为 ≥ 0.1 以保持椭圆表现良好。每个特征向量按 sqrt(2λ) 缩放以产生渲染器预期的半轴。

颜色评估

  • SH 模式 —— evaluate_sh(dir, sh_deg) 加速高达 3 阶的预计算基多项式。当 settings.maxSHDegree 较低时,阶数阈值让着色器跳过高阶项,减少 ALU 开销。
  • 原始 RGB 模式 —— 当使用 useRawColor=true 创建管道时,特化常量绕过 SH 数学运算,并将 SH 缓冲区视为直接 RGBA(仍存储为打包的 f16 对)。

可见性与剔除

  1. 剔除框:拒绝世界空间中 settings.clipping_box_min/max 之外的 Splat。
  2. 视锥体测试:投影后,确保 0 < z < 1-1.2w < x,y < 1.2w 以保留小的安全余量。
  3. Mip-splatting (可选):如果启用,比较添加 kernel_size 前后的行列式以调制不透明度并减少闪烁。

所有测试通过 return 短路,以最大限度地减少 warp 内的分支差异。

与排序器的原子握手

let output_idx = atomicAdd(&sort_infos.keys_size, 1u);
points_2d[output_idx] = packed_splat;
sort_keys[output_idx] = bitcast<u32>(zfar - pos_ndc.z);
sort_payload[output_idx] = output_idx;

let KEYS_PER_WORKGROUP = 256u * 15u;
if (output_idx % KEYS_PER_WORKGROUP) == 0u {
  atomicAdd(&sort_dispatch.dispatch_x, 1u);
}

这些原子操作使 sort_infos.keys_size 与可见 Splat 的数量保持同步,并在每次填满新的键块时递增间接分发计数器。渲染器稍后将 sort_infos.keys_size 复制到绘制间接缓冲区中以设置最终的 instanceCount

多模型 / 全局路径

B 阶段在不改变着色器契约的情况下整合了跨模型的资源使用:

  • 渲染器分配一个大的 splat2D 缓冲区和一个 PointCloudSortStuff 结构,其容量与所有点云计数之和匹配。
  • 对于每个模型:计算 baseOffset,绑定共享缓冲区,并调用 dispatchModel(根据 pointCloud.colorMode 使用 SH 或 RGB 预处理器)。
  • 当所有模型处理完毕后,调用一次 recordSortIndirect 并使用全局排序负载发出单个间接绘制。

这种方法消除了冗余的排序/绘制通道,保持绑定组布局稳定,并为未来的批处理或遮挡感知调度奠定了基础。

性能与诊断

  • 暂存缓冲区 —— 相机和设置 Uniform 复用预分配的 ArrayBuffer 以避免每帧分配。
  • 双管道 —— 渲染器创建两个 GaussianPreprocessor 实例(SH + RGB),以便着色器可以在编译时专注于颜色模式。
  • 精度元数据 —— 动态点云在分发前调用 setPrecisionForShader,以便着色器知道如何解码 INT8/FP16 存储。
  • 原子热点 —— 极其密集的场景可能会受到原子操作的限制;在密集场景中减少 gaussianScaling、收紧剔除框或启用 mip-splatting 有助于降低每帧可见计数。
  • 调试工具 —— debugCountValues() 利用 debugCountPipeline(参见 src/utils/debug-gpu-buffers.ts)将 ONNX 计数缓冲区与模型参数 Uniform 进行比较,使间接计数问题更容易诊断。

参考

  • src/preprocess/gaussian_preprocessor.ts —— 权威的 TypeScript 实现。
  • src/shaders/preprocess.wgsl —— 上文讨论的计算着色器。
  • src/renderer/gaussian_renderer.ts —— 展示渲染器如何在 prepareMulti 期间实例化预处理器并调用 dispatchModel
  • src/sort/radix_sort.ts —— 详细说明预处理消耗的排序器资源 (sorter_bg_pre, sorter_uni, sorter_dis)。