已合并
aclnnNormalTensorTensor A2/A3分支内存优化 #4137
aclnnNormalTensorTensor A2/A3分支内存优化 #4137
已合并
Ding_Jing创建于 7月18日
Ding_Jing
Ding_Jing成员
7月18日

描述

aclnnNormalTensorTensor 在 A2/A3(ASCEND910B / DAV_2201)分支计算 out = stateLessOut * std + mean 时,Mul、Add 使用非 inplace 接口会各自新分配一块中间 tensor 内存。本 PR 复用计算过程中新建的中间 tensor 内存,减少内存分配:

  • Mul:当 stateLessOut 的 shape 与 out 一致、且 std 的 dtype 与 stateLessOut 一致时,改用 MulInplace(std, stateLessOut) 复用 stateLessOut 内存作为输出;否则回退 Mul(stateLessOut, std)。注意 MulInplace 输出为第二入参、Mul 输出 dtype 取第一入参,故两分支入参顺序相反。
  • Add:类型提升已保证 meanCastmulOutCast 的 dtype 一致,且 mean 恒能广播进 mulOutCast,直接改用 AddInplace(meanCast, mulOutCast) 复用 mulOutCast 内存作为输出。

优化仅改变内存使用方式,最多减少 2 次中间内存分配,对功能和精度无影响。

关联的Issue

#2260

测试

test_aclnn_normal_tensor_tensor_l2.cpp 新增 2 个 L2 UT 覆盖两条分支(UT 默认 soc 为 ASCEND910B/DAV_2201,走 A2/A3 分支):

  1. case_broadcast_mean_ND_normal:mean{1,3} 向 std{2,3} 广播,stateLessOut 取 mean 的 shape{1,3} 与 out{2,3} 不一致,Mul 走非 inplace 回退分支;Add 侧 meanCast{1,3} 广播进 mulOutCast{2,3}。
  2. case_broadcast_std_ND_normal:std{1,3} 向 mean{2,3} 广播,stateLessOut 取 mean 的 shape{2,3} 与 out 一致且 dtype 相同,Mul 走 MulInplace 复用内存分支。

两个用例均校验 aclnnNormalTensorTensorGetWorkspaceSize 返回 ACL_SUCCESS;原有 UT(含 case_mixed_dtype_ND_normal 等)保持通过。另对测试文件做了少量格式清理(SetUpTestCase/TearDownTestCase 单行化、补齐文件末尾换行)。

文档更新

无文档变更(本 PR 为内部内存优化,不涉及接口/README 变更)。

类型标签

likedislike
Pull Request已成功合入, 合并人@CANN-robot
(感谢 Ding_Jing 的贡献)
Ding_JingDing_Jing成员
7月18日 创建了 pull request,commit df524010
atomgit-bot
atomgit-bot
7月18日 评论:

变更摘要

此 PR 对 aclnnNormalTensorTensor 算子的 A2/A3 分支 (CommonLogicGeneralNormal) 进行了内存优化:在 MulAdd 调用处,当中间 tensor 的 shape/dtype 满足条件时改用 Inplace 版本复用其内存作为输出,从而减少一次内存分配。

主要改动

  • Mul 调用改为条件 MulInplace:新增 canMulInplace 判断逻辑,当中间 tensor stateLessOut 的 shape 与 out 一致且 dtype 与 std 相同时,调用 l0op::MulInplace(std, stateLessOut, ...) 复用 stateLessOut 内存作为输出;否则回退到原来的 l0op::Mul(stateLessOut, std, ...)。注意两分支入参顺序相反,以适配 MulInplace 输出为第二个入参、Mul 输出 dtype 取第一个入参的语义差异。
  • Add 调用改为 AddInplace:将 l0op::Add(mulOutCast, meanCast, ...) 替换为 l0op::AddInplace(meanCast, mulOutCast, ...),直接复用 mulOutCast 的中间内存作为输出,省略一次内存分配。依赖上方类型提升逻辑已保证 dtype 一致、且 mean 恒可广播进 mulOutCast
  • 新增广播场景单测:在测试文件 test_aclnn_normal_tensor_tensor_l2.cpp 中增加 case_broadcast_mean_ND_normalcase_broadcast_std_ND_normal 两个用例,分别覆盖 mean 向 std 广播(Mul 无法 inplace 的回退路径)和 std 向 mean 广播(Mul 可 inplace 的优化路径),确保 A2/A3 分支的内存优化逻辑正确。
likedislike
atomgit-bot
atomgit-bot
7月18日 评论:

代码审查

✅ 未发现问题

likedislike
CANN-robotCANN-robot成员
7月18日 添加了label:cann-cla/yes
CANN-robot
CANN-robot成员
7月18日 评论:

CLA Signature Pass

Ding_Jing, thanks for your pull request. All authors of the commits have signed the CLA. 👍

likedislike
此处折叠了72条消息 查看更多
RuiWang_成员
7月22日 评论:

/lgtm
/approve

likedislike
CANN-robot
CANN-robot成员
7月22日 评论:

The following users do not have permission to comment /lgtm or /approve on any module in this PR:
刘伟

likedislike
CANN-robotCANN-robot成员
7月22日 添加了label:lgtm
CANN-robotCANN-robot成员
7月22日 关闭了关联的issue
CANN-robotCANN-robot成员
7月22日 合入了pull request