已关闭
【AICPU】优化 Expand 广播过程临时内存占用 #3254
zhaowenrui创建于  23 天前关闭于  23 天前
zhaowenrui
zhaowenrui
23 天前 创建

问题背景

ops-math 的 Expand AICPU kernel 在广播过程中会先复制完整输入 Tensor,再按广播轴逐层生成中间 std::vector,最终再复制到输出。临时存储量随输出 shape 和 dtype 字节宽度增长,大 shape 下会产生不必要的堆内存占用和重复数据搬运。

以 INT64 [1,34,34] -> [32,34,34] 为例,单次 kernel 执行窗口内 C++ new/new[] 累计请求量为 626,816 B、93 次。

优化目标

  1. 从原输入直接写入最终输出,删除完整输入副本和逐层广播中间结果。
  2. 合并输入、输出相同的连续尾维,按连续块复制。
  3. 广播维通过 stride 和坐标进位映射输入位置。
  4. 将 dtype 分派由运行时 map<std::function> 改为 switch。
  5. 在复制循环外完成 shape 乘法溢出和输出容量校验。
  6. 保持 Expand 接口、广播语义、支持类型和数值位模式不变。

验收标准

  • 覆盖 BOOL/FLOAT16/BFLOAT16/FLOAT/INT8/UINT8/INT32/INT64 和 INT32/INT64 shape。
  • 覆盖 rank 0~8、-1、identity、标量、前/中/后轴广播、交替广播、连续尾块和大 shape。
  • 内存统计证明临时存储不再随输出元素数增长。
  • NPU AICPU 全量性能测试无确认回退。
  • 原有 UT、专项边界 UT、ASAN 和 pre-commit 通过。

已完成验证

  • 8 种 dtype × 5 档 shape 共 40 个内存场景,优化后均为 1,504 B、61 次申请。
  • INT64 [1,34,34] -> [32,34,34]:626,816 B -> 1,504 B,减少 99.76%。
  • 844 个 NPU 性能点全部完成,确认性能回退 0 个;全部点 speedup 中位数 3.119x。
  • 256 MiB 大 shape 提升 8.17x~54.85x。
  • AICPU 日志与 msprof 探针确认实际进入 Expand AICPU,三种计时中位数最大相对差 6.03%。

关联

likedislike
zhaowenruizhaowenrui
23 天前 修改了issue 的描述
zhaowenruizhaowenrui
23 天前 关联了pull request:[性能优化] 降低 Expand AICPU 广播临时内存
zhaowenrui
zhaowenrui
23 天前 评论:

/assign

likedislike
CANN-robotCANN-robot成员
23 天前 将 zhaowenrui666 设为负责人
CANN-robotCANN-robot成员
23 天前 关闭了 issue
CANN-robotCANN-robot成员
22 天前 添加了label:resolved