ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

CANN 昇腾 3DGS 视锥剔除优化:GaussianFilter 融合算子设计与实现解析

CANN 昇腾 3DGS 视锥剔除优化:GaussianFilter 融合算子设计与实现解析 CANN 昇腾 3DGS 视锥剔除优化GaussianFilter 融合算子设计与实现解析【免费下载链接】cann-recipes-spatial-intelligence本项目针对空间智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-spatial-intelligence导读本文围绕 CANN 昇腾平台上 3D Gaussian Splatting3DGS训推优化实践中的视锥剔除Culling融合算子优化展开系统讲解投影预处理阶段四类高斯剔除判据、基于 int8 位图掩码的GaussianFilter算子设计、8 元素对齐的多核 tiling 策略、基于SyncAll的两阶段分核计算流程以及反向传播中的 Scatter 数据还原与Transpose搬运优化。读者阅读后可掌握该算子从 tiling 到 kernel 再到反向还原的完整实现链路并可直接在 algorithms/gaussian_splatting 样例中验证其效果。一、算子背景为什么渲染前必须做视锥剔除在 3DGS 中Culling剔除是渲染前的关键优化步骤目的是提前过滤掉对当前帧渲染无贡献的 3D 高斯基元减少后续光栅化、Alpha blending 的计算量大幅提升渲染效率。它的核心逻辑很简单只保留当前相机能看到的高斯丢弃看不到的高斯。在投影预处理阶段采用的是视锥体剔除具体包含四类判据剔除类型判据逻辑目的无效高斯剔除高斯必须是合法的正定矩阵非退化、非扁平剔除协方差矩阵行列式非正的高斯视锥体深度方向剔除剔除深度小于近裁切面太近或大于远裁切面太远的高斯只保留视锥体深度范围内的高斯视锥体水平/垂直方向剔除剔除高斯投影到屏幕空间后的 2D 椭圆 x/y 方向半径为 0 的高斯剔除投影退化的高斯屏幕外剔除仅保留投影椭圆与屏幕有重叠的高斯剔除完全在屏幕外的高斯从仓库测试用例 test_gaussian_filter.py 中的 CPU 参考实现可以精确还原这四类判据的组合方式valid (det 0) (depths near_plane) (depths far_plane) radius[~valid] 0.0 inside ( (means2d[..., 0] radius[..., 0] 0) (means2d[..., 0] - radius[..., 0] width) (means2d[..., 1] radius[..., 1] 0) (means2d[..., 1] - radius[..., 1] height) ) proj_filter torch.logical_and(inside, valid)其中det为 2D 协方差矩阵行列式depths为投影深度means2d/radius为投影到屏幕空间的 2D 椭圆中心与半径width/height为屏幕分辨率。GatherMask核内实现gaussian_filter.h通过CompareScalar与And逐条完成x±rx、y±ry与屏幕边界的比较、det0、near_planedepthfar_plane的掩码累积与参考实现严格对应。二、算子实现与优化总览GaussianFilter算子的实现与优化分为两步前向融合实现GaussianFilter将视锥剔除融合为一个算子避免相当耗时的取 index tensormove操作并返回filter——其中比特位为 1 代表对应高斯有效比特位为 0 代表对应高斯无效。测试参考实现test_gaussian_filter.py展示了 8 个高斯打包进 1 个 uint8 的编码方式filter_bool proj_filter.bool() remainder gaussian_nums % 8 if remainder ! 0: pad_size 8 - remainder filter_bool F.pad(proj_filter, (0, pad_size), modeconstant, valueFalse) matrix_r (gaussian_nums 7) // 8 filter_reshaped filter_bool.reshape(batch_dims, camera_nums, matrix_r, 8) powers torch.tensor([1, 2, 4, 8, 16, 32, 64, 128], dtypetorch.uint8, ...) filter_uint8 (filter_reshaped.to(torch.uint8) * powers).sum(dim-1, dtypetorch.uint8)反向还原投影预处理反向算子计算前进行反视锥剔除还原出原始 Tensor无效高斯的梯度计算中自然置为 0。三、实现难点比特位与元素的映射对齐int8 类型每个元素可存储 8 个高斯的掩码状态因此在计算时必须严格保证高斯数据与 int8 的比特位一一映射。这要求在核内迭代计算时输入高斯数据的分片必须是 8 的整数倍否则会出现比特位跨元素映射的混乱导致掩码判断错误。核间数据依赖与全局同步由于算子的切分涉及核间数据依赖必须引入全局同步操作但同步操作会带来性能开销需要在保证精度正常的情况下减少全局同步的次数。反向索引构建反向操作的核心是将过滤后的数据分散还原到原始位置核心难点在于对于无规则的filter构建GatherMask的数据分散索引。四、tiling 分核8 元素对齐与尾核反推考虑到高斯数量普遍为十万以上测试用例覆盖 [1,1,10000]、[2,7,18471]、[1,1,117611]、[1,1,188152] 等规模见 test_gaussian_filter.py因此在 tiling 阶段并不对 batchsize 和相机数进行分核而是对高斯N 维进行分核。具体 tiling 需要考虑以下 2 点filter 类型为 int8需要对头核进行 8 元素对齐对分核进行 8 元素对齐后尾核可能计算出负数因此需要重新反推使用的核数再进行尾核计算。假设 NPU 的核数为coreNum、高斯总数为N、头核处理的高斯数为perCoreN、尾核处理的高斯数为lastCoreN分核计算如下$$ perCoreN \left \lceil \left \lceil N / coreNum \right \rceil / 8 \right \rceil \cdot 8 $$$$ usedCoreNum \left \lceil N / perCoreN \right \rceil $$$$ lastCoreN N - usedCoreNum \cdot perCoreN $$这与 host 侧 tiling 实现 gaussian_filter.cpp 完全对应int64_t blockLength ((gaussNum maxCoreNum - 1) / maxCoreNum / SIZE_OF_FILTER SIZE_OF_FILTER - 1) * SIZE_OF_FILTER; maxCoreNum (gaussNum blockLength - 1) / blockLength; int64_t lastBlockLength gaussNum - blockLength * (maxCoreNum - 1); tiling.set_blockLength(blockLength); tiling.set_lastcoreNum(lastBlockLength);其中SIZE_OF_FILTER 8gaussian_filter.cpp即掩码位图的对齐粒度。尾核lastBlockLength由gaussNum - blockLength * (maxCoreNum - 1)计算避免了对齐后尾核为负数的问题。此外tiling 还根据 UB 大小计算单次迭代可处理的高斯数perloopNumgaussian_filter.cpp并预留 workspace 用于核间计数交换RESERVED_WORKSPACE_SIZE 16MB。tiling 数据字段定义见 gaussian_filter_tiling.h包括needCoreNum、batchNum、cameraNum、gaussNum、width、height、nearPlane、farPlane、blockLength、lastcoreNum、perloopNum、hasCompensations。五、kernel 实现GaussianFilter中 kernel 部分主要考虑的点是分核计算完后的搬运到 GM 上时的偏移计算。kernel 入口gaussian_filter.cpp按 tiling key 分发到模板类GaussianFiltertrue/falsetrue/false表示是否存在可选的compensations输入。5.1 偏移计算因为 UB 大小限制每个核一次迭代中无法过滤出所有有效高斯因此引入两个变量第 $i$ 个核$i$ 取值从 0 到 $usedCoreNum-1$计算出来的有效高斯数量为 $cntPerCore_i$第 $i$ 个核的第 $j$ 次迭代$j$ 取值从 0 到 $loopN_i$$loopN_i$ 取决于 UB 大小计算出来的有效高斯数量为 $cntPerCore_{i,j}$。那么第 $i$ 个核的第 $j$ 次迭代中需要把计算结果搬运出去的偏移为$$ offset \sum_{k0}^{i-1}{cntPerCore_k} \sum_{k0}^{j-1}{cntPerCore_{i,k}} $$即前面所有核的有效高斯总数加上本核前面所有迭代的有效高斯总数。该偏移公式在 kernel 的SubProcess中通过offsetFilterCore_核间偏移由cntPerCoreLocal.GetValue(k)累加得到与offsetFilterLoop核内迭代偏移由cntPreSum_.GetValue(k-1)得到相加实现gaussian_filter.h。5.2 两阶段计算流程SyncAll 同步分阶段显然第 $i$ 个核的计算偏移依赖于第 $0$ 到 $i-1$ 核的计算结果而所有的核又是同步计算的无法一次直接算出每次迭代的偏移。因此这里使用SyncAll同步分阶段来实现 Culling 的分核计算全程只做一次全局同步兼顾精度与性能。Phase1过滤与计数首先执行 $loopN$ 次循环迭代在每次迭代中先完成filter的搬入、计算与搬出同时统计当前迭代内 cnt 结果的累积值将核内迭代计数数组cntPerLoop并同步更新核内总计数cntPerCore。待该阶段循环结束后将cntPerCore存入位于workspace上的暂存空间执行全局同步操作SyncAll以确保所有计算核完成第一阶段计算随后更新并迁出 GM 上的输出偏移量coreOffset仅当当前计算核为最后一个核时执行全局计数的汇总计算并将最终计数结果迁出至 GM 上。对应源码gaussian_filter.h的执行顺序为循环内CopyInFilterSource→CalcFilter→CalcCntPerLoop→CopyOutFilter每轮将cntPerLoop写入cntPerLoops_、将累积值写入cntPreSum_循环结束后CopyOutCnt(cntPerCoreGm_[blockIdx_], cntPerCore)迁出到 workspace随后SyncAll()全局同步再从 workspace 搬回所有核的计数计算核间偏移offsetFilterCore_尾核负责把总有效高斯数写入cntGm_。Phase2掩码过滤与输出将计算出的filter重新搬入再次执行 $loopN$ 次循环迭代在每次迭代中依次搬入需要进行视锥剔除的输入以filter作为 mask 使用GatherMask执行视锥剔除随后结合第一阶段得到的核间偏移量coreOffset与对应迭代的核内计数cntPerLoop确定数据迁出的全局内存偏移地址。使用 3 个TQue保证 MTE2、Vector、MTE3 流水可以同步进行。对应源码中SubProcess的 Phase2 循环gaussian_filter.h依次调用CopyInFilter、ProcessMeans2dAndRadius、ProcessMeansAndDepth、ProcessCovars2d、ProcessColors、ProcessConics、ProcessOpacities(AndCompensations)七类数据通道的处理函数每个函数内部都是DataCopyIn 搬入 →GatherMask按掩码收集 → DataCopyOut 按偏移迁出输出地址统一使用(b*C_*N_ c*N_)*dim offsetNDim定位。当存在compensations输入时ProcessOpacitiesAndCompensations还会先执行Mul(opacities, compensations)完成不透明度补偿再过滤gaussian_filter.h。5.3 反向梯度还原与数据搬运Scatter 实现本算子的反向处理需将过滤后的数据恢复到原来位置需要进行数据分散操作。为了避免通过循环 标量计算带来的较长耗时本算子使用GatherMask、Sort以及Gather等多 API 结合的方式实现数据分散的优化。根据cntPerLoop对齐的Filter的掩码进行恢复恢复的目的数据分散如图Phase1构建新的排序索引newIndex。先构造与待处理元素等长的元素索引Index(0,1,...,cntPerLoop)通过GatherMask接口对Filter的比特位为 1 的元素进行Index收集存入头核headIndexTensor再对Filter的比特位为 0 的元素进行Index收集存入尾核tailIndexTensor。将两者拼接得到新的、用于Sort排序的索引newIndex。由于 Tensor 操作都需要 32 字节对齐拼接实现时需注意对齐问题这里采用的是先将两部分拷贝到 GM 上连续的地址再重新拷贝到 UB 上Phase2Sort排序 Gather重排恢复。Sort接口对newIndex排序后得到排序值score以及每一个score对应的Sort前所在的位置索引。Sort接口提供Extract接口实现排序值和排序索引的分离这里得到的索引正是用于gather进行元素收集的取元素地址偏移。在具体实现时还需注意Sort接口只支持降序排序需构造等差数列对Extract得到的索引再重排得到升序排序索引5.4 数据搬运优化用 Transpose 替代构造索引数据在LocalTensor以 N 维度切分处理搬运至GlobalMemory上在计算时为方便 vector 运算将高斯球维度放到 Tensor 的最后而为了得到原本输入的梯度的真实 shape数据输出要求改变最后一维和倒数第二维的顺序。以维度是(B,3,N)的数据搬运成维度是(B,N,3)为例在进行数据切分后的 UB 上的 shape 为(3, perLoopN)原本需要构造形如[0, perLoopN, perLoopN*2, 1, 1perLoopN, 1perLoopN*2, ...]的 index然后使用GatherMask收集对应索引的数据而构建 index 需要大量 Scalar 运算性能较差。因此这里使用Transpose接口搬运替代构造 index 的搬运进行优化思路如下图六、算子接口、封装与验证6.1 算子定义与属性GaussianFilter算子在 gaussian_filter.cpp 中以 OpDef 方式注册输入输出均为 ND 格式的 float32输入9 个必需 1 个可选means(B,3,N)、colors(B,3,N)、det(B,C,N)、opacities(B,N)、means2d(B,C,2,N)、depths(B,C,N)、radius(B,C,2,N)、conics(B,C,3,N)、covars2d(B,C,3,N)以及可选的compensations(B,C,N)输出means_culling、colors_culling、means2d_culling、depths_culling、radius_culling、covars2d_culling、conics_culling、opacities_culling八个过滤后的稠密张量外加filteruint8 位图掩码B*C*ceil(N/8)与cntint32 有效计数属性width、height屏幕分辨率Int、near_plane、far_plane近远裁切面Float平台AICore配置支持ascend910b与ascend910_93。6.2 Python 侧调用算子通过 pybind 与 Python 封装暴露给上层C host 层 GaussianFilter.cpp 完成输入维度校验如means.dim()3、means2d.dim()4、输出张量分配含filter按(gaussianNum7)/8对齐后调用aclnnGaussianFilterPython 封装 gaussian_filter.py 透传全部输入与width/height/near_plane/far_plane返回 8 个 culling 输出 proj_filtercnt。6.3 精度验证测试用例 test_gaussian_filter.py 针对[1,1,10000]、[2,7,18471]、[1,1,117611]、[1,1,188152]四组形状分别用 CPU 参考实现_gaussian_filter纯 torch 算子复现判据与位图编码和 NPU 自定义算子执行逐项assertRtolEqual对比 8 个 culling 输出、proj_filter与cnt保证融合算子与标准实现逐比特等价。七、在 3DGS 训推样例中的使用该算子属于 algorithms/gaussian_splatting 样例的四大 NPU 优化点之一视锥剔除融合算子优化。使用前需按样例 README 完成 CANN 环境准备CANN 8.2.RC1 与 torch/torch_npu 2.1.0/2.1.0.post12编译并安装自定义算子包cd ops/ascendc bash build.sh --python3.8 # 生成 dist/*.whl pip install dist/*.whl --force-reinstall随后即可在单卡 Atlas A2 环境启动 3DGS 训练或推理视锥剔除算子作为投影预处理链路的前置步骤自动参与每次前向渲染与反向梯度计算帮助过滤掉对当前帧无贡献的高斯基元从而降低下游光栅化与 Alpha blending 的计算量。总结GaussianFilter视锥剔除融合算子通过 int8 位图掩码、8 元素对齐 tiling、单次SyncAll的两阶段分核流程以及GatherMaskSortGather的反向 Scatter 还原将过滤无效高斯、压缩渲染负载这一关键优化完整下沉到 NPU 算子层配合Transpose搬运优化规避标量索引构建开销最终以融合算子的形态接入 3DGS 训推链路是理解昇腾平台上空间智能算法算子化改造的典型范例。【免费下载链接】cann-recipes-spatial-intelligence本项目针对空间智能业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-spatial-intelligence创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表