perf(arg_max_grad): arch35 轴下标改寄存器生成并开双缓冲, UB 记账由容量反解-ops-nn-AtomGit

Merged

perf(arg_max_grad): arch35 轴下标改寄存器生成并开双缓冲, UB 记账由容量反解 #9546

zl_hwAdd on Aug 31Merged fromzl_hw/ops-nn： fix/arg-max-grad-perf into cann/ops-nn： master

perf(arg_max_grad): arch35 轴下标改寄存器生成并开双缓冲, UB 记账由容量反解 #9546

Merged

zl_hwAdd on Aug 31

Merged fromzl_hw/ops-nn fix/arg-max-grad-perfintocann/ops-nn master

zl_hwCollaborators

Aug 31

描述

ArgMaxGrad 在 arch35 上改为不带 D 的原型后，轴下标（A2 融合 pass 里的 assist 张量）由内核自生成。原实现每段都要把轴下标物化到 UB 再读回比较，既占 UB（每元素 4B）又多走一趟写；同时 BUFFER_NUM 固定为 1，MTE2→V→MTE3 全串行。本 PR 按「腾地方 + 用地方」两步优化，并把 UB 记账从「先定尺寸再回头检查」改为「按容量反解」。

改动方法

多行合并按行长分四档分派（arg_max_grad_nd.h）
- MERGE_COMPACT_DIRECT 逐行直算：行长是 32B 整数倍且不短于「车道数/铺设遍数」时，轴下标退化成该行的标量常量留在寄存器，indices/updates 只驻留一行并原地重复读；
- MERGE_COMPACT_FILL 紧排铺 tile：行长对齐但过短时改回 UB→UB 倍增铺设；
- MERGE_PAD 按行补齐：行长非 32B 整数倍且长于一个寄存器时每行一个 burst；
- MERGE_COMPACT_STREAM 判据已实现但以常量 false 关闭，原因见代码注释（非对齐流式写会覆盖 32B 块内前导车道）。
轴下标改为寄存器内生成（arg_max_grad_vf.h）
VF 入口新增 AssistSrc 模板档：ARANGE（inner==1，沿被选轴连续，一条 Reg::Arange）/ SCALAR（inner>1 单行段，一条 Reg::Duplicate）/ UB（仅剩「一个寄存器块跨多行」的紧排与补齐两档，k 在寄存器内算不出闭式）。ComputeSeg、ComputeOuters、ComputeRows 三条路径不再调 GenAssist。
UB 记账由容量反解（arg_max_grad_tiling_arch35.cpp）
新增 UbLayout，六个字段与内核 InitBuffer 的六次调用一一对应，Total() 即一段的真实 UB 需求；SolveColsPerChunk 对其单调性二分，求满足 Total() <= ubSize 的最大整宽段长。i32BufBytes 按档记账：inner==1 或逐行直算档记 0，每元素少 4B。
开启双缓冲：BUFFER_NUM 1 → 2。上一条腾出的 UB 正好抵消队列翻倍（fp16 每元素 ~8.1B → ~4.1B → ~8.2B，UB 占用基本持平），换来 MTE2 搬入与 V 计算重叠。
Process() 拆成 ProcessOuterSeg / ProcessRowSeg（纯重构，语义不变）；host UT 补 10 类拒收分支用例。

性能实测（Ascend950PR，`ttk -d false -b release`，`--run 12`）

两步收益必须分开说明，不能合并成一句「优化」：

阶段	大形状 6 例	小形状 8 例
仅第 2、3 条（轴下标寄存器化）	几何均值 0.98x（无收益）	0.99x
叠加第 4 条（双缓冲）vs 原始	中位 1.13x，最好 1.35x	中位 1.05x，最好 1.42x

收益集中在 inner>1 的三档（DIRECT/FILL/PAD，1.11~1.42x）；inner==1 基本持平（0.98~0.99x，一进一出没有可重叠的流）。
轴下标寄存器化单独看没有收益，它的作用是腾出 UB 让双缓冲开得起；两步缺一不可。
已知退化 1 例：(2,128,129) 0.88x（两次测一致）。段数不足 2 时双缓冲多付一次同步，后续可按段数决定 buffer 深度收回。

关联的Issue

关联 Issue #5418

https://gitcode.com/cann/ops-nn/issues/5418

测试

Ascend950PR 真机，干净包（无插桩符号），kernel .o 与 op_tiling.so 同步部署；每份结果条数与用例数逐份核对一致。

集合	例数	结果
泛化（R7 派生量全值域集）	578	578/578 PASS
泛化（交付主集）	1224	1224/1224 PASS
GE 图通路	56	56/56 PASS
三方精度	100	100/100 PASS
A/B 对照	42	42/42 PASS
DFX（7 例预期拒收 + 2 例行为记录）	9	与改造前基线逐例一致
host UT（`--soc=ascend950`）	28	28/28 PASS
性能抽检（精度/越界列）	14	14/14 PASS

