Shader API 参考

本参考文档重点介绍 Visionary 使用的 WGSL 入口点 (entry points) 和数据布局。每一节都列出了相关的结构体、绑定 (bindings) 和特化选项，以便您可以安全地修改或扩展着色器。

1. 预处理 (preprocess.wgsl)

入口点

@compute @workgroup_size(256, 1, 1)
fn preprocess(@builtin(global_invocation_id) gid: vec3<u32>,
              @builtin(num_workgroups) wgs: vec3<u32>)

- 每次调用处理一个高斯点。如果 gid.x >= uModel.num_points 则提前退出。 - 写入共享的 Splat 缓冲 (points_2d)、排序器缓冲和间接计数器。

绑定 (Bindings)

Group.Binding	资源
0.0	`CameraUniforms` (视图/投影矩阵 + 视口/焦距)
1.0	`gaussians_packed` : `array<u32>` (原始高斯数据，格式由 `uModel.gaussDataType` 选择)
1.1	`color_buffer` : `array<u32>` (SH 系数或 RGB 负载，格式由 `uModel.colorDataType` 选择)
1.2	`points_2d` : `array<Splat>` (读/写 Splat 缓冲)
2.0	`SortInfos` (原子计数器)
2.1	`sort_depths` : `array<u32>`
2.2	`sort_indices` : `array<u32>`
2.3	`DispatchIndirect`
3.0	`RenderSettings`
3.1	`ModelParams` (每模型变换 + 精度元数据)

关键结构体

CameraUniforms – 视图/投影矩阵、逆矩阵、视口、焦距。
Splat – 打包的特征向量、NDC 位置、高精度 Z 值、打包的 RGBA。
ModelParams – 变换矩阵、基础偏移量 (baseOffset)、点数 (num_points)、高斯缩放 (gaussianScaling)、最大 SH 阶数 (maxShDeg)、内核大小 (kernelSize)、不透明度缩放 (opacityScale)、截断缩放 (cutoffScale)、渲染模式 (rendermode)，以及量化字段 (gaussDataType, colorDataType, 缩放/零点)。

辅助函数

read_gaussian_pos_opacity(idx) / read_gaussian_cov(idx) – 根据存储精度 (FP32/FP16/INT8/UINT8) 进行分支处理。
read_color_channel / sh_coef – 根据 USE_RAW_COLOR 和布局覆盖读取原始 RGB 或 SH 系数。
evaluate_sh(dir, idx, sh_deg) – 真正的 SH 评估，最高支持 3 阶（每个颜色通道 16 个系数）。
applyDistanceScaling, applyPanning, applyRotation 位于轨道数学模块中，不在这里。

特化覆盖 (Specialization overrides)

名称	默认值	效果
`MAX_SH_DEG`	(注入值)	每个 Splat 评估的最大 SH 阶数 (0–3)。
`USE_RAW_COLOR`	`false`	将颜色缓冲视为 RGB 而非 SH 系数。
`SH_LAYOUT_CHANNEL_MAJOR`	`false`	在交错存储和通道优先 (channel-major) SH 存储之间切换。
`DISCARD_BY_WORLD_TRACE`	`false`	启用基于世界空间协方差迹 (trace) 的剔除。
`MAX_WORLD_TRACE`	`25.0`	启用上述功能时的阈值。

2. 基数排序 (radix_sort.wgsl)

入口点

@compute @workgroup_size(histogram_wg_size) fn zero_histograms(...)
@compute @workgroup_size(histogram_wg_size) fn calculate_histogram(...)
@compute @workgroup_size(prefix_wg_size)    fn prefix_histogram(...)
@compute @workgroup_size(scatter_wg_size)   fn scatter_even(...)
@compute @workgroup_size(scatter_wg_size)   fn scatter_odd(...)

工作组大小和基数参数在编译时预置（参见 GPURSSorter.processShaderTemplate）。

