已关闭
[Bug] Greater/SplitV/Mul AICPU kernel 存在可消除的额外内存占用,且 Greater 对 0 元素输入、Mul 对过小输出存在未定义行为 #3024
Ding_Jing创建于  9月3日关闭于  9月19日
Ding_Jing
Ding_Jing成员
9月3日 创建

问题描述

Greater 与 SplitV 两个 AICPU kernel 存在可消除的额外内存占用,且 Greater 在 0 元素输入下存在未定义行为。

1. Greater:Eigen 广播表达式带来固定约 33 KB 的库内部临时分配

greater_aicpu.cpp 用 TensorMap + .reshape().broadcast() 承载逐元素比较,Eigen 的 block scratch 在所有 shape 上恒定占用约 33 KB。同时 BroadcastCompute<T, RANK> 被 8 个 RANK × 11 个 dtype 实例化共 88 份,生产编译口径下 .text 达 920,888 字节。

补充:x > y 展开的 scalar_cmp_op 其 functor_traits::PacketAccess 对 Greater 支持的全部 dtype 均为 false,即当前实现并未从 Eigen 获得向量化收益。

2. Greater:0 元素输入触发 Eigen 断言

Compute 无空张量早返回(同仓 Mul、RealDiv 均有),零长度维被直接交给广播评估器,触发 TensorBroadcasting.h:146 的 input_dims[i] > 0 断言。生产构建定义 EIGEN_NO_DEBUG 使断言被编掉,于是变为无任何防护的未定义行为。

3. SplitV:与 num_split 线性相关的冗余显式申请

每次 Compute 有四处 8×num_split 的分配,其中 output_data_vec 与成员 output_ptr_vec_ 内容完全重复,三个 compute 函数又按值接收 std::vector。峰值约 32 × num_split + 200 字节。

4. SplitV:格式串与实参类型不匹配

split_v_aicpu.cpp 中用 %ld 打印 int32_t 的 split_dim_,属未定义行为(同文件另一处对同一变量用的是正确的 %d)。

期望行为

在保持功能、dtype 支持列表、错误码语义完全不变的前提下降低两个算子的 kernel 内存占用,并保证性能不劣化。

环境信息

  • 芯片:Ascend 910B3
  • CANN:9.1.0
  • 架构:aarch64

补充:Mul(2026-09-08 追加,扩大本 Issue 范围)

Mul 的广播路径存在同一族问题,另有两处正确性问题,故并入本 Issue 一并跟踪。

5. Mul:Eigen 广播表达式带来 29-43 KB 的库内部临时分配

mul_aicpu.cpp 同样用 TensorMap + .reshape().broadcast()。广播路径上 Eigen 的 block scratch 恒定占用 29-43 KB(随 dtype 元素宽度变化),与输出规模无关:[1] x [1, 268435456](2.68 亿元素)实测峰值 33,136 B,[64] 小 shape 也在同一量级。同 shape 与标量快路径仅约 500 B,问题只在广播路径。

与 Greater 不同的是,scalar_product_op 的 PacketAccess 为 true,即 Mul 确实从 Eigen 拿到了向量化,因此替换必须先证明不劣化,而不是像 Greater 那样稳赚。

6. Mul:MulCalculate<RANK, T, OPTION> 被 8 RANK x 14 dtype x 2 对齐 = 224 份实例化

生产编译口径下 .text 达 3,712,488 B / 3,465 个函数符号。对照实验:把 RANK 分派从 8 路砍到 2 路,.text 即降到 969,884 B —— RANK 这一个维度贡献了 74% 的膨胀。

7. Mul:输出张量容量不足时写穿缓冲区

输出元素数由两个输入 shape 推导得出,out.reshape(result_shape) = ... 不校验推导值是否等于 output_tensor->NumElements()。调用方传入比广播结果更小的 output 时按推导值写入即越界;生产构建定义 EIGEN_NO_DEBUG,Eigen 自身的断言被编掉,于是没有任何防护。

8. Mul:复数含 NaN/Inf 时结果取决于元素落在哪个 lane

复数乘法有朴素式(四乘两加)与 C99 附录 G(__mulsc3,对含无穷分量的操作数做恢复)两套算法,只在操作数含 NaN 或 Inf 时不同。Eigen 的广播求值器在 2 元素 packet 跨越广播边界时退回标量 coeff() 走附录 G,其余位置走 packet 的朴素式。实测 [N,1] x [N,3] complex64:150,528 个输出里有 12 个与朴素式不同,全部满足 下标 % 3 == 2,即每个长度 3 的广播 run 的最后一个。同一对操作数,落在哪个 lane 决定它得到哪个答案。

likedislike
Ding_Jing
Ding_Jing成员
9月3日 评论:

/assign