真机功能用例合计 1839 例全过。

文档更新

无。

类型标签

Bug修复
新特性
性能优化
文档更新
其他，请描述：

AI/Agent生成声明

AI辅助编写

Pull Request has been successfully merged by CANN-robot.

（Thank you for zl_hw's contribution.）

Z zl_hwCollaborators

Aug 31 Created pull request, commit 42a596f7

atomgit-bot

Aug 31 Comments：

变更摘要

本 PR 对 arch35 上 ArgMaxGrad 算子的多行合并路径做性能优化：将「多行合并」按行长形态分成四档（按行补齐、紧排逐行直算、紧排铺 tile、紧排流式铺设），并在「行长是 32B 整数倍」的档位上采用逐行直算，不再物化轴下标、不再把 indices/updates 复制成多行，从而去掉对整段操作数的冗余 UB 读写。同时把行跨度、行数、合并模式等布局参数由 host 一次算准写入 tiling 数据，内核直接取用、不再自行推导，规避了 inner 非 8 整数倍时 assist 逐行写导致的 UB 地址未对齐问题（errcode 340）。arg_max_grad_proto.h 另为 REG_OP(ArgMaxGrad) 增加了 OPS_PROTO_DEF_ARGMAXGRAD 宏保护，防止重复定义。

主要改动

多行合并四档分派: host 侧新增 MERGE_PAD、MERGE_COMPACT_DIRECT、MERGE_COMPACT_FILL、MERGE_COMPACT_STREAM 四种合并形态常量及 FILL_PASSES 代价阈值判据（rowDirectShape），按行长与车道数关系选出档位并写入 tilingData_->mergeMode，内核 Process 据此分别走 ComputeRowsDirect / ComputePacked / ComputePad 路径。
逐行直算去掉操作数复制: 新增 ComputeRowsDirect 与 VF 侧 ArgMaxGradSelectRowsVF——行长 32B 对齐时轴下标以标量常量留在寄存器（Duplicate 不进 UB），indices/updates 每行从同一份原地重复读，相比「先复制再整段选择」省掉三遍对整段数据的 UB 读写；同时 UB 记账将该档的 indices/updates 从每元素占用挪成与段长无关的 fixedBytes 固定块。
布局参数由 host 统一下发: arg_max_grad_tiling_data.h 新增 rowsPerChunk、rowElems、dstStrideT、dstStrideI、mergeMode 及分列的 idxBufBytes/updBufBytes 字段，host 一次算准后内核 Init 直接取用，不再按 inner_ 自行推导行跨度（修复非对齐写入问题）；inner==1 时还修正了段长上限口径（按整张量 totalElems 计算）。
inner==1 多 outer 合并路径: 新增 ComputeOuters，一段内连续搬入 m 个 outer 的 var 与各自标量 indices/updates，轴下标 0..D-1 只生成一次，逐 outer 一条 VF，摊薄每段的 MTE2→V→MTE3 固定开销。
assist 倍增铺 tile: GenAssist 多行填充由逐行 Duplicate 改为「第 0 行铺底 + log2 倍增复制 + 一次整体加常量」，指令数由 O(rows) 降为 O(log2 rows)；非对齐场景新增 TileAppendAddsUnalignVF，把复制与加常量合成一条非对齐流式写。

Incorrect?

atomgit-bot

Aug 31 Comments：

代码审查

关闭总结

各变更文件审查结论

index/arg_max_grad/op_graph/arg_max_grad_proto.h — 无问题。OPS_PROTO_DEF_ARGMAXGRAD 防重注册守卫是标准做法, 不影响现有 include 守卫, 未发现缺陷。
index/arg_max_grad/op_host/arch35/arg_max_grad_tiling_arch35.cpp — 发现 P1: inner==1 场景新增的 idx/upd 缓冲未纳入 budget 预留, 导致沿短轴(如 D≤8, int32)做 argmax 的大张量 tiling 返回 GRAPH_FAILED, 属对旧行为的功能回归; 且 rowsPerChunk==1(多 outer 路径未启用)时仍按 maxOuters≈colsPerChunk 分配未用缓冲, 加剧该问题。
index/arg_max_grad/op_kernel/arch35/arg_max_grad_nd.h — 无独立缺陷。新增 CopyInRowsPad/ComputeRowsDirect/ComputeOuters 及其在 Process 中的四档分派, 经核对 buffer 记账与内核 InitBuffer 一一对应、行布局/对齐条件与 host 下发参数一致、GenAssist 倍增铺设无重叠且偏移对齐、MTE2→S 同步覆盖标量读取; 非对齐流式写分支当前不可达(mergeMode 恒不取 COMPACT_STREAM), 属防御性代码。另附一条 P3 低置信耦合性提示。
index/arg_max_grad/op_kernel/arch35/arg_max_grad_tiling_data.h — 发现 P3: idxBufBytes/updBufBytes 插入结构体中段, 使 maskBufBytes/selBufBytes 历史偏移后移, 违反"新字段只能追加在末尾"的规则, 存在跨版本 tiling 数据 ABI 隐患。
index/arg_max_grad/op_kernel/arch35/arg_max_grad_vf.h — 无问题。新增 ArgMaxGradSelectRowsVF/TileAppendAddsUnalignVF 与既有已验证的 VF 范式(车道推进、UpdateMask 收窄、B16 拆打包、地址按各自类型元素推进)完全同构, 未发现正确性问题。