绑定 (Bindings - Group 0)

0.0 – SortInfos / GeneralInfo (原子计数器)
0.1 – Histogram buffer (atomic<u32>)
0.2 – Key buffer A (深度键)
0.3 – Key buffer B (乒乓缓冲)
0.4 – Payload buffer A (Splat 索引)
0.5 – Payload buffer B (乒乓缓冲)

注入的常量

const histogram_wg_size : u32 = ...;
const histogram_sg_size : u32 = ...;
const prefix_wg_size    : u32 = ...;
const scatter_wg_size   : u32 = ...;
const rs_radix_log2     : u32 = 8u;   // 256 桶
const rs_keyval_size    : u32 = 4u;   // 32位键需要 4 次 Pass

工作流

zero_histograms 清除直方图并重置 Pass 元数据。
calculate_histogram 使用共享内存并发填充所有 Pass 的直方图。
prefix_histogram 为每个数字位生成独占前缀和。
scatter_even/scatter_odd 使用前缀偏移量将键/负载移动到排序位置，并更新乒乓缓冲。

DispatchIndirect.dispatch_x 决定了在使用间接调度时启动多少个工作组；预处理会在写入 Splat 时增加它。

3. 高斯渲染器 (gaussian.wgsl)

顶点着色器

@vertex
fn vs_main(@builtin(vertex_index) vertex_id: u32,
           @builtin(instance_index) instance_id: u32) -> VertexOutput

- 通过 points_2d[indices[instance_id]] 获取 Splat。 - 为每个实例生成四个顶点（由特征向量 × CUTOFF 缩放的屏幕对齐四边形）。 - 输出裁剪空间位置、局部屏幕坐标和颜色。

片元着色器

@fragment
fn fs_main(in: VertexOutput) -> @location(0) vec4<f32>

- 计算 r² = dot(screen_pos, screen_pos)。 - 丢弃截断圆 (cutoff circle) 之外的片元。 - 评估高斯衰减 exp(-r²) 并乘以存储的 alpha（上限为 0.99）以避免完全不透明的钳位问题。 - 返回预乘颜色以进行正确的 Alpha 混合。

绑定 (Bindings)

Group.Binding	资源
0.2	`points_2d` (只读 Splat)
1.4	`indices` (由基数排序写入的已排序负载)

常量

CUTOFF = sqrt(log(255)) ≈ 2.3539 – 确保一旦高斯不透明度低于约 1/255 就丢弃片元。

4. 实用内核 (Utility kernels)

compress_gaussians.wgsl – 读取 FP32 Splat 并写入量化版本的计算着色器（离线使用或用于测试）。共享 Gaussian、ModelParams 和 RenderSettings 结构体。
convert_precision.wgsl – 用于 ONNX 精度管道的类似转换内核（采用现有的 GPU 缓冲，写入新的缓冲）。
debug-helpers.wgsl – 通过开发者工具公开的小型计算函数，用于复制/检查 GPU 缓冲。

5. TypeScript 集成 (src/shaders/index.ts)

export { default as preprocessShader } from './preprocess.wgsl?raw';
export { default as gaussianShader } from './gaussian.wgsl?raw';
export { default as radixSortShader } from './radix_sort.wgsl?raw';

渲染器和预处理管线在编译 WebGPU 管线之前会替换占位符字符串（例如 <injected> 替换为 MAX_SH_DEG）。

6. 使用检查清单

如果添加新的每模型 Uniform 字段，请更新 ModelParams（预处理直接读取它们）。
更改 SH 布局或精度模式时，请保持 read_color_channel / sh_coef_* 辅助函数与加载器同步。
任何新的着色器都必须遵循现有的绑定顺序，以便渲染器/预处理不需要额外的绑定组。
对于间接工作负载，确保预处理同时写入 SortInfos.keys_size 和 DispatchIndirect.dispatch_x；基数排序和渲染器无需 CPU 干预即可使用它们。