企业官网建设流程全解析

• 计算机视觉
• 图形学
• 深度学习

导读：本文以 crates/brush-sort/README.md 为核心，深入讲解 brush（3D Reconstruction for all，一个基于 Rust + Burn/CubeCL 的 3D 高斯溅射（3D Gaussian Splatting）渲染与训练项目）中用于 GPU 端基数排序（radix sort）的brush-sortcrate。你将理解它如何支持"按任意位数排序、以间接调度（indirect dispatch）处理 GPU 已知元素数量的数组"两大能力，掌握其基于 FidelityFX Radix Sort 的 5 阶段内核流水线（count → reduce → scan → scan-add → scatter）的完整原理，并看到它在渲染器中深度排序与 Tile 排序两个关键环节的真实调用方式，以及配套的测试与基准使用方法。

一、模块定位：为什么 3D 高斯溅射渲染需要 GPU 排序

brush 是一个面向 WebGPU 的 3D 高斯溅射渲染与训练项目（项目描述为 "3D Reconstruction for all"），渲染核心位于 crates/brush-render。在它的光栅化流水线中，排序是一个无法回避的基础操作，涉及两个典型场景：

1. 深度排序（Depth Sort）：高斯溅射以半透明 blob 的方式绘制，必须按到相机的深度从远到近（或近到远）排序，才能得到正确的混色结果；
2. Tile 排序（Tile Sort）：渲染器把屏幕划分成若干 tile，需要把每个高斯与它覆盖的 tile 的对应关系（intersection 记录）按 tile id 排序，以便按 tile 顺序紧凑地打包光栅化任务。

这两处排序的数据规模非常可观——单个场景可能包含上百万个高斯、数千万条 intersection 记录（见后文测试用例），而排序结果必须在 GPU 上立即被后续 kernel 消费，不能来回拷贝到 CPU。因此，brush 把排序实现为纯 GPU 的 WebGPU 兼容基数排序，也就是 crates/brush-sort 这个 crate。

二、brush-sort 是什么：WebGPU 兼容的基数排序实现

crates/brush-sort/README.md 对该模块的定位非常精炼：

WebGPU compatible radix sort. It's based on this implementation, which in turn is based on FidelityFX Radix sort. It allows sorting up to a given number of bits, and sorting an array with a GPU known number of elements using indirect dispatches.

翻译过来即：一个 WebGPU 兼容的基数排序。它基于 googlefonts/compute-shader-101 的一个实现（PR #31），而后者又源自 AMD 的FidelityFX Radix Sort算法。其两大核心能力是：

• 支持按任意给定的位数排序（up to a given number of bits）：排序范围可以是任意 bit 数，最多 32 位；
• 支持对 GPU 已知元素数量的数组进行间接调度排序：元素数量由 GPU 端 buffer 决定（运行时才知道），通过间接调度（indirect dispatch）启动排序 kernel，无需 CPU 回读。

在 brush 的整体架构中，brush-sort与 crates/brush-scan（GPU 前缀和/扫描原语）、crates/brush-cube（CubeTensor 创建与设备抽象）协作，构成 crates/brush-render 渲染流水线的底层基础设施。

2.1 依赖关系与构建配置

从 crates/brush-sort/Cargo.toml 可以看到它的依赖设计：

• brush-cube：提供CubeTensor、CubeDevice、create_tensor/create_tensor_from_slice等张量与设备工具；
• brush-scan：提供block_scan、cube_exclusive_sum、cube_sum、lds_index等扫描原语以及BLOCK_SIZE、ELEMENTS_PER_THREAD、WG常量；
• burn、burn-wgpu、burn-cubecl：通过 Burn 框架获得 WebGPU/CubeCL 后端与张量类型系统；
• bytemuck：用于安全地做零拷贝位模式转换；
• tracing：用于性能 trace 跨度（如Radix sort、DepthSort、Tile sort等 trace_span）。

dev-dependencies 则包括rand（随机数据生成）、tokio（异步测试）、divan（基准框架）以及wasm-bindgen-test（WebAssembly 环境测试）。它还声明了一个独立基准目标：