likedislike
CANN-robotCANN-robot成员
9月3日 将 Ding_Jing 设为负责人
Ding_JingDing_Jing成员
9月3日 关联了pull request:Greater/SplitV AICPU kernel 内存优化
Ding_JingDing_Jing成员
9月4日 关联了pull request:Greater AICPU kernel 内存优化
Ding_JingDing_Jing成员
9月4日 关联了pull request:SplitV AICPU kernel 内存优化
Ding_Jing
Ding_Jing成员
9月4日 评论:

修复已按算子拆为两个 PR 提交:

  • 本 Issue 第 1、2 点(Greater 的 Eigen 内部占用与 0 元素未定义行为)→ #5364
  • 本 Issue 第 3 点(SplitV 与 num_split 线性相关的显式申请)→ #5365

另外补充一条评估结论:同批排查的 Mul 算子不做修改。其 kernel 源码中显式内存申请为 0 处,少量分配全部来自框架 BCalcInfo + Bcast 的 O(rank) 容器;Eigen 内部占用约 38–43 KB 但恒定,实测 [1] × [1,268435456](2.68 亿元素)峰值仍仅 33,136 B,无任何 O(N) 项。且 scalar_product_op 的 functor_traits::PacketAccess 为 true,Eigen 对乘法确实做了向量化——实测以 stride 迭代替换后,-O2 下慢 1.35–5.0 倍,-O3 下仍慢 1.35–2.58 倍,故不替换。

likedislike
Ding_Jing
Ding_Jing成员
9月5日 评论:

更正前一条评论中关于 Mul 的性能结论。

前一条评论里我给出的「以 stride 迭代替换 Eigen 会慢 1.35–5.0 倍,-O3 下仍慢 1.35–2.58 倍」是错误的。错因是那次判别实验的 stride 实现没有做维度合并:[1,224,224,3] → [32,224,224,3] 场景下最内层维保持为 C=3,内层循环被调用 160 万次、每次只处理 3 个元素,向量化无从谈起。而同一手法在 Greater 上的实际交付实现是带 DropAndCollapse 的(合并后内层为 150,528)。等于用一个残废版本做了判据。

补做了带维度合并的门槛验证:14 个 dtype × 4 条路径 × 含 C=3/8/32 的 shape,共 84 档,输出全部 bit-exact。协议为绑核 + 预热 + 中位数,三个变体每轮轮换顺序。

路径 纯 -O2 过/总 加 optimize("O3") 过/总 未过的 dtype
same-shape 5/14 13/14 complex64 1.23
x 标量 6/14 13/14 complex64 1.22
外层广播 5/14 13/14 complex64 1.28
内层广播 C=3 1/14 3/14 11 个,最差 int8 5.41×
内层广播 C=8 5/14 9/14 float 1.79、int8 1.89、uint8 1.75 等
内层广播 C=32 6/14 14/14 无

修正后的结论:

  1. 纯 -O2 下 stride 确实打不过 Eigen(仅 5–6/14 过)——-O2 GCC 不向量化这些循环,而 Eigen 用显式 packet。这一点原结论方向没错。
  2. 但给热函数加 optimize("O3") 属性后,三条主路径均 13/14 通过,且普遍快 2–4 倍(int8 same-shape 0.436、uint8 标量路径 0.257)。原结论「必然劣化」不成立。
  3. 失效只剩两类,边界清晰:complex64(全路径慢 1.16–1.45×,因 Eigen 对 complex<float> 有专门 packet),以及最内层广播维过小(C=3 惨败,C=32 全过)。

因此本算子由「不宜优化」改判为「可优化,待立项」,但需要带守卫,且 optimize("O3") 是必需项而非可选项。另外此前低估了收益规模:真正的奖品不是那 43 KB 堆,而是 .text —— Mul 的 .text 为 2,633,412 B / 3,004 符号,其中约 85% 可归因于 Eigen 模板实例(8 RANK × 14 dtype × 2 对齐 = 224 份 MulCalculate 实例),是 AICPU 进程的常驻代码内存。

尚未验证、决定是否立项的关键点:加守卫(小 C 与 complex64 保留 Eigen)之后,.text 还能降多少。若回退路径能写成 rank 无关的扁平调用,实例数可由 224 降到 14 量级;若不能,.text 收益会被守卫吃掉大半。这一项需要原型实测后再定。

likedislike
Ding_JingDing_Jing成员
9月8日 修改标题为 “[Bug] Greater/SplitV/Mul AICPU kernel 存在可消除的额外内存占用,且 Greater 对 0 元素输入、Mul 对过小输出存在未定义行为”,原标题为“[Bug] Greater/SplitV AICPU kernel 存在可消除的额外内存占用,且 Greater 对 0 元素输入存在未定义行为”
Ding_JingDing_Jing成员
9月8日 修改了issue 的描述
Ding_JingDing_Jing成员
9月8日 关联了pull request:Mul AICPU kernel 内存优化
songkai111成员
9月19日 评论:

你好,当前issue较长时间未贡献,我们将先关闭该issue,后续如有需要请重新打开。

likedislike
Ssongkai111成员
9月19日 issue状态由 进行中 改变为 已拒绝
Ssongkai111成员
9月19日 关闭了 issue