统计

P1: 1 条(P0: 0)
P3: 2 条(其中 1 条为低置信耦合性提示)
无 P0/P2。

总体风险判断

中高风险。主要风险是 host tiling 预算计算的回归: 对 inner==1 且被选轴较短(≤8)的大张量(如 [N,8] 沿末轴 argmax), 新引入的 idx/upd 缓冲会使 ubUsed 超限、tiling 直接失败, 影响常见形状的可用性; 现有 UT 形状(inner 均>1)无法覆盖该回归, 建议修复后补充对应用例。内核侧新路径的搬运/布局/同步逻辑经核对自洽, 未发现独立正确性缺陷; tiling data 字段顺序问题为 ABI 约定层面的低危隐患。

类型	数量
🔴 阻塞	1
🟡 建议	0

⛔ 需要修改

Incorrect?

CANN-robotCollaborators

Aug 31 Added label: cann-cla/yes

CANN-robotCollaborators

Aug 31 Comments：

Thanks for your pull-request.
The full list of commands accepted by me can be found at here.
You can get sig-info at here.
You can self-configure the PR merge rules for this repository. For more details, please refer to Here.
For more, you also can visit HICANN.

PR Approval Progress

✅ Congratulations! All modules have met the lgtm and approve requirements.

Module Approval Details

module	lgtm status	approve status
//op_graph/*_proto.h	✅ 王永光, 汤平川 (2/2)	✅ 王永光 (1/1)
index	✅ 陈风, 汤平川 (2/2)	✅ 汤平川, 陈风 (2/1)

💡 Tip:

Committer can comment /approve or /lgtm

Commenting /approve implies both code review (lgtm) and intent to merge (approve)

CLA Signature Pass

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

CANN-robotCollaborators

Aug 31 Set zengjuan,chenfeng61,chaotang233,kevin_huang1234,FelixTang7,crystalhu,yangyang016,fanqirui,wangrui_,su-yueming,tang-lei01,chenqi317,wang-xing001,liubo75,tangweiwei2,pingchuantang,qianzehong,wangyongguang as approver

CANN-robotCollaborators

Aug 31 Set chenfeng61,kevin_huang1234,FelixTang7,wangrui_,su-yueming,tang-lei01,wang-xing001,pingchuantang,qianzehong,wangyongguang as approver

atomgit-botAug 31Start code review2

index/arg_max_grad/op_host/arch35/arg_max_grad_tiling_arch35.cpp

Outdated

282+ int64_t d = tilingData_->dimSize > 0 ? tilingData_->dimSize : 1;

283+ int64_t maxOuters = colsPerChunk / d + 1;

284+ tilingData_->idxBufBytes = Ops::Base::CeilAlign(maxOuters * DTYPE_LEN_INT32, vlInt32_ * DTYPE_LEN_INT32);

285+ tilingData_->updBufBytes = Ops::Base::CeilAlign(maxOuters * varDtypeLen_, vlMax_ * varDtypeLen_);

atomgit-botAug 31Comments：

🟠 High Priority

changed line → affected behavior → failure mode → suggested fix。

本 PR 在内核里对 inner==1 也新增了 idxQue_/updQue_ 两块缓冲(Init 中无条件 InitBuffer), host 侧在 280-285 行按 maxOuters = colsPerChunk/d + 1 分配 idxBufBytes/updBufBytes, 并在 343 行把这两块计入 ubUsed。但 257 行 budget = (ubSize - fixedBytes) / bytesPerPoint 在 inner==1 时 bytesPerPoint=fullPoint=basePoint(不含 idx/upd, 见 240 行 innerIsOne ? 0 : ...), 也就是说 colsPerChunk 仍按"只用满整块 UB"计算, 完全没有为新增的 idx/upd 缓冲预留空间。