[[bench]] name = "sort_bench" harness = false path = "benches/sort_bench.rs"

三、算法架构：基于 FidelityFX Radix Sort 的五阶段内核流水线

brush-sort的核心入口是 src/lib.rs 中的radix_argsort函数——注意它是argsort：除了返回排序后的 keys，还会同步返回对应的 values（即"原下标"，因为 values 通常被初始化为 0..n 的索引），这正符合渲染器"拿到排序后的下标再做 gather"的需求。

3.1 常量与每趟位宽

在 src/kernels.rs 中定义了关键的算法参数：

pub const WG_USIZE: usize = WG as usize; // 每 workgroup 线程数，WG = 256 pub const BITS_PER_PASS: u32 = 4; // 每趟处理 4 bit pub const BIN_COUNT: u32 = 1 << BITS_PER_PASS; // 16 个 bin（0x0 ~ 0xF） pub const BIN_COUNT_USIZE: usize = BIN_COUNT as usize;

其中WG（256）、ELEMENTS_PER_THREAD（4）、BLOCK_SIZE（= WG × ELEMENTS_PER_THREAD = 1024）这些扫描相关的常量从brush-scan复用（见 crates/brush-scan/src/kernels.rs）。每趟处理 4 位（16 个 bin），因此：

• 排序sorting_bits位需要ceil(sorting_bits / 4)趟（见for pass in 0..sorting_bits.div_ceil(4)）；
• 排序 32 位 key 需要 8 趟，排序 10 位 key（渲染器 Tile 排序场景）需要 3 趟（见 benches/sort_bench.rs 中的注释）。

3.2 入口函数radix_argsort的调用约束

radix_argsort(input_keys, input_values, sorting_bits)首先做一系列防御性断言（src/lib.rs）：

• keys 与 values 的长度必须相等；
• sorting_bits <= 32，即最多排序 32 位；
• keys 与 values 必须是连续内存（contiguous）的CubeTensor。

随后它计算出调度规模：

let max_n = input_keys.shape()[0] as u32; let max_needed_wgs = max_n.div_ceil(BLOCK_SIZE); let num_wgs_count = max_n.div_ceil(BLOCK_SIZE); let num_reduce_wgs_count = num_wgs_count.div_ceil(BLOCK_SIZE) * BIN_COUNT;

这些数量被写入num_keys_buf（I32 张量），供各 kernel 以间接调度的方式自行计算边界，这正是 README 所说 "sorting an array with a GPU known number of elements using indirect dispatches" 的实现方式——CPU 侧只需要按最大可能规模发起一次调度，kernel 内部通过读取num_keys_arr[0]并配合terminate!()提前退出多余 workgroup（见各 kernel 开头的if group_id >= num_wgs { terminate!(); }）。

3.3 五阶段内核：一趟排序的完整数据流

每一趟（一个 4-bit pass）依次启动 5 个 kernel（src/lib.rs）：

阶段 1 ——sort_count_kernel：构建直方图

每个 workgroup 负责BLOCK_SIZE（1024）个元素，用共享内存中的 16 个原子计数器（Shared::<[Atomic<u32>]>::new_slice(BIN_COUNT_USIZE)）统计本块内每个 bin 的元素个数，然后把直方图按counts[bin * num_wgs + group_id]的布局写入全局计数 buffer（src/kernels.rs）。

阶段 2 ——sort_reduce_kernel：块直方图规约

把每个 bin 的各块计数累加为该 bin 的全局块级总量。由于num_wgs可能很大，这里再做一层"块的块"（BIN_COUNT * ceil(num_wgs / BLOCK_SIZE)个 reduce workgroup），每个 reduce workgroup 用cube_sum把ELEMENTS_PER_THREAD个块的计数求和，结果写入reducedbuffer（src/kernels.rs）。

阶段 3 ——sort_scan_kernel：跨块前缀和

对reducedbuffer 做分块 exclusive scan。它使用block_scan以BLOCK_SIZE为 chunk 循环推进，并用carry跨 chunk 传递累加值——这也是解决大规模数据（reduce 块数超过 BLOCK_SIZE）时前缀和正确性的关键（src/kernels.rs）。

阶段 4 ——sort_scan_add_kernel：扫描结果回填

把上一阶段得到的每个 bin 的全局偏移加回对应的块直方图，得到每个块内各 bin 的排他性起始偏移（写入 counts buffer），供 scatter 阶段直接使用（src/kernels.rs）。

阶段 5 ——sort_scatter_kernel：就地重排（scatter）

这是最复杂的内核（src/kernels.rs）。每个 workgroup 处理自己负责的BLOCK_SIZE个元素：

1. 先用cube_exclusive_sum配合打包计数技巧（把每个线程 4-bit key 的每 2 bit 打包进一个 u32 的不同字节，一次 cube scan 同时得到 2 个 2-bit bin 的排他偏移，见packed_input = 1u32 << (bit_key * 8u32)与后续bin_offsets解包）在共享内存中完成块内重排；
2. 再用块内局部直方图与bin_offset_cache（阶段 4 得到的全局偏移）合成total_offset = global_offset + local_offset，把 (key, value) 对写入输出 buffer 的正确位置（out[total_offset] = local_key; out_values[total_offset] = local_value）。

一趟结束后，cur_keys/cur_vals更新为输出张量，进入下一趟，直到sorting_bits全部处理完，返回(cur_keys, cur_vals)。

3.4 与 brush-scan 的复用关系

brush-sort大量复用brush-scan提供的原语：block_scan（共享内存块扫描）、cube_exclusive_sum/cube_sum（cube 级排他前缀和/求和，内部采用 plane/subgroup 两阶段扫描，MAX_PLANES = 32适配不同硬件 wave 宽度，见 crates/brush-scan/src/kernels.rs）以及lds_index（冲突无关的共享内存布局索引）。这也解释了为什么brush-scan的源码注释明确写着"GPU scans shared across brush：a cube-wide block scan for kernels that need a prefix sum in shared memory (the radix sort)..."。

四、实战调用：渲染器中的深度排序与 Tile 排序

在 crates/brush-render/src/render.rs 中，radix_argsort被调用了两次，是观察它真实用途的最佳样例。

4.1 深度排序（Depth Sort）

渲染器先把所有高斯投影到相机并求交，得到depths（每个可见高斯的深度）与global_from_presort_gid（排序前的全局 id），随后以32 位为排序范围进行深度排序（crates/brush-render/src/render.rs）：

let (_, global_from_compact_gid) = tracing::trace_span!("DepthSort") .in_scope(|| radix_argsort(depths, global_from_presort_gid, 32));

返回的global_from_compact_gid是"按深度排序后的紧凑全局 id"，之后用它通过int_gather重新排列intersect_counts，并交给prefix_sum计算每个高斯的累计 tile 命中数（crates/brush-render/src/render.rs）。注意深度是 f32，这里以 32 位排序意味着算法对 f32 的位模式直接按无符号整数语义排序——这对应基准中Random32场景的注释（"what the gaussian-splat depth sort does internally for non-negative floats"）。

4.2 Tile 排序（Tile Sort）

对于每条 intersection 记录，渲染器先生成tile_id_from_isect和compact_gid_from_isect，然后根据 tile 总数计算所需的排序位数（crates/brush-render/src/render.rs）：

let bits = u32::BITS - num_tiles.leading_zeros(); let (tile_id_from_isect, compact_gid_from_isect) = tracing::trace_span!("Tile sort") .in_scope(|| radix_argsort(tile_id_from_isect, compact_gid_from_isect, bits));

这里bits = u32::BITS - num_tiles.leading_zeros()是覆盖全部 tile id 所需的最小位数，充分体现了 README 所说"支持排序任意位数"的实用价值——例如 512×512 图像约 1024 个 tile，只需 10 位（3 趟）即可完成排序，相比固定 32 位省下大量调度开销。排序后的数组再交给get_tile_offsets生成每个 tile 的起始偏移，供光栅化 kernel 按 tile 顺序消费（crates/brush-render/src/render.rs）。

五、质量保障：测试与基准

5.1 正确性测试

src/lib.rs 内置了 4 组测试，全部通过wasm_bindgen_test(unsupported = tokio::test)同时支持 WebAssembly 浏览器环境与原生 tokio 异步环境，使用的设备是CubeDevice::Wgpu(brush_cube::test_helpers::test_device())，即真实的 WebGPU 设备：

• test_sorting：对 15 个包含小值、大值（如2^24 + 123）、重复值、0、递增/递减模式的 key 数组迭代 128 次，与 CPU 参考 argsort 逐元素比对 keys 和 values；
• test_sorting_big：模拟"一批高斯"的数据形态——按区间随机生成约上万条带重叠的区间键（类似高斯投影到 tile 的分布），验证真实负载下的正确性；
• test_sorting_large：3000 万（30_000_000）个元素，验证大规模下 2D dispatch 的正确性，并抽查若干下标确认 values 确实指向原始索引；
• test_sorting_above_scan_block_size：7000 万（70_000_000）个元素——这是针对一个真实回归 bug的测试：排序在约 67M keys 时存在静默数据损坏（silent corruption），根因是num_reduce_wgs > BLOCK_SIZE时sort_scan_kernel的扫描越界。修复后该测试逐元素验证 7000 万个 key 与对应 value 的完全正确（注释见 src/lib.rs）。

5.2 微基准（divan）

benches/sort_bench.rs 提供了独立于渲染的微基准，用divan（harness = false，手动main）测量纯排序开销。要点包括：

• 两种 key 形态：TileIds（key 范围 [0, 1024)，匹配渲染器 512×512 图像的 tile 预算，10 bit → 3 趟）和Random32（全 32 位随机 key，8 趟）；
• 四个规模档位：1M（普通帧）、10M（中等帧）、30M（旧版本约最大可用规模）、70M（刚越过 ~67M 扫描悬崖，见 benches/sort_bench.rs）；
• 计时方式：输入 buffer 在计时区外预分配，计时区内只包含radix_argsort调度 + 最小化回读（只读最后一个 key）强制 GPU 完成，从而测得纯粹的排序 dispatch 时间；
• u32 直传：用create_tensor_from_slice(.., DType::U32)直接构造CubeTensor，避免 Burn 的 i32 类型构造器对高位为 1 的 key 报错，从而覆盖完整的 u32 取值范围。

运行基准的方式：

cargo bench -p brush-sort

运行测试的方式（原生设备）：

cargo test -p brush-sort

六、小结与使用建议

brush-sort是 brush 渲染流水线中"GPU 侧有序性"的基石：它基于 FidelityFX Radix Sort 的 5 阶段流水线（count → reduce → scan → scan-add → scatter），每趟处理 4 位（16 bin），通过间接调度支持 GPU 已知元素数量的大规模数组，并通过可配置的sorting_bits在"排序位数"与"调度趟数"之间做出精细权衡——深度排序用满 32 位，Tile 排序则只使用覆盖 tile id 所需的最小位数。

从其实现与测试中可以提炼出几条实用经验：

1. 按需选择排序位数：bits = u32::BITS - num_tiles.leading_zeros()这类最小位数计算能显著减少趟数，是"支持任意位数排序"这一设计最直接的收益；
2. 大规模数据需关注扫描悬崖：sort_scan_kernel以BLOCK_SIZE为 chunk 的循环扫描设计，正是为了避免 reduce 结果超过单块容量时的越界损坏，测试test_sorting_above_scan_block_size（70M 元素）验证了这一点；
3. 间接调度依赖 GPU 已知的计数 buffer：所有 kernel 通过num_keys_arr[0]读取元素数量并terminate!()空 workgroup，CPU 侧无需回读即可安全处理运行时才知道的数组长度；
4. 输入必须 contiguous：radix_argsort对 keys/values 有连续性断言，调用方需先into_contiguous或确保切片来源连续。

如果要在自己的模块中使用brush-sort，可以参考渲染器的调用模式：准备一对等长的CubeTensor（keys + 原始索引 values），计算需要的排序位数，然后调用radix_argsort获得排序后的 (keys, values) 对，再以 values 作为 gather 索引驱动后续 kernel。

• 计算机视觉
• 图形学
• 深度学习