CANN广播OneDim优化

发布时间:2026/7/15 8:57:32
CANN广播OneDim优化 实现④OneDim 合轴塌一维快路径标量广播【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills主文档broadcast_design.md | 原语code/broadcast_common.hTryOneDim/ComputeOneDimTiling/OneDimCalcCore/OneDimLoadInput 端到端样例Host Kernelz x yy 为标量思路很多广播在合轴把相邻、广播状态一致的轴乘到一起后会塌成一维此时每个输入要么是满 shape 的连续大块要么是纯标量一个值。这类场景没必要套多维 NDDMA 参数装配、也不必调 UBBroadcast指令——满 shape 输入一维连续逐块DataCopyPad搬入即可标量输入首块用Duplicate把那一个值铺满 UB后续块直接复用完全不走逐块 DMA。这是z x lr、z x * scale、alpha·x y这类「标量参与 elementwise」以及同 shape elementwise 的最快写法。固定成本也最低tiling 退化成一维只需 1D 总长 / 块大小 / 核数切分让 kernel 运行期自己推。对标 atvossSCH_MODE_ONE_DIM_ADVANCEschMode 202tiling 极简dimLen/tileNum/blockNum/scalarFlag、kernel 运行期派生ubOuter/ubTail/blockFormer/blockTail省去结构体拷贝与多维下标还原。atvoss 在入出参 4 或 DAG 带 Var 标量属性时退回普通SCH_MODE_ONE_DIM201Host 预切分本指导只吸收 Advance 形态——纯 AscendC 下没有 Var 概念运行期标量属性由算子自行经 tiling 传入即可。适用合轴后能塌成一维即每个输入都满足「纯标量」或「满 shape 连续」之一✅z x scalar标量广播✅z x op yx、y 同 shape无广播纯 elementwise✅ 多输入混合部分满 shape、部分标量。❌[M,1] → [M,N]、[1,N] → [M,N]这类部分轴广播合轴后仍是多维不属本路径回到 ①②③ 选型。判定见TryOneDim逐输入算元素总数1记为标量置scalarFlag对应位dimLen是满 shape二者皆非则不能塌一维。Host选型 极简 TilingOneDim 是在三类选型之前的前置分流先TryOneDim命中就走 OneDim否则才进 §1.3 的PickBroadcastMode。// 节选完整见 code/onedim_add_tiling.cpp。在通用切分 / PickBroadcastMode 之前先试 OneDim int64_t dimLen; int32_t scalarFlag; if (TryOneDim(td, dimLen, scalarFlag)) { // 合轴塌一维每输入满 or 标量 OneDimTilingData ot; ComputeOneDimTiling(ot, dimLen, scalarFlag, dtSize, coreNum, ubSize, /*aliveBuf*/IN_NUM 1); // 写 ot 进 tilingDatatilingKey 标记走 OneDim 分支 return; } // 否则ComputeTiling(...) 逐输入 PickBroadcastMode(...)见 broadcast_design.md §1ComputeOneDimTilingcode/broadcast_common.h只算三个量不预切分tileNum AlignDown(ubSize / (aliveBuf * dtSize * DB_LOOP), alignEle); // 单块元素数 ubOuter ceil(dimLen / tileNum); // 共几块 blockNum min(coreNum, ubOuter); // 块数不足核数则减核小 shape 自适应 // dimLen / tileNum / blockNum / scalarFlag 即全部 tilingubTail/blockFormer/blockTail 交给 kernel极简 TilingData 是 Advance 形态的精髓搬运/切分细节全在 kernel 运行期由这几个标量推出tiling 体积小、下发开销低。Kernel 写法每核循环参数由OneDimCalcCore运行期推导标量输入用单块TBuf首块Duplicate一次、后续复用满 shape 输入用TQue双缓冲。// 节选完整见 code/onedim_add_kernel.cpp。z x yx 满 shape、y 可能是标量scalarFlag bit1 __aicore__ inline void Process() { bool xScalar ot_.scalarFlag (1 0); bool yScalar ot_.scalarFlag (1 1); OneDimCoreParam cp OneDimCalcCore(ot_.dimLen, ot_.tileNum, ot_.blockNum); // baseOffset/loops/tailLen int64_t off cp.baseOffset; for (int64_t loop 0; loop cp.loops; loop) { int64_t len (loop cp.loops - 1) ? cp.tailLen : ot_.tileNum; // 最后一块可非对齐 LocalTensorT xl xScalar ? xBuf_.GetT() : qx_.AllocTensorT(); OneDimLoadInput(xl, xGm_, off, len, xScalar, /*firstTile*/loop 0); // 标量仅首块 Duplicate if (!xScalar) qx_.EnQue(xl); LocalTensorT yl yScalar ? yBuf_.GetT() : qy_.AllocTensorT(); OneDimLoadInput(yl, yGm_, off, len, yScalar, loop 0); if (!yScalar) qy_.EnQue(yl); LocalTensorT xv xScalar ? xl : qx_.DeQueT(); LocalTensorT yv yScalar ? yl : qy_.DeQueT(); LocalTensorT zv qz_.AllocTensorT(); AscendC::Add(zv, xv, yv, len); qz_.EnQue(zv); if (!xScalar) qx_.FreeTensor(xv); if (!yScalar) qy_.FreeTensor(yv); LocalTensorT zo qz_.DeQueT(); AscendC::DataCopyExtParams ext{1, (uint32_t)(len * sizeof(T)), 0, 0, 0}; AscendC::DataCopyPad(zGm_[off], zo, ext); // 一维连续写回 qz_.FreeTensor(zo); off ot_.tileNum; } }OneDimLoadInput内部code/broadcast_common.hif (isScalar) { if (firstTile) AscendC::Duplicate(dst, gm.GetValue(0), len); // 标量值铺满 UB之后块复用同一块 } else { DataCopyPadCompact(gm[offset], dst, len); // 满输入一维连续搬非对齐 Pad 兜底 }注意标量不进双缓冲队列标量首块铺满后整轮复用应放单块TBuf若误塞进 depth2 的TQue并每轮 Alloc会丢掉只填一次的收益还要处理 ping/pong 两块各填一次。标量首块的len首块用Duplicate(dst, val, len)铺当前块长度即可后续块直接复用dst无需重铺各块长度只有最后一块可能更短复用块多出的尾部不参与计算/写回无害。尾块非对齐一维总长dimLen常非alignEle倍数最后一块tailLen非对齐——DataCopyPad自动兜底计算指令count传真实len即可同 advanced_tiling.md §3。小 shape 减核ubOuter coreNum时ComputeOneDimTiling把blockNum降到ubOuter避免空核OneDimCalcCore的blockFormer/blockTail据此自洽。UB 预算标量块按 1 份非双缓冲、满输入/输出按双缓冲计入aliveBuf示例取IN_NUM1为保守上界。doubleGetValue/Duplicate标量对 double 需按位转 int64 处理同 atvoss 对double的 reinterpret其余 dtype 直接用。与 ①②③ 的关系OneDim 是合轴塌一维时的前置快路径优先级高于三类选型——命中即用命中不了才进PickBroadcastMode。三类处理的是合轴后仍多维的部分轴广播[M,1]/[1,N]等。决策树见 broadcast_design.md §3。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考