对 inner==1 且 dimSize 较小的形状, colsPerChunk 会逼近 UB 上限, 此时 idxBuf+updBuf 约 colsPerChunk*(4+varDtypeLen) 字节, 344 行 ubUsed > ubSize_ 检查必然触发, CalUbSplit 返回 GRAPH_FAILED, 算子整体执行失败。具体算例(int32、arch35 UB 192KB-32KB=163840 字节、vlMax_=64): 形状 [N,8] dimension=1(outer=N, D=8, inner=1, N 较大时 colsPerChunk=12544): tBuf=50176, i32Buf=50176, mask=1568, idxBuf=6272, updBuf=6272 → ubUsed=165632 > 163840 → 失败; 而旧代码(不为 inner==1 分配 idx/upd)ubUsed=152096 < 163840 可正常运行。D=1/2/4/8(int32)、D=1/2/4(fp16) 同理失败(d 更大如 ≥16 时因为 idx/upd 变小才通过)。这是对"沿长度为 ≤8 的末轴做 argmax 的二维大张量"这类常见形状的功能回归, 现有 UT 形状({2,3,4}/{2,3,4,5}/{64,128,32}, inner 均>1)无法覆盖。注意即使 dAligned=false、rowsPerChunk_=1(多 outer 路径根本不启用、内核里 idx/upd 缓冲完全用不上)时, host 仍按 maxOuters≈colsPerChunk 分配, 同样导致失败, 属于双重问题: 既未预留预算, 又对不启用的路径过量分配。

建议：在计算 budget/colsPerChunk 前, 把 inner==1 场景需要新增的 idx/upd 缓冲字节数先计入 fixedBytes(与 rowDirect 分支同法), 使 colsPerChunk 按剩余空间推算; 同时仅在 rowsPerChunk>1(多 outer 路径启用)时按 maxOuters 分配 idxBufBytes/updBufBytes, rowsPerChunk==1 时分配最小块, 避免对未启用的路径过量分配导致 ubUsed 超限。

Incorrect?

System MessageSystem

Aug 31 Comments：

changed this line on dfcba76f view diff detail

Z zl_hwCollaborators

Aug 31 Pre merge succeeded (commit_id: efec791ffb4d108abd2024e73580fafb7473acd1)

10 event messages are collapsed here。 See more

CANN-robotCollaborators

Aug 31 Set chenxingyu18 as approver

zl_hwCollaborators

Aug 31 Comments：

compile

Z zl_hwCollaborators

Aug 31 Pre merge succeeded (commit_id: 8c553b34b68ef3fcd18358cbf023716d86fad66a)

CANN-robotCollaborators

Aug 31 Added label: ci-pipeline-running

CANN-robotCollaborators

Aug 31 Removed label: ci-pipeline-running

CANN-robotCollaborators

Aug 31 Added label: ci-pipeline-passed

TangPCCollaborators

Aug 31 Change title to "perf(arg_max_grad): arch35 轴下标改寄存器生成并开双缓冲, UB 记账由容量反解", originally titled "perf(arg_max_grad): arch35 多行合并按行长分四档分派, 去掉不必要的操作数复制".

TangPCCollaborators

Aug 31 Modify description of pull request

yuning_chenCollaborators

Aug 31 Comments：

您好，pr合入已准备就绪，请尽快联系committer进行检视，谢谢！

TangPCCollaborators

Sep 1 Comments：

/lgtm
/approve

wangyongguangCollaborators

Sep 1 Comments：

/approve

CANN-robotCollaborators

Sep 1 Added label: approved

chenfengCollaborators

Sep 1 Comments：

/lgtm
/approve

CANN-robotCollaborators

Sep 1 Added label: lgtm

CANN-robotCollaborators

Sep 1 Comments：

The MR can not be merged, because of CodeReview discussion not resolved

If you want to solve this problem, you can click here to do it in the FAQs.

Z zl_hwCollaborators

Sep 1 Resolved final issue

CANN-robotCollaborators

Sep 1 Closed associated issue

CANN-robotCollaborators

Sep 1 Merged pull request

CANN-robotCollaborators

Sep 1 Comments：

Pull Request 已合并或已关闭。

If you want to solve this problem, you can click here to do it in the FAQs.

描述

改动方法

性能实测（Ascend950PR，ttk -d false -b release，--run 12）

关联的Issue

测试

文档更新

类型标签

AI/Agent生成声明

Pull Request has been successfully merged by CANN-robot.

变更摘要

代码审查

关闭总结

各变更文件审查结论

统计

总体风险判断

PR Approval Progress

Module Approval Details

CLA Signature Pass

性能实测（Ascend950PR，`ttk -d false -b release`，`--run 12`）