CANN广播UB指令设计

实现③UB 内 Broadcast 指令【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills主文档broadcast_design.md | 参考代码code/broadcast_common.hUbBroadcast/FillUbShape思路先把紧凑源DataCopyPad搬进 UB再用AscendC::Broadcast矢量指令在 UB 内展开到目标 shape。占 Vector 单元但对尾轴对齐 小字节类型吞吐最高。还是唯一能广播中间计算结果的实现① ② 只能广播 GM 输入。适用前提tile 内要有 UB Broadcast 能展开的维dShapesShape。两种会出 nan 的反例cann-bench Maximum case 18 实测广播轴严格在外层brcInTilefalseUB 展不了外层轴 → 必须走 ②offset 寻址。广播轴是切分轴但ubFormer1tile 内仅 1 行dShapesShape空广播 → 必须走 ① NDDMA。即需inTileExpandable切分轴以下某轴广播或切分轴广播且ubFormer1。满足后任一成立即用 ③尾轴对齐且 dtype ∈ {INT8, UINT8, FP16, BF16, INT16, UINT16}B8/B16或 BigNLast非尾轴广播且尾轴字节 ≥nddma dcache/2源码注释 ~4096B或广播的是上游计算结果而非 GM 输入。Kernel 写法// 节选完整见 code/broadcast_add_kernel.cpp LocalTensorT src bufSrc_.GetT(); // 源长度按源 shape算不能用展开后的输出 tile 长度否则 [M,1] 尾轴广播会多搬 N 倍 int64_t srcLen UbSrcLen(inputStrides[i], inputDims[i], ubSplitAxis, shapeLen, rows); DataCopyPadCompact(yGm[off], src, srcLen); // ① 紧凑搬入 UbBroadcast(dst, src, outputDims, inputDims[i], ubSplitAxis, shapeLen, rows); // ② UB 内展开UbBroadcast内部code/broadcast_common.huint32_t dShape[8], sShape[8]; FillUbShape(dShape, sShape, outputDims, inDims, ubSplitAxis, shapeLen, curRows); // 首轴随主/尾块变 int64_t rank shapeLen - ubSplitAxis; // ★退化守卫所有 dShapesShape无可展开维如尾块 rows1→ 直接 UB→UB 拷贝否则空广播出 nan bool anyExpand false; int64_t cnt 1; for (int k 0; k rank; k) { if (dShape[k] sShape[k]) anyExpand true; cnt * dShape[k]; } if (!anyExpand) { // 无可展开维 → dtype 无关 UB→UB 拷贝 int32_t alignEle 32 / sizeof(T); int32_t cntAlign ((int32_t)cnt alignEle - 1) / alignEle * alignEle; // 向上对齐bufferelemNum 不越界 AscendC::DataCopy(dst, src, cntAlign); return; } AscendC::BroadcastTiling bt; AscendC::GetBroadcastTilingInfoT(rank, dShape, sShape, false, bt); AscendC::BroadcastT(dst, src, dShape, sShape, bt); // ← UB 矢量广播指令rankRR shapeLen - ubSplitAxisUB tile 内参与广播的维数。静态已知且 ≤4 → 可用编译期BroadcastT, R/GetBroadcastTilingInfoT, R性能更好动态或超阈值 → 运行期BroadcastT如样例R 作为运行期实参。尾块需把dstShape[0]/srcShape[0]改成ubTail/ 源尾块长度FillUbShape已按curRows处理。中间结果广播若广播的是上游计算节点的输出不是 GM 输入跳过DataCopyPadCompact直接对上游 UB tensor 调UbBroadcast。这是 ③ 相对 ① ② 的独占能力。注意退化守卫UbBroadcast内先检查是否存在dShapesShape的维若无如尾块rows1或选型漏判直接DataCopy(dst,src,对齐cnt)做 dtype 无关的 UB→UB 拷贝不用Adds避免 int64 等不支持而非调Broadcast防空广播出 nan。这是兜底正确选型应优先在 host 侧避免见上前提。UB 预算要多留一块紧凑源 buffer样例的bufSrc_③ 是 src dst 两块① 只需一块。tiling 的aliveBuf要含这块。每个走 ③ 的输入要各自一块bufSrc若多个输入都可能 ③如 binary 的 x1/x2 都广播共用一块TBuf会 RAW 冒险第二个输入的DataCopyPadCompact覆盖第一个的紧凑源。binary 算子应配bufSrc1_/bufSrc2_aliveBuf相应增大cann-bench Maximum 落地经验。GetBroadcastTilingInfo的 rank 必须与Broadcast调用一致。src/dst shape 的首轴是 UB 切分轴方向的行数后续轴是尾轴及以下广播轴在 src 上填 1。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考