跳转至

Shaders模块架构

本文档深入探讨在 GPU 上运行的 WGSL 程序。重点介绍每个着色器的结构、阶段之间的数据流向以及它们期望的绑定组/资源。

阶段图

Gaussians (GPU buffers)
  │  ├─ preprocess.wgsl        -> Splat 缓冲 + 排序键/负载 + 间接计数器
  │  ├─ radix_sort.wgsl        -> 排序后的负载 (索引) + 更新后的间接调度
  │  ├─ gaussian.wgsl (VS/FS)  -> 颜色缓冲 (预乘 RGBA)
  │  └─ compositor (显示着色器位于此模块之外)

实用计算内核 (compress_gaussians.wgsl, convert_precision.wgsl, debug-helpers.wgsl) 使用相同的缓冲布局，但由工具/诊断按需调用。

绑定布局摘要

Group 0 : 相机 Uniforms (视图/投影 + 视口/焦距)
Group 1 : 点云数据 (打包的高斯数据, SH/原始颜色, Splat 缓冲)
Group 2 : 排序缓冲 (SortInfos, 深度键, 负载, 调度)
Group 3 : 渲染/模型设置 (RenderSettings, ModelParams)

渲染器在调用预处理时，会在 Group 3 中转发 ModelParams（变换、基础偏移量、每模型缩放、精度元数据）。gaussian.wgsl 只需要 Group 0 和 Group 1（相机 + Splat）以及排序后的负载索引。

预处理 (preprocess.wgsl)

主要职责

  • 读取打包的高斯数据；格式取决于 uModel.gaussDataType (0=f32, 1=f16, 2=int8, 3=uint8)。
  • 应用相机矩阵生成裁剪空间位置。
  • 将 3×3 协方差转换为 2×2 屏幕空间椭圆（通过雅可比矩阵 J 和相机旋转 W）。
  • 通过球谐函数 (evaluate_sh) 或直接 RGB 读取 (USE_RAW_COLOR) 来评估颜色。
  • 将 Splat 结构体写入共享缓冲，并将深度键/负载索引写入排序器缓冲。
  • 原子地递增 sort_infos.keys_size, drawIndirect.instance_count, 和 DispatchIndirect.dispatch_x。

重要结构体

struct CameraUniforms { view, view_inv, proj, proj_inv : mat4x4<f32>; viewport : vec2<f32>; focal : vec2<f32>; }
struct ModelParams { model : mat4x4<f32>; baseOffset : u32; num_points : u32; gaussianScaling : f32; maxShDeg : u32; ... precision fields }
struct Splat { v_0:u32; v_1:u32; pos:u32; posz:f32; color_0:u32; color_1:u32; }
uModel.num_points 会被 ONNX 计数缓冲（GPU 复制）覆盖，以便预处理在 gid.x >= num_points 时提前退出。

覆盖 (Overrides)

  • USE_RAW_COLOR (默认 false) – 跳过 SH 评估，将颜色缓冲视为 RGB 通道。
  • SH_LAYOUT_CHANNEL_MAJOR (默认 false) – 为 SH 系数选择通道优先 (channel-major) 布局。
  • DISCARD_BY_WORLD_TRACE 和 MAX_WORLD_TRACE – 可选保护，用于丢弃世界空间协方差迹 (trace) 过大的 Splat。

精度辅助函数

read_gaussian_pos_opacity, read_gaussian_cov, 和 read_color_channel 根据 uModel.gaussDataType/colorDataType 进行分支。INT8 路径使用 colorScale/colorZeroPoint 对 ONNX 精度转换器提供的数据进行反量化。

基数排序 (radix_sort.wgsl)

  • histogram_wg_size, histogram_sg_size, rs_radix_log2 等常量从 TypeScript 注入，以匹配 GPU。
  • 每个 Pass 有四个阶段：zero_histograms, calculate_histogram, prefix_histogram, scatter_even/scatter_odd。
  • 每个 even/odd scatter 分别处理两个 Pass (0,2) 和 (1,3)，对键/负载缓冲进行乒乓操作。
  • 使用回溯 (look-back) 机制来协调工作组之间的分区，同时最大限度地减少原子操作。
  • SortInfos 和 DispatchIndirect 通过绑定组 2 与预处理/渲染器共享。

高斯光栅化器 (gaussian.wgsl)

顶点着色器： - 通过 points_2d[indices[instance_id]] 读取 Splat 记录（indices 缓冲来自排序器负载）。 - 通过构建一个由特征向量缩放的屏幕对齐四边形 (mat2(v1, v2) × CUTOFF)，为每个实例生成四个顶点。 - 发送 NDC 位置（带有钳制的 z）并将屏幕空间坐标 + 颜色传递给片元着色器。

片元着色器： - 计算 r² = dot(screen_pos, screen_pos)。 - 丢弃 r² > 2*CUTOFF（紧密包围圆）的片元以节省带宽。 - 评估高斯衰减 exp(-r²) 并乘以 Alpha 通道（上限为 0.99）以避免全不透明伪影。 - 返回预乘颜色 vec4(rgb,1) * weight 以进行正确混合。

实用内核

  • compress_gaussians.wgsl – 读取 FP32 高斯数据并写入打包的 FP16/INT8 版本。由导入工具或离线管线使用。
  • convert_precision.wgsl – 类似于压缩器，但对现有的 GPU 缓冲进行操作，因此 ONNX 驱动的转换可以完全在 GPU 上运行。
  • debug-helpers.wgsl – 用于将计数器/缓冲复制到暂存区域以进行检查的小型计算内核（由开发者调试菜单使用）。

绑定更新

所有着色器都依赖于一致的绑定组顺序。渲染器/预处理器确保： - Group 0: 相机 Uniforms (跨阶段共享)。 - Group 1: 点云缓冲 (高斯, SH/原始颜色, Splat)。 - Group 2: 排序数据 (Infos, 深度键, 负载, 间接调度)。 - Group 3: 渲染/模型设置 (RenderSettings + ModelParams)。gaussian.wgsl 只需要 group 0/1；排序只需要 group 0。

间接调度和绘制计数完全在 GPU 上生成： 1. 预处理每 256*15 个 Splat 增加 sort_infos.keys_size 和 sort_dispatch.dispatch_x。 2. 基数排序通过 dispatchWorkgroupsIndirect 消费 sort_dispatch。 3. 渲染器在调用 pass.drawIndirect 之前将 sort_infos.keys_size 复制到间接绘制缓冲中。

这种架构使 CPU-GPU 同步保持最小：一旦缓冲分配完毕，GPU 决定每帧启动多少工作组和实例。