cannbot-skills广播算子优化指南 AscendC Broadcast 性能优化与实现选型指导总览【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills广播算子的性能关键在于按场景选对实现前置 OneDim 快路径 tile 内 NDDMA / DataCopyPad / UB 三类——选错直接慢甚至 nan选对才快。本文即该选型的最佳实践。纯AscendC API指导不依赖任何模板调度框架。提炼自 atvoss broadcast 的设计经验只保留两块可直接迁移的内容Tiling 选型/切分逻辑与OneDim 快路径 三类 Kernel 广播写法。 适用DAV_3510A5 / Ascend950C310。广播 某输入在某些轴上 shape1 / stride0需铺开到输出 shape。文档结构文件内容broadcast_design.md本文总览、实现速查、Tiling 选型与切分逻辑、决策树、约束onedim_design.md实现④ OneDim 合轴塌一维快路径标量广播★前置优先nddma_design.md实现① NDDMA 搬入即广播datacopypad_design.md实现② DataCopyPad 紧凑搬入 复用ub_broadcast_design.md实现③ UB 内 Broadcast 指令advanced_tiling.md进阶通用切分任意 ubSplitAxis / 大 shape / 尾块非对齐 / cann-bench casecode/broadcast_common.h常量 / TilingData /ComputeTiling/ 切分 helper / 四类广播原语含 OneDimcode/broadcast_add_tiling.cpp端到端样例①②③ 多维选型Host选型 切分code/broadcast_add_kernel.cpp端到端样例①②③ 多维选型Kernel 算子入口code/onedim_add_tiling.cpp端到端样例④ OneDim 标量广播HostTryOneDim 极简切分code/onedim_add_kernel.cpp端到端样例④ OneDim 标量广播Kernel 算子入口0. 实现速查一个前置快路径 tile 内三类#实现广播在哪发生核心 AscendC API占用单元典型适用④OneDim 合轴塌一维★前置标量首块 Duplicate / 满输入直搬Duplicate(ub,scalar,len)DataCopyPadVector(标量铺) / MTE2合轴后塌一维每输入要么满 shape 要么纯标量标量广播、同 shape elementwise①NDDMA 搬入即广播GM→UB 搬运中DataCopyT,DIM,cfg(ub,gm,MultiCopyParams)广播轴loopSrcStride0MTE2DMA不占 Vector大 dtypeFP32/INT32等、或 B8/B16 尾轴非对齐切分轴本身是广播轴②DataCopyPad外层广播offset 寻址不在搬运中靠 GM offset外层广播轴 stride0DataCopyPad(ub,gm,ext,pad)每轮搬一段连续数据MTE2不占 Vector广播轴全在切分轴之上外层、切分轴连续③UB 内 Broadcast 指令UB 内BroadcastT,R(dst,src,dstShape,srcShape,tiling)Vector前提广播轴在 tile 内(brcInTile)且尾轴对齐且 dtype∈{INT8,UINT8,FP16,BF16,INT16,UINT16}或非尾轴广播且尾轴≥dcache/2~4096B或广播中间计算结果判断口诀按优先级先试 ④——合轴后能塌一维每个输入要么满 shape、要么纯标量就走 OneDim标量首块 Duplicate、tiling 退化 1D、固定成本最低命中即用。塌不成一维部分轴广播如[M,1]/[1,N]才进 ①②③先看广播轴在不在 UB tile 内。全在 tile 外外层循环轴→ 一律 ②offset 寻址每轮连续搬UB 指令和 NDDMA 都展不了外层轴。在 tile 内才在 ③/① 间选尾轴对齐的小字节类型B8/B16或尾轴非广播的 BigNLast → ③Vector 吞吐高其余 → ①DMA 不占 Vector。1. Tiling 逻辑Host1.1 统一的切分模型无论哪类实现切分都是两层与普通 elementwise 一致对象始终是输出 shape多核切分blockNum 个核每核处理 blockFormer 个「UB 外循环单元」尾核 blockTail UB 切分选 ubSplitAxis 轴每次处理 ubFormer 行尾块 ubTail共 ubOuter 次外循环TilingData 字段见code/broadcast_common.h的BroadcastTilingData可直接照搬。关键字段inputStrides广播轴0、outputDims/Strides、ubSplitAxis/ubFormer/ubOuter、blockNum/blockFormer、dimProductBeforeUbInnerUB tile 总数、elemNum。ubSplitAxis不要写死成 0。尾轴常超 UB 容量如 8192 fp32必须按输出 shape 动态选切分轴、必要时切到尾轴内部。通用选轴 完整公式含dimProductBeforeUbInner通用式、大 shape、尾块非对齐见 advanced_tiling.md 与ComputeTiling。下文 §1.2–1.4 给要点。1.2 广播轴识别广播轴 该输入inputStrides[i][axis]0且outputStrides[axis]!0。这是所有判断的基础不是只看 dim1。必须扫描所有轴得到三个标志不能只取第一个广播轴否则同时有外层广播轴和 tile 内广播轴时会误判bool anyBrc false, brcInTile false, nonLastBrc false, lastAxisBrc false; for (int j 0; j shapeLen; j) { if (inputStrides[i][j] 0 outputStrides[j] ! 0) { anyBrc true; if (j ubSplitAxis) brcInTile true; // 广播轴落在 UB tile 内含尾轴 if (j shapeLen - 1) nonLastBrc true; // 存在非尾轴广播 else lastAxisBrc true; // 尾轴本身是广播轴 } }1.3 三类选型决策每个广播输入独立判断前置先试 OneDim④。进入下面的逐输入PickBroadcastMode之前先TryOneDim——若合轴后塌成一维每个输入要么满 shape 要么纯标量整算子走 OneDim 快路径ComputeOneDimTiling 标量Duplicate不再做三类切分。详见 onedim_design.md。塌不成一维部分轴广播才执行下面的三类选型。int PickBroadcastMode(int i) { // 承接 §1.2 扫描得到的 anyBrc/brcInTile/nonLastBrc/lastAxisBrc if (!anyBrc) return 0; // 非广播输入普通 CopyIn if (!brcInTile) return 2; // ★广播轴严格在外层 → ② DataCopyPadoffset 寻址切分轴非广播len 公式成立 // ★「在 tile 内」≠「UB 能展开」切分轴广播但 ubFormer1 时 tile 内仅 1 行(dShapesShape) // UB Broadcast 退化空广播 → nan。只有存在真正可展开维(dstsrc)才可选 ③否则走 ① NDDMA。 bool inTileExpandable false; for (int j ubSplitAxis1; j shapeLen; j) if (inputStrides[i][j]0 outputDims[j]1) inTileExpandable true; // 切分轴以下广播 if (inputStrides[i][ubSplitAxis]0 outputDims[ubSplitAxis]1 ubFormer1) inTileExpandable true; // 切分轴广播且多行 int64_t lastDim outputDims[shapeLen - 1], dt sizeof(T_i), alignEle 32 / dt; static const setDataType kUb {DT_INT8,DT_UINT8,DT_FLOAT16,DT_BF16,DT_INT16,DT_UINT16}; if (inTileExpandable) { if (lastDim % alignEle 0 kUb.count(dtype_i)) return 3; // ①尾轴对齐B8/B16 → ③ if (nonLastBrc !lastAxisBrc lastDim*dt GetNddmaDcacheSize()/2) return 3; // ②BigNLast → ③ } if (shapeLen - ubSplitAxis NDDMA_DIM) return 3; // rank5 without-loop NDDMA 装不下 → ③(含退化守卫兜底) return 1; // 其余 tile 内广播 → ① NDDMA }brcMode逐输入写进 TilingData / tilingKeykernel 据此分发。完整实现见code/broadcast_add_tiling.cpp。两道关键门控cann-bench Maximum case 18 实测!brcInTile广播轴严格在外层→ 一律 ②UB/NDDMA 只在 tile 内展开外层轴只能靠 offset 寻址。brcInTile但无可展开维切分轴广播且ubFormer1且rank≤5→ ① NDDMA此时 UB Broadcast 会dShapesShape空广播出 nanNDDMA 在ubFormer1下搬 inner 连续数据、外层靠 offset 广播仍正确。rank5时 without-loop NDDMA 装不下 → 仍走 ③由 kernelUbBroadcast退化守卫拷贝兜底。kernelUbBroadcast的退化守卫无可展开维则拷贝对两种情形都做双保险。1.4 切分大小与选轴ComputeTilingelemNum AlignDown( ubSize / (aliveBufferNum * maxDtypeBytes * DB_LOOP), alignEle ); // 从尾轴向外累乘选切分轴能整轴塞进 elemNum 就继续向外放不下则切该轴 // 连尾轴一行都放不下 → 切尾轴本身ubFormer elemNum // ubFormer/ubOuter/ubTail 沿 ubSplitAxis 算ubTail 可非对齐DataCopyPad 兜底 dimProductBeforeUbInner (∏ outputDims[0 .. ubSplitAxis-1]) × ubOuter // UB tile 总数通用公式完整实现ComputeTiling见code/broadcast_common.h推导与 cann-bench case 对照见 advanced_tiling.md。aliveBufferNum③ 需额外一块放广播前的紧凑源srcdst 两块①② 直接铺到目标块一块——选 ③ 的输入要把这块计入预算。2. 各实现各实现的适用、Kernel 写法、注意见分册④ OneDim前置、① NDDMA、② DataCopyPad、③ UB Broadcast。原语统一在code/broadcast_common.h① ② ③ 为BroadcastNddma/DataCopyPadCompact/UbBroadcast通用切分ComputeTilinghelperGetAxesIndices/UpdateAxesIndices/GetGmOffset/FillUbShape/UbSrcLen④ 为TryOneDim/ComputeOneDimTiling/OneDimCalcCore/OneDimLoadInput。端到端串接见样例 Host Kernel。3. 选型决策树★ 前置合轴后塌成一维吗(每个输入要么满 shape、要么纯标量 → TryOneDim) ├─ 是 → ④ OneDim 快路径标量首块 Duplicate、满输入直搬tiling 退化 1D ——命中即用整算子走此路 └─ 否部分轴广播如 [M,1]/[1,N]↓ 进入逐输入三类选型 该输入是广播输入吗(存在 stride0 轴) ——按从上到下的优先级匹配 ├─ 否 → 普通 DataCopyPad 搬入(brcMode0) └─ 是 ├─ 广播轴严格在外层(brcInTilefalse) → ② DataCopyPad ★必须先判外层轴只能 offset 寻址 └─ 广播轴在 UB tile 内(brcInTiletrue) ├─ rank(shapeLen-ubSplitAxis) 5 → ③without-loop NDDMA 装不下UB 退化守卫兜底无可展开维的情形 ├─ 有可展开维(inTileExpandable切分轴以下广播或切分轴广播且 ubFormer1) │ ├─ 尾轴对齐 且 dtype∈{INT8,UINT8,FP16,BF16,INT16,UINT16} → ③ UB-BRC │ ├─ 尾轴非广播 且 存在非尾轴广播 且 尾轴字节 ≥ dcache/2(~4096B) → ③ UB-BRC (BigNLast) │ └─ 其余 → ① NDDMA └─ 无可展开维(切分轴广播且 ubFormer1tile 内仅 1 行) → ① NDDMA ★走 ③ 会 dShapesShape 空广播出 nan 注实际代码 rank5 在 ①② 之后判可展开时①②仍优先但对无可展开维分支也生效——见 broadcast_add_tiling.cpp。 官方 atvoss rank5 走 with-loop NDDMA(SCH_MODE_NDDMA_WITH_LOOP_CACHE_BRC)本样例未实现故回退 ③。4. 自检清单④ OneDim 前置TryOneDim命中每输入满 shape 或纯标量即走快路径标量走单块TBuf首块Duplicate、不进双缓冲队列ComputeOneDimTiling只给dimLen/tileNum/blockNum切分由 kernelOneDimCalcCore运行期推广播轴判定基于inputStrides[i][axis]0 outputStrides[axis]!0不是只看 dim1brcMode逐输入计算并写入 TilingData/tilingKeykernel 分支与之一一对应① NDDMA广播轴loopSrcStride0、非广播轴填真实 stride维数 ≤ 5否则套循环② DataCopyPad广播靠GetGmOffset外层广播轴 stride0每轮连续搬即可复用为可选优化且判据用 offset 比较非切分轴 stride后者是死路径切分用ComputeTiling动态选 ubSplitAxisdimProductBeforeUbInner∏outputDims[0..split-1]×ubOuter尾轴超 UB 时切到尾轴尾块非对齐由 DataCopyPad 兜底③ UB-BRCelemNum预算多留一块紧凑源rank 静态(≤4)可走编译期BroadcastT,R尾块修正 shape 首轴DoubleBufferelemNum已按aliveBufferNum * dtypeBytes * 2反推精度与 naive 全展开实现对比FP32 1e-5、FP16 1e-35. 约束与适用范围平台DAV_3510C310。MultiCopy/NDDMA 多维 DMA 与Broadcast指令均为该架构 AscendC API。维度上限建议BROADCAST_MAX_DIMS 8单次 NDDMANDDMA_DIM 5。BLOCK_LENGTH 32字节对齐元素数 32 / sizeof(dtype)。只覆盖连续布局广播转置/非连续布局需搬运前做 layout 处理不在本文范围。④ OneDim 是整算子级的前置分流合轴塌一维则全算子走它①②③ 是塌不成一维时逐输入混用不同输入走不同 brcMode均由 tiling 决定。对标 atvoss 连续布局四个调度模式① ↔ NDDMA without/with-loopschMode 1/2、③ ↔ UB_BROADCAST101~109、④ ↔ ONE_DIM_ADVANCE/ONE_DIM202/201非连续布局族transpose 等不在本文范围。代码为范式级AlignDown/CeilDiv/GET_TILING_DATA/DTYPE_X等按工程宏/惯例占位落地需接 op proto / REGISTER / tiling 注册 / 对齐校验等工程件。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考