NPU 3DGS 渲染优化基于贪心负载均衡的 Ascend C Alpha Blending 算子切分策略【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai本文聚焦 CANN 平台下 3DGS3D Gaussian Splatting光栅化核心算子 Alpha Blending 的 SIMD 化改造难题由于各像素块tile上投影的高斯球数量分布差异极大传统的按像素块均匀切分会让 vector core 之间负载严重不均产生单核过载、多核闲置的拖尾效应。文章在 gaussian_splatting_load_balance_optimization.md 的基础上结合仓库中 GetRenderSchedule.cpp、GaussianSort 等真实源码完整讲解先重后轻、分给最轻的贪心负载均衡算法、调度表schedule的数据结构、算子 Tiling 与 Workspace 工程细节以及实测收益数据。读完本文你将掌握该优化方案从原理、实现到验证的完整闭环并能在自己的 3DGS NPU 移植工作中直接复用这套负载均衡思路。一、为什么 SIMD 架构下 3DGS 需要负载均衡1.1 Alpha Blending 算子的计算流程Alpha blending 是 3DGS 光栅化过程中的核心步骤给定一个特定的相机视角和对应的投影平面将空间中的 3D 高斯球颜色表达渲染到各个像素上从而构建成对应视角的渲染图像。其计算公式如下$$ C \sum_{i \in N} c_i \alpha_i \prod^{i-1}_{j1}(1-\alpha_j) $$$$ \alpha_i o_i \exp(-\frac{1}{2} \delta ^ T * \Sigma * \delta) $$其中 $c_i$ 为高斯球颜色$\alpha_i$ 为该高斯球对该像素的透射不透明度$o_i$ 为高斯球整体不透明度$\delta$ 为像素中心到高斯球中心的偏移$\Sigma$ 为二维协方差矩阵。由于 $\prod(1-\alpha_j)$ 的存在渲染结果依赖严格按深度排序的高斯球累乘顺序。整体算法流程如下按 tile 对 vector 核进行切分每次调用渲染一个 tile 上的所有像素块由于累乘操作的存在必须保证对高斯球排序后依次渲染所以用 for 循环遍历每个高斯球进行渲染对每个高斯球并行计算其对 tile 上所有像素块的影响计算出 alpha 值和颜色贡献对每个像素块更新累积颜色和透明度值。1.2 GPU 与 SIMD 架构的调度差异在开源 GPU 实现中硬件支持 Warp 级别的线程调度GPU 会以极小粒度动态平衡各线程束的负载因此无需额外的负载均衡策略。但在 SIMD 架构如 NPU 的 vector core / AIV下实现该算法时通常采用按像素块对 vector core 进行均匀切分的方式。由于每个像素块上投影的高斯球数分布差异较大——有的 tile 只被几个高斯球覆盖有的 tile 则被成百上千个高斯球覆盖——均匀切分策略可能导致每个 core 上处理的负载不均衡造成个别 core 计算拖尾整体渲染耗时随之增加。从源码结构看仓库对 vector 核数量的获取方式为Tiling 阶段通过platformInfo.GetCoreNumAiv()取得 AIVAI Vector core数量见 gaussian_sort.cpp运行时侧则通过acl.get_device_capability(0, 1)[0]查询见 rasterizer.py。负载均衡的均衡对象正是这些 AIV 计算核。二、贪心负载均衡策略先重后轻分给最轻2.1 优化目标将所有像素块的高斯球分配给固定的 vector 核最终使各 vector 核的负载尽可能均衡避免出现单核过载多核闲置的情况。2.2 策略原则采用贪心算法遵循先重后轻分给最轻的原则通过局部最优的选择以达到全局较优的结果。即负载越重的像素块越优先被分配且每次都分配给当前累计负载最轻的 vector 核。2.3 实现流程排序按照像素块对应的高斯球数将像素块从大到小进行排序。贪心分配将排序好的像素块优先分配到当前高斯球数最小的 vector 核上以此继续分配像素块直到将所有像素块分配完。2.4 一个四 tile、双核的完整示例以如下场景为例四个像素块每个像素块高斯球数分别为 50、100、10、20分配到 2 个 vector 核。目标分配状态如下图所示采用像素块平均切分策略会造成倾斜问题core-0 分配 150 个高斯球tile-0 的 50 tile-1 的 100core-1 仅分配 30 个高斯球tile-2 的 10 tile-3 的 20采用负载均衡策略则分配相对均匀core-0 分配 100 个高斯球core-1 分配 80 个高斯球。两种切分策略的结果对比如下图第一步排序。先将四个像素块按照像素块对应高斯球数从大到小排序。排序结果为 tile-1100、tile-050、tile-320、tile-210。第二步贪心分配。将排序好的像素块优先分配到当前高斯球数最小的 vector 核上分配 tile-1当前 core-0 和 core-1 上高斯球数均为 0将 tile-1 分配到 core-0core-0 负载更新为 100分配 tile-0当前 core-1 负载为 0core-0 负载为 100将 tile-0 分配到 core-1core-1 负载更新为 50分配 tile-3当前 core-1 负载为 50core-0 负载为 100将 tile-3 分配到 core-1core-1 负载更新为 70分配 tile-2当前 core-1 负载为 70core-0 负载为 100将 tile-2 分配到 core-1core-1 负载更新为 80。完成所有像素块分配后core-0 需处理 100 个高斯球core-1 需处理 80 个高斯球。具体过程如下图所示可以看到贪心策略将最大负载从均匀切分的 150 降到 100两个核的负载差从 120 缩小到 20渲染耗时由负载最高的核决定因此整体渲染效率得到显著提升。三、源码级实现GetRenderSchedule 调度表生成负载均衡的调度逻辑在仓库中由GetRenderSchedule算子实现核心算法位于 GetRenderSchedule.cpp 的get_render_schedule_impl函数L44-L100。该实现与文档中的策略一一对应排序阶段先对nums每个 tile 的高斯球数升序排序得到索引数组tile_idxes随后从rbegin()反向遍历即实现了从大到小分配L52-L55。贪心分配阶段使用小顶堆优先队列std::priority_queuestd::pairint, int, std::vectorstd::pairint, int, std::greater...维护每个 bin即每个 vector 核的当前累计负载每次取出负载最小的核pq.top()把当前最大的 tile 分配给它再将其负载更新后压回队列L63-L75。这正是分给最轻原则的程序化表达。边界处理代码中还处理了两个实际工程细节——其一effective_bins std::min(T, num_bins)动态收缩 bin 数量以避免空核L51其二RESERVE_CAPACITY_MARGIN 2为每个 bin 预分配额外容量以应对贪心算法分配不均的边界情况L32。四、调度表的数据结构与下游消费4.1 schedule 张量布局get_render_schedule输出的调度表是一维数组按如下三段布局SCHEDULE_DATA_SECTION_PER_TILE 2即每个 tile 需要存储 tile_id 与累积 offset 两个字段schedule: [num_bins, tile_ids(T), tile_offsets(T)]前num_bins个元素每个核的结束索引bin 边界schedule[i] - schedule[i-1]即第 i 个核分到的 tile 数空核的结束索引与前一核相同中间T个元素按调度顺序排列的tile_id 序列即每个核实际负责渲染的像素块编号最后T个元素按原始 tile 编号排列的高斯球数累积偏移cumsum用于在排序后的全局高斯球数组中定位每个 tile 的数据区间L76-L98。4.2 GaussianSort 算子的调度消费调度表由GaussianSort算子消费用于将全局高斯球按深度优先顺序重排这是 Alpha Blending 累乘正确性的前提。从 gaussian_sort.cpp 的LoopProcessL355-L392可以看出其消费逻辑coreOffsetsGM_读取当前核blockIndex_的调度区间[startScheduleIdx, endScheduleIdx)明确本核负责哪些调度位置scheduleGM_按调度位置取出tileIdx得到本核负责的像素块tileOffsetsGM_与gaussian_cnt每个 tile 的高斯球数共同定位该 tile 在 depths/gs_ids 数组中的数据区间随后TileSort()对该区间的深度值排序得到sorted_gs_ids。也就是说排序算子的并行任务划分完全由负载均衡调度表决定核与核之间互不重叠、负载接近从数据准备阶段就保证了后续CalcRender前向/反向计算阶段的核间均衡。4.3 上层渲染流水线的调用链调度表的生成与消费被封装为get_render_schedule与gaussian_sort两个 Python 接口见 get_render_schedule.py、gaussian_sort.py在两条渲染链路中被使用在 gaussian_splatting/rasterization/rasterizer.py 中先由tile_sums每个 tile 的高斯球数调用get_render_schedule生成调度表再调用gaussian_sort得到重排后的高斯球 ID最后将lb_sched与sorted_gs_id一并交给calc_render完成逐像素渲染在 SLARM 的 npu_rendering.py 中_get_render_schedule从tile_offsets差分还原出每个 tile 的高斯球数再调用get_render_schedule(nums.cpu(), get_num_vector_core())生成调度随后传入CalcRender.apply完成颜色、alpha、flow 的渲染。五、工程落地Tiling 与 Workspace 设计负载均衡带来的一个关键工程问题是核间负载不均导致无法按固定 tile 数预估每核需要的资源。仓库在 gaussian_sort.cpp 的 Tiling 阶段做了相应处理核数确定与 blockDim 裁剪vectorNum platformInfo.GetCoreNumAiv()获取 AIV 核数blockDim min(tileNum, vectorNum)当 tile 数少于核数时只启动必要的核L91-L95max_tile_gauss 与 mask 对齐从算子属性读取max_tile_gauss全局 tile 中最大高斯球数按ALIGN_NUM 32向上对齐得到maxMaskNum作为每核排序缓冲的上界L81-L84Workspace 按全局最大负载申请源码注释明确说明——引入 B、C 维度外加负载均衡策略无法精细化处理故选择全局 tile 中最大高斯球数作为空间申请L113-L119即userWorkspaceSize maxMaskNum * sizeof(float) * WS_TENSOR_NUM * blockDim外加系统库所需的LibApiWorkSpaceSizeUB 内排序规模自适配maxSortNum由 UB 容量与排序所需张量数计算得出SortInUB与SortInGM子序列排序 冒泡归并两条路径按sortProcessNum_ maxSortNum_自动选择兼顾小 tile 的快速路径与大 tile 的正确性L341-L353。以上内容为排序阶段的负载适配渲染阶段CalcRender同样以lb_sched驱动逐核任务划分详见 CalcRender.cpp 中calc_render_fwd_double_clip_gsids等入口。六、正确性验证多规模测试用例负载均衡改变了核间数据划分但绝不能改变渲染结果。仓库在 test_gaussian_sort.py 中提供了 CPU 参考实现与 NPU 结果的逐元素对比验证测试规模覆盖[[2, 8], [2, 123], [10, 1234], [112, 23456], [222, 234567]]五组tile 数, 高斯球数组合横跨小场景到 23 万高斯球的大场景L69CPU 参考实现_gaussian_sort完全按调度表顺序对每个 tile 用torch.argsort(depths_t, stableTrue)稳定排序并写入对应 offsetL39-L62其语义即按深度升序全局重排NPU 端将lb_scheds、tile_sums转 int32、tile_depths、tile_gaussian_ids、sorted_offset、max_tile_gauss送入gaussian_sort算子L106-L115其中lb_scheds由get_render_schedule(tile_sums, num_bins)生成num_bins min(acl 查询到的 vector_num, T)判定标准assertRtolEqual对比 CPU 与 NPU 的sorted_gs_ids确保负载均衡重排后的高斯球序列与标准排序完全一致。七、实验数据与收益文档在一个包含13 万个高斯球的场景实际上经过过滤后需要渲染的高斯球数仅为67917 个上评估了按块均匀切分与负载均衡设计对渲染效率的提升效果。实验结果如下优化方法前向 device 耗时 (ms)反向 device 耗时 (ms)按块切分5.97315.908负载均衡4.59912.223根据表中数据计算前向耗时由 5.973ms 降至 4.599ms反向耗时由 15.908ms 降至 12.223ms正、反向耗时收益均约 23%。该收益正是来自单核过载、多核闲置问题的消除——负载最高的核决定了整体渲染耗时而贪心均衡把每核负载压到了接近理论平均值。八、小结与可复用的设计要点总结这套 NPU 3DGS 负载均衡方案的四个可复用要点以每个任务单元的计算量而非任务单元个数作为切分依据3DGS 中计算量与 tile 内高斯球数成正比必须按高斯球数加权分配贪心算法的两个原则缺一不可先重后轻保证大负载任务不被挤到末位分给最轻保证每一步都缩小核间差距最终逼近全局较优调度表schedule与数据重排sort解耦调度表只描述哪个核负责哪些 tile排序算子按调度表执行二者通过lb_sched/sorted_offset衔接渲染算子无需感知均衡细节资源申请按最坏情况兜底由于均衡后核间 tile 数不固定Workspace 按全局最大 tile 负载申请gaussian_sort.cpp并用多规模测试test_gaussian_sort.py保证正确性。如需进一步了解本仓库 3DGS 相关的其他优化手段如 Alpha Blending、裁剪、精确求交等可继续阅读 docs/3d_vision/gaussian_splatting 目录下的对应文档以及算子实现与测试所在的 ops/ascendc 目录。【免费下载链接】cann-recipes-embodied-ai本项目针对具身智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-embodied-ai创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考