CANN/cannbot-skills: 进阶通用切分技术

进阶通用切分任意 ubSplitAxis / 大 shape / 尾块非对齐【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills主文档broadcast_design.md | 参考代码code/broadcast_common.hComputeTiling入门样例broadcast_add已切换到本文的通用切分。本文解释为什么不能写死ubSplitAxis0以及通用切分的完整公式——这是把指导用到真实大 shape如 cann-bench case 4 int32 8192×8192、case 10 fp16 363×367×373的关键。1. 为什么不能写死 ubSplitAxis0写死「尾轴整条进 UB、沿 0 轴切」只在尾轴元素数 ≤ elemNum时成立。真实 case 的尾轴常远超 UB场景问题尾轴 rowLen elemNum如 8192 fp32一整行放不进 UB必须切到尾轴内部大 shape 无广播BRC_NONE同样会尾轴超 UB需与广播输入一样的多轴切分高维 shape切分轴可能落在中间某轴dimProductBeforeUbInner不再等于ubOuter2. 通用切分算法ComputeTiling切分对象是输出 shape。从尾轴向外累乘找能塞进 UB 的最内切分轴若连尾轴一行都放不下则切尾轴本身。// 1) 选切分轴从尾轴向外尽量把内层整轴塞进 UB int splitAxis 0; int64_t innerBelow 1; // innerBelow ∏ outputDims[splitAxis1 .. rank-1] for (int ax rank - 1; ax 0; ax--) { if (innerBelow * outputDims[ax] elemNum) { splitAxis ax; break; } innerBelow * outputDims[ax]; splitAxis ax; } // 2) 切分大小 if (innerBelow elemNum) { // 连一行都放不下 → 切尾轴本身 ubSplitAxis rank - 1; innerBelow 1; ubFormer elemNum; // 每块 elemNum 个元素 } else { ubSplitAxis splitAxis; ubFormer elemNum / innerBelow; // 一次放几行 } ubFormer clamp(ubFormer, 1, outputDims[ubSplitAxis]); ubOuter ceil(outputDims[ubSplitAxis] / ubFormer); ubTail outputDims[ubSplitAxis] - (ubOuter - 1) * ubFormer;关键dimProductBeforeUbInner通用公式dimProductBeforeUbInner (∏ outputDims[0 .. ubSplitAxis-1]) × ubOuter // UB tile 总数ubSplitAxis 0∏(空)1 → 退化为ubOuter即入门样例的特例。ubSplitAxis rank-1切尾轴∏ outputDims[0..rank-2] × ubOuter。这个值喂给GetAxesIndices/GetGmOffsetubSplitAxis≠0 时必须用通用公式否则下标还原全错。完整实现见ComputeTiling。多核totalTiles dimProductBeforeUbInner; blockFormer ceil(totalTiles / coreNum); // 先按核数算每核最多几个 tile blockNum ceil(totalTiles / blockFormer); // 再回收用不上的核 blockTail totalTiles - (blockNum-1)*blockFormer; // ∈[1, blockFormer]顺序很关键若先blockNummin(coreNum,totalTiles)再blockFormerceil(totalTiles/blockNum)会出现blockTail0如 totalTiles100、coreNum64 → blockFormer2、blockTail-26非尾核处理越界 tile。必须先定 blockFormer 再回收核。3. 尾块非 32B 对齐ubTail splitDim - (ubOuter-1)*ubFormer往往不是 alignEle 的倍数如质数 shape、大 shape 展平。处理方式搬运DataCopyPad的blockLen以字节计内部自动 pad 到 32B非对齐尾块直接传真实长度即可无需手动 mask/paddingDataCopyPadCompact已如此。计算Add/Broadcast等矢量指令的count传真实元素数AscendC 处理非对齐尾数不要按对齐数多算否则踩到相邻 tile。输出DataCopyPad写回同样传真实tile字节不会越界写。即连续布局下尾块非对齐全部由 DataCopyPad 兜底开发者不需要额外分支。仅当用DataCopy非 Pad搬非对齐数据时才需自己 padding。4. 大 shape 无广播BRC_NONE无广播的大 shapecann-bench case 4/5/6和广播输入共用同一套切分ComputeTiling只看输出 shape与是否广播无关。BRC_NONE 输入在 kernel 走CopyInPlain其len rows * inputStrides[ubSplitAxis]因 stride 与输出一致即为tile天然支持尾轴被切。所以无广播但 shape 很大无需特殊处理——切分通用化后自动覆盖。5. NDDMA 在 ubSplitAxis≠0 的参数装配MakeNddmaParams/BroadcastNddma已按ubSplitAxis参数化ubSplitSize语义是「切分轴方向本 tile 的元素/行数」ubSplitAxis 0ubSplitSize rows行数。ubSplitAxis rank-1切尾轴ubSplitSize ubFormer尾轴一段的元素数axisInsideUb NDDMA_DIM-1只填切分轴一维 高位补 1。axisInsideUb NDDMA_DIM - (shapeLen - ubSplitAxis)rank 越大、切分轴越靠外需补的 size1 高位轴越少。前置条件rank shapeLen - ubSplitAxis ≤ NDDMA_DIM(5)超过时 Host 选型回退 ③PickBroadcastMode官方 atvoss 则走 with-loop NDDMA本样例未实现。6. cann-bench case 对照caseshape / dtype切分结果ComputeTiling4int32 8192×8192尾轴 8192 elemNum → 切尾轴(ubSplitAxis1)ubFormerelemNum无广播走 BRC_NONE7/8/9质数非对齐 shapeubTail 非对齐 → DataCopyPad 自动兜底10fp16 363×367×373尾轴 373 若 ≤ elemNum 则切中间轴(ubSplitAxis1)dimProductBeforeUbInner363×ubOuter这些 case 现在都能由通用切分覆盖写死 ubSplitAxis0 的旧样例会在 case 4/10 上因「一行放不进 UB」失败。【免费下载链接】cannbot-skillsCANNBot 是面向 CANN 开发的用于提升开发效率的系列智能体本仓库为其提供可复用的 Skills 模块。项目地址: https://gitcode.com/cann/cannbot-skills创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考