Star1393
2139
代码
代码
Issues846
Pull Requests694
Actions
项目讨论
Wiki
分析
Star1393
2139
已合并
perf(arg_max_grad): arch35 轴下标改寄存器生成并开双缓冲, UB 记账由容量反解 #9546
zl_hw创建于 8月31日
perf(arg_max_grad): arch35 轴下标改寄存器生成并开双缓冲, UB 记账由容量反解 #9546
已合并
zl_hw创建于 8月31日
从zl_hw/ops-nnfix/arg-max-grad-perf合入到cann/ops-nnmaster
共 6 个文件变更+551-152
Mindex/arg_max_grad/op_graph/arg_max_grad_proto.h+3-0
@@ -46,6 +46,8 @@ namespace ge {
46* @attention Constraints:46* @attention Constraints:
47* The operator will not be enhanced in the future.47* The operator will not be enhanced in the future.
48*/48*/
49+#ifndef OPS_PROTO_DEF_ARGMAXGRAD
50+#define OPS_PROTO_DEF_ARGMAXGRAD
49REG_OP(ArgMaxGrad)51REG_OP(ArgMaxGrad)
50 .INPUT(var, TensorType({DT_FLOAT, DT_FLOAT16, DT_INT32, DT_INT8}))52 .INPUT(var, TensorType({DT_FLOAT, DT_FLOAT16, DT_INT32, DT_INT8}))
51 .INPUT(indices, TensorType({DT_INT32}))53 .INPUT(indices, TensorType({DT_INT32}))
@@ -53,6 +55,7 @@ REG_OP(ArgMaxGrad)
53 .OUTPUT(y, TensorType({DT_FLOAT, DT_FLOAT16, DT_INT32, DT_INT8}))55 .OUTPUT(y, TensorType({DT_FLOAT, DT_FLOAT16, DT_INT32, DT_INT8}))
54 .REQUIRED_ATTR(dimension, Int)56 .REQUIRED_ATTR(dimension, Int)
55 .OP_END_FACTORY_REG(ArgMaxGrad)57 .OP_END_FACTORY_REG(ArgMaxGrad)
58+#endif // OPS_PROTO_DEF_ARGMAXGRAD
56} // namespace ge59} // namespace ge
57 60 
58#endif // OPS_INDEX_ARG_MAX_GRAD_PROTO_H_61#endif // OPS_INDEX_ARG_MAX_GRAD_PROTO_H_
Mindex/arg_max_grad/op_host/arch35/arg_max_grad_tiling_arch35.cpp+180-38
@@ -23,6 +23,10 @@
23#include "../op_kernel/arch35/arg_max_grad_tiling_key.h"23#include "../op_kernel/arch35/arg_max_grad_tiling_key.h"
24 24 
25namespace optiling {25namespace optiling {
26+// 说明: 本文件里 `OPS_CHECK_NULL_WITH_CONTEXT` 展开出的拒收分支均**不可达**(输入描述/形状/属性
27+// 由 GE 框架保证非空, faker 也造不出), 保留为防御。真正可达的拒收分支都有对应 UT:
28+// dtype 非法 / updates dtype 不一致 / indices 非 int32 / updates 与 indices 形状不一致 / rank 为 0 /
29+// dimension 越界 / indices 元素数不匹配 / 平台核数为 0 / UB 不大于 dcache 预留 / UB 装不下一个向量整宽。
26using namespace Ops::Base;30using namespace Ops::Base;
27 31 
28namespace {32namespace {
@@ -43,12 +47,42 @@ constexpr int64_t INT8_SEL_BUF_NUM = 3;
43// 掩码 1 bit/元素, 记账按 1 字节保守预留47// 掩码 1 bit/元素, 记账按 1 字节保守预留
44constexpr int64_t MASK_BYTES_PER_POINT = 1;48constexpr int64_t MASK_BYTES_PER_POINT = 1;
45// 与 op_kernel/arch35/arg_max_grad_nd.h 的 BUFFER_NUM 一致49// 与 op_kernel/arch35/arg_max_grad_nd.h 的 BUFFER_NUM 一致
46-constexpr int64_t BUFFER_NUM = 1;50+constexpr int64_t BUFFER_NUM = 2;
47// UB / GM 的搬运块大小: DataCopyPad 写不足一块时按块读-改-写51// UB / GM 的搬运块大小: DataCopyPad 写不足一块时按块读-改-写
48constexpr int64_t BITS_PER_BYTE = 8;52constexpr int64_t BITS_PER_BYTE = 8;
49constexpr int64_t UB_BLOCK_BYTES = 32;53constexpr int64_t UB_BLOCK_BYTES = 32;
54+// 多行合并的四种形态, 与 op_kernel 的同名常量一一对应:
55+// PAD 行在 UB 里补齐到 32B 边界, 每行一个 burst
56+// COMPACT_DIRECT紧排 + 逐行直算(轴下标是标量常量留在寄存器, indices/updates 原地重复读, 不复制操作数)
57+// COMPACT_FILL 紧排 + 先把操作数铺成多行再整段选择(UB→UB 拷贝)
58+// COMPACT_STREAM紧排 + 铺法改用非对齐流式写(行长不是 32B 整数倍时行起点落不到寄存器块边界)
59+constexpr int64_t MERGE_PAD = 0;
60+constexpr int64_t MERGE_COMPACT_DIRECT = 1;
61+constexpr int64_t MERGE_COMPACT_FILL = 2;
62+constexpr int64_t MERGE_COMPACT_STREAM = 3;
63+// 铺 tile 路径要对整段数据走的遍数: indices 铺 + updates 铺 + 轴下标铺 + 比较选择。
64+// 逐行直算每元素代价 ∝ 1/min(inner, 车道数)(一行一组指令, 车道用不满按比例摊薄),
65+// 铺 tile 每元素代价 ∝ FILL_PASSES/车道数, 两者相等处即 min(inner, 车道数) = 车道数 / FILL_PASSES。
66+constexpr int64_t FILL_PASSES = 4;
50// SIMD/SIMT 共用的 dcache, 与参照算子一致预留67// SIMD/SIMT 共用的 dcache, 与参照算子一致预留
51constexpr uint64_t SIMD_SIMT_DCACHE_SIZE = static_cast<uint64_t>(32 * 1024);68constexpr uint64_t SIMD_SIMT_DCACHE_SIZE = static_cast<uint64_t>(32 * 1024);
69+ 
70+// 一段(colsPerChunk 个元素)在 UB 上的完整占用。字段与内核 InitBuffer 的调用一一对应,
71+// Total() 即该段的真实 UB 需求 —— 段长由它反解出来, 不再"先定尺寸再回头检查是否超"。
72+struct UbLayout {
73+ int64_t tBufBytes = 0;
74+ int64_t i32BufBytes = 0;
75+ int64_t idxBufBytes = 0;
76+ int64_t updBufBytes = 0;
77+ int64_t maskBufBytes = 0;
78+ int64_t selBufBytes = 0;
79+ 
80+ int64_t Total() const
81+ {
82+ return BUFFER_NUM * (2 * tBufBytes) + i32BufBytes + maskBufBytes + selBufBytes +
83+ BUFFER_NUM * (idxBufBytes + updBufBytes);
84+ }
85+};
52} // namespace86} // namespace
53 87 
54class ArgMaxGradTiling {88class ArgMaxGradTiling {
@@ -63,6 +97,9 @@ private:
63 ge::graphStatus CheckShape();97 ge::graphStatus CheckShape();
64 ge::graphStatus CheckAttrAndSplitLayout();98 ge::graphStatus CheckAttrAndSplitLayout();
65 ge::graphStatus CalUbSplit();99 ge::graphStatus CalUbSplit();
100+ UbLayout MakeUbLayout(int64_t cols, bool rowDirect) const;
101+ int64_t SolveColsPerChunk(bool rowDirect, int64_t upperCols) const;
102+ void CalMergeParams(int64_t cols);
66 void CalCoreSplit();103 void CalCoreSplit();
67 void PrintTilingData();104 void PrintTilingData();
68 105 
@@ -204,53 +241,155 @@ void ArgMaxGradTiling::CalCoreSplit()
204 realCoreNum_ = Ops::Base::CeilDiv(total, perCore);241 realCoreNum_ = Ops::Base::CeilDiv(total, perCore);
205}242}
206 243 
244+// 一段(cols 个元素)在 UB 上的真实占用: 每个字段对应内核 InitBuffer 的一次调用。
245+UbLayout ArgMaxGradTiling::MakeUbLayout(int64_t cols, bool rowDirect) const
246+{
247+ const int64_t innerElems = tilingData_->inner > 0 ? tilingData_->inner : 1;
248+ UbLayout layout;
249+ layout.tBufBytes = cols * varDtypeLen_;
250+ // 轴下标(assist)只有"一个寄存器块跨多行"的两档才需要物化到 UB(紧排铺 tile / 按行补齐);
251+ // inner==1 与逐行直算档的 k 由内核在寄存器内生成(Reg::Arange / Duplicate), 这块整块省掉,
252+ // 每元素少 4B —— fp16 形状的段长因此接近翻倍。
253+ const bool assistInUb = (tilingData_->inner != 1) && !rowDirect;
254+ layout.i32BufBytes = assistInUb ? cols * DTYPE_LEN_INT32 : 0;
255+ if (rowDirect) {
256+ // 逐行直算: indices/updates 只驻留一行, 是与段长无关的固定块
257+ layout.idxBufBytes = Ops::Base::CeilAlign(innerElems * DTYPE_LEN_INT32, vlInt32_ * DTYPE_LEN_INT32);
258+ layout.updBufBytes = Ops::Base::CeilAlign(innerElems * varDtypeLen_, vlMax_ * varDtypeLen_);
259+ } else if (tilingData_->inner == 1) {
260+ // inner==1: 一段装 m 个 outer, 每个 outer 只带一个标量 indices/updates(m 的上界与 rowsPerChunk 同口径)
261+ const int64_t d = tilingData_->dimSize > 0 ? tilingData_->dimSize : 1;
262+ const int64_t maxOuters = cols / d + 1;
263+ layout.idxBufBytes = Ops::Base::CeilAlign(maxOuters * DTYPE_LEN_INT32, vlInt32_ * DTYPE_LEN_INT32);
264+ layout.updBufBytes = Ops::Base::CeilAlign(maxOuters * varDtypeLen_, vlMax_ * varDtypeLen_);
265+ } else {
266+ layout.idxBufBytes = cols * DTYPE_LEN_INT32;
267+ layout.updBufBytes = cols * varDtypeLen_;
268+ }
269+ layout.maskBufBytes = Ops::Base::CeilAlign(cols / BITS_PER_BYTE, UB_BLOCK_BYTES);
270+ layout.selBufBytes = isInt8_ ? (INT8_SEL_BUF_NUM * cols * DTYPE_LEN_HALF) : 0;
271+ return layout;
272+}
273+ 
274+// 由 UB 容量反解段长: Total() 对 cols 单调不减(各块要么正比于 cols, 要么与 cols 无关),
275+// 故可二分取满足 Total() <= ubSize_ 的最大整宽段长。段长必须是 vlMax_ 的整数倍, 否则后续
276+// buffer 的起点会偏离向量寄存器整宽边界(errcode 340)。返回 0 表示一个整宽都放不下。
277+int64_t ArgMaxGradTiling::SolveColsPerChunk(bool rowDirect, int64_t upperCols) const
278+{
279+ int64_t lo = 0; // 单位: vlMax_ 个元素
280+ int64_t hi = upperCols / vlMax_;
281+ while (lo < hi) {
282+ const int64_t mid = lo + (hi - lo + 1) / 2;
283+ if (MakeUbLayout(mid * vlMax_, rowDirect).Total() <= static_cast<int64_t>(ubSize_)) {
284+ lo = mid;
285+ } else {
286+ hi = mid - 1;
287+ }
288+ }
289+ return lo * vlMax_;
290+}
291+ 
292+// ── 多行合并的 UB 布局参数: host 一次算准, 内核直接取用 ──────────────────────────────
293+void ArgMaxGradTiling::CalMergeParams(int64_t cols)
294+{
295+ int64_t alignElems = UB_BLOCK_BYTES / DTYPE_LEN_INT32;
296+ if (UB_BLOCK_BYTES / varDtypeLen_ > alignElems) {
297+ alignElems = UB_BLOCK_BYTES / varDtypeLen_;
298+ }
299+ // 行跨度取 T 域与 int32 域的公共 32B 对齐粒度, 使同一 lane 在两个域里指向同一个元素;
300+ // 由此 assist 逐行写的目的地址 r*rowElems*4 必然落在 32B 边界(内核侧曾自行按 inner 推导,
301+ // inner 非 8 整数倍时触发 aivec errcode 340: VEC 访问 UB 地址未对齐)。
302+ if (tilingData_->inner != 1) {
303+ const int64_t inner = tilingData_->inner > 0 ? tilingData_->inner : 1;
304+ const int64_t rowElems = Ops::Base::CeilAlign(inner, alignElems);
305+ const int64_t occT = Ops::Base::CeilAlign(inner * varDtypeLen_, UB_BLOCK_BYTES);
306+ const int64_t occI = Ops::Base::CeilAlign(inner * DTYPE_LEN_INT32, UB_BLOCK_BYTES);
307+ tilingData_->rowElems = rowElems;
308+ tilingData_->dstStrideT = (rowElems * varDtypeLen_ - occT) / UB_BLOCK_BYTES;
309+ tilingData_->dstStrideI = (rowElems * DTYPE_LEN_INT32 - occI) / UB_BLOCK_BYTES;
310+ if (rowElems == inner) {
311+ // 行长是 32B 整数倍 → 行起点落在寄存器块边界, 可以逐行直算; 但一行短于 车道数/FILL_PASSES 时
312+ // 直算的车道浪费超过铺 tile 多走的几遍, 这一档仍走铺 tile(相等处判给铺 tile: 它每条指令都是满车道)。
313+ tilingData_->mergeMode = (inner * FILL_PASSES > vlInt32_) ? MERGE_COMPACT_DIRECT : MERGE_COMPACT_FILL;
314+ // ⚠️ 紧排 + 非对齐流式铺设当前不启用: 该铺法写入目的地址所在 32B 块时会连带覆盖块内前导
315+ // 车道。indices/updates 各行内容相同, 覆盖不可见; 轴下标每行取值不同, 会在行边界处产生少量
316+ // 错值(实测 binary_equal 下每例错 1~3 个元素, 集中在 inner 小且 D 不大的形状)。要启用需先把
317+ // 轴下标改成"对齐单元倍增 + 头部算术构造", 使所有写入偏移都是向量整宽的整数倍。判据本身成立。
318+ } else if (false && inner <= vlInt32_) {
319+ // 两条路径的代价都是"每行若干条指令", 按每行指令数选:
320+ // 按行补齐 = 每行 1 个 DMA burst(与 inner 无关);
321+ // 紧排 = 每行 ceil(inner / 车道数) 条向量指令铺 tile, int32 域车道最窄故为绑定项。
322+ // inner 不超过一个向量寄存器的 int32 车道数时, 紧排每行只需一条指令, 不多于补齐,
323+ // 且省掉把整段拆成 rows 个 burst; 超过则补齐更省。车道数由平台 VRegSize 得出, 跨芯片自适应。
324+ tilingData_->mergeMode = MERGE_COMPACT_STREAM;
325+ } else {
326+ tilingData_->mergeMode = MERGE_PAD;
327+ }
328+ // 紧排时行在 UB 里不留空隙, 合并行数按 inner 算; 补齐时按补齐后的行跨度算。
329+ const int64_t stride = (tilingData_->mergeMode == MERGE_PAD) ? rowElems : inner;
330+ tilingData_->rowsPerChunk = (stride > 0 && stride <= cols) ? (cols / stride) : 1;
331+ } else {
332+ // inner==1: 一段能装几个 outer。每个 outer 在 UB 里的起点是 j*D, 要求 D 在 T 域与 int32 域
333+ // 都落在 32B 边界, 否则逐 outer 的向量读写会落到非对齐地址上(errcode 340)。
334+ const int64_t d = tilingData_->dimSize > 0 ? tilingData_->dimSize : 1;
335+ const bool dAligned = ((d * varDtypeLen_) % UB_BLOCK_BYTES == 0) &&
336+ ((d * DTYPE_LEN_INT32) % UB_BLOCK_BYTES == 0);
337+ tilingData_->rowElems = 1;
338+ tilingData_->rowsPerChunk = (dAligned && d <= cols) ? (cols / d) : 1;
339+ tilingData_->dstStrideT = 0;
340+ tilingData_->dstStrideI = 0;
341+ tilingData_->mergeMode = MERGE_COMPACT_DIRECT;
342+ }
343+}
344+ 
207ge::graphStatus ArgMaxGradTiling::CalUbSplit()345ge::graphStatus ArgMaxGradTiling::CalUbSplit()
208{346{
209- bool innerIsOne = (tilingData_->inner == 1);347+ const bool innerIsOne = (tilingData_->inner == 1);
210- // 每个元素的 UB 占用: var + 轴下标(int32, 自生成) + out + 掩码; inner>1 时另加 indices(int32) 与 updates348+ int64_t alignElems = UB_BLOCK_BYTES / DTYPE_LEN_INT32;
211- // 记账份数与内核的 BUFFER_NUM 一致(当前 1, 见本文件顶部常量与 op_kernel 同名常量)349+ if (UB_BLOCK_BYTES / varDtypeLen_ > alignElems) {
212- int64_t bytesPerPoint = BUFFER_NUM * (2 * varDtypeLen_ + DTYPE_LEN_INT32) + MASK_BYTES_PER_POINT;350+ alignElems = UB_BLOCK_BYTES / varDtypeLen_;
213- if (!innerIsOne) {
214- bytesPerPoint += BUFFER_NUM * (DTYPE_LEN_INT32 + varDtypeLen_);
215- }
216- if (isInt8_) {
217- bytesPerPoint += INT8_SEL_BUF_NUM * DTYPE_LEN_HALF;
218 }351 }
352+ const int64_t innerElems = tilingData_->inner > 0 ? tilingData_->inner : 1;
353+ // 逐行直算档: 行长是整宽的整数倍(行起点落在寄存器块边界)且一行够肥, 此时 indices/updates
354+ // 只需一行常驻, 不随段长增长。
355+ const bool rowDirectShape = !innerIsOne && (Ops::Base::CeilAlign(innerElems, alignElems) == innerElems) &&
356+ (innerElems * FILL_PASSES > vlInt32_);
219 357 
220- OP_CHECK_IF((static_cast<int64_t>(ubSize_) <= bytesPerPoint * vlMax_),358+ // 单次驻留 UB 的上限 = 一个 outer 覆盖的元素数(再多也用不上, indices/updates 按 outer 变):
359+ // inner == 1 : 一段可以装多个 outer(每个 outer 只带一个标量 indices/updates), 上限按整张量算;
360+ // inner > 1 : indices/updates 是整行数据, 一段不能跨 outer, 上限是 D×inner —— 【不能】按
361+ // "一行(inner)"设上限: 实测 inner=16 时按行设限会把段长钳到 64 个元素, 每段只够
362+ // 并 4 行, 搬运与同步的固定开销吃掉全部收益(合并前后 G/N 中位仅 1.04x)。
363+ const int64_t spanElems = innerIsOne ? tilingData_->totalElems : tilingData_->dimSize * tilingData_->inner;
364+ const int64_t upperCols = Ops::Base::CeilAlign(spanElems > 0 ? spanElems : vlMax_, vlMax_);
365+ 
366+ // 先按"操作数整段驻留"的保守口径反解一次: 只有当一行本身装得进这个段长时, 才谈得上逐行直算。
367+ const int64_t colsProbe = SolveColsPerChunk(false, upperCols);
368+ OP_CHECK_IF((colsProbe <= 0),
221 OP_LOGE(context_->GetNodeName(), "ub size %lu is too small for one vector of elements", ubSize_),369 OP_LOGE(context_->GetNodeName(), "ub size %lu is too small for one vector of elements", ubSize_),
222 return ge::GRAPH_FAILED);370 return ge::GRAPH_FAILED);
223- 371+ bool rowDirect = false;
224- int64_t budget = static_cast<int64_t>(ubSize_) / bytesPerPoint;372+ int64_t colsPerChunk = colsProbe;
225- int64_t colsPerChunk = budget / vlMax_ * vlMax_; // 向下取整宽对齐(见 vlMax_ 说明)373+ if (rowDirectShape && innerElems <= colsProbe) {
226- // 单次驻留 UB 的上限 = 一个 outer 覆盖的元素数(再多也用不上, indices/updates 按 outer 变):374+ const int64_t directCols = SolveColsPerChunk(true, upperCols);
227- // inner == 1 : 一个 outer 就是 D 个元素;375+ if (directCols > 0) {
228- // inner > 1 : 一个 outer 是 D×inner 个元素 —— 核内会把同一个 outer 的多行并成一段处理,376+ rowDirect = true;
229- // 所以这里【不能】按"一行(inner)"设上限。实测 inner=16 时按行设限会把377+ colsPerChunk = directCols;
230- // colsPerChunk 钳到 64 个元素, 每段只够并 4 行, 搬运与同步的固定开销378+ }
231- // 吃掉全部收益(合并前后 G/N 中位仅 1.04x)。
232- int64_t spanElems = innerIsOne ? tilingData_->dimSize : tilingData_->dimSize * tilingData_->inner;
233- int64_t alignSpan = Ops::Base::CeilAlign(spanElems > 0 ? spanElems : vlMax_, vlMax_);
234- if (colsPerChunk >= alignSpan) {
235- colsPerChunk = alignSpan; // 一个 outer 全载, 再大也用不上
236 }379 }
237 tilingData_->colsPerChunk = colsPerChunk;380 tilingData_->colsPerChunk = colsPerChunk;
238 381 
239- // 各 buffer 字节数在此一次算准, 内核直接透传给 InitBuffer, 不再做任何对齐/补齐。382+ // 各 buffer 字节数即反解时用的那份布局, 内核直接透传给 InitBuffer, 不再做任何对齐/补齐;
240- tilingData_->tBufBytes = colsPerChunk * varDtypeLen_;383+ // 段长由 Total() <= ubSize_ 反解而来, 故无需再回头校验总量。
241- tilingData_->i32BufBytes = colsPerChunk * DTYPE_LEN_INT32;384+ const UbLayout layout = MakeUbLayout(colsPerChunk, rowDirect);
242- tilingData_->maskBufBytes = Ops::Base::CeilAlign(colsPerChunk / BITS_PER_BYTE, UB_BLOCK_BYTES);385+ tilingData_->tBufBytes = layout.tBufBytes;
243- tilingData_->selBufBytes = isInt8_ ? (INT8_SEL_BUF_NUM * colsPerChunk * DTYPE_LEN_HALF) : 0;386+ tilingData_->i32BufBytes = layout.i32BufBytes;
387+ tilingData_->idxBufBytes = layout.idxBufBytes;
388+ tilingData_->updBufBytes = layout.updBufBytes;
389+ tilingData_->maskBufBytes = layout.maskBufBytes;
390+ tilingData_->selBufBytes = layout.selBufBytes;
244 391 
245- // 总量自检: 记账口径与内核 InitBuffer 的调用一一对应, 超了直接失败而不是让内核去踩 UB。392+ CalMergeParams(colsPerChunk);
246- int64_t ubUsed = BUFFER_NUM * (2 * tilingData_->tBufBytes) + tilingData_->i32BufBytes + tilingData_->maskBufBytes +
247- tilingData_->selBufBytes;
248- if (!innerIsOne) {
249- ubUsed += BUFFER_NUM * (tilingData_->i32BufBytes + tilingData_->tBufBytes);
250- }
251- OP_CHECK_IF((ubUsed > static_cast<int64_t>(ubSize_)),
252- OP_LOGE(context_->GetNodeName(), "ub buffers %ld bytes exceed ub size %lu", ubUsed, ubSize_),
253- return ge::GRAPH_FAILED);
254 return ge::GRAPH_SUCCESS;393 return ge::GRAPH_SUCCESS;
255}394}
256 395 
@@ -282,8 +421,10 @@ ge::graphStatus ArgMaxGradTiling::Init()
282 if (vlMax_ < vlInt32_) {421 if (vlMax_ < vlInt32_) {
283 vlMax_ = vlInt32_;422 vlMax_ = vlInt32_;
284 }423 }
424+ // 不可达(平台常量): GetVRegSize 在 arch35 上是编译期常量 256, 保留仅为防御跨代改动
285 OP_CHECK_IF((vlInt32_ <= 0), OP_LOGE(context_->GetNodeName(), "GetVRegSize failed"), return ge::GRAPH_FAILED);425 OP_CHECK_IF((vlInt32_ <= 0), OP_LOGE(context_->GetNodeName(), "GetVRegSize failed"), return ge::GRAPH_FAILED);
286 426 
427+ // 以下两条不可达(框架保证): tiling data 缓冲由框架按注册大小分配, memset_s 参数固定合法
287 tilingData_ = context_->GetTilingData<ArgMaxGradArch35TilingData>();428 tilingData_ = context_->GetTilingData<ArgMaxGradArch35TilingData>();
288 OP_CHECK_IF((tilingData_ == nullptr), OP_LOGE(context_->GetNodeName(), "get tilingdata ptr failed"),429 OP_CHECK_IF((tilingData_ == nullptr), OP_LOGE(context_->GetNodeName(), "get tilingdata ptr failed"),
289 return ge::GRAPH_FAILED);430 return ge::GRAPH_FAILED);
@@ -326,6 +467,7 @@ ge::graphStatus ArgMaxGradTiling::DoTiling()
326 467 
327static ge::graphStatus Tiling4ArgMaxGrad(gert::TilingContext* context)468static ge::graphStatus Tiling4ArgMaxGrad(gert::TilingContext* context)
328{469{
470+ // 不可达(框架保证 context 非空), 保留为接口层防御
329 if (context == nullptr) {471 if (context == nullptr) {
330 return ge::GRAPH_FAILED;472 return ge::GRAPH_FAILED;
331 }473 }
Mindex/arg_max_grad/op_kernel/arch35/arg_max_grad_nd.h+226-106
@@ -41,6 +41,9 @@ public:
41 __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,41 __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR y,
42 const ArgMaxGradArch35TilingData* tiling, TPipe* pipe);42 const ArgMaxGradArch35TilingData* tiling, TPipe* pipe);
43 __aicore__ inline void Process();43 __aicore__ inline void Process();
44+ // 每次处理一段, 返回本段消费的元素数(0 不可能: 段长至少 1)
45+ __aicore__ inline int64_t ProcessOuterSeg(int64_t g, int64_t o, int64_t k, int64_t rem);
46+ __aicore__ inline int64_t ProcessRowSeg(int64_t g, int64_t o, int64_t inRow, int64_t rem);
44 47 
45private:48private:
46 using SelT = typename SelectTypeTrait<T>::Type;49 using SelT = typename SelectTypeTrait<T>::Type;
@@ -49,12 +52,15 @@ private:
49 __aicore__ inline void CopyInCommon(int64_t varOffset, int64_t len);52 __aicore__ inline void CopyInCommon(int64_t varOffset, int64_t len);
50 __aicore__ inline void CopyInIdxUpd(int64_t idxOffset, int64_t len);53 __aicore__ inline void CopyInIdxUpd(int64_t idxOffset, int64_t len);
51 __aicore__ inline void GenAssist(int64_t kStart, int64_t rows, int64_t rowLen, int64_t rowStride);54 __aicore__ inline void GenAssist(int64_t kStart, int64_t rows, int64_t rowLen, int64_t rowStride);
55+ __aicore__ inline void CopyInRowsPad(int64_t varOffset, int64_t rows);
52 __aicore__ inline void CopyInIdxUpdRepeat(int64_t idxOffset, int64_t rows);56 __aicore__ inline void CopyInIdxUpdRepeat(int64_t idxOffset, int64_t rows);
53 __aicore__ inline void CopyOutRowsPad(int64_t varOffset, int64_t rows);57 __aicore__ inline void CopyOutRowsPad(int64_t varOffset, int64_t rows);
54 __aicore__ inline void ComputePad(int64_t rows);58 __aicore__ inline void ComputePad(int64_t rows);
55 __aicore__ inline void ComputePacked(int64_t rows, int64_t totalLen);59 __aicore__ inline void ComputePacked(int64_t rows, int64_t totalLen);
56- __aicore__ inline void ComputeRows(int64_t rows, int64_t rowLen, int64_t alignRowLen);60+ __aicore__ inline void ComputeRowsDirect(int64_t kStart, int64_t rows);
57- __aicore__ inline void ComputeSeg(int64_t len, int64_t alignLen, int32_t idxValue, T updValue);61+ __aicore__ inline void ComputeOuters(int64_t outers);
62+ __aicore__ inline void ComputeRows(int64_t rows, int64_t rowLen, int64_t alignRowLen, int64_t kStart);
63+ __aicore__ inline void ComputeSeg(int64_t len, int64_t alignLen, int32_t idxValue, T updValue, int64_t kStart);
58 __aicore__ inline void CopyOut(int64_t varOffset, int64_t len);64 __aicore__ inline void CopyOut(int64_t varOffset, int64_t len);
59 65 
60 __aicore__ inline int64_t CeilAlign(int64_t x, int64_t base) const { return (x + base - 1) / base * base; }66 __aicore__ inline int64_t CeilAlign(int64_t x, int64_t base) const { return (x + base - 1) / base * base; }
@@ -82,15 +88,22 @@ private:
82 int64_t totalElems_ = 0;88 int64_t totalElems_ = 0;
83 int64_t colsPerChunk_ = 0;89 int64_t colsPerChunk_ = 0;
84 int64_t rowElems_ = 0; // 行在 UB 里的跨度(元素): 使 T 域与 int32 域的行起点同时落在 32B 边界90 int64_t rowElems_ = 0; // 行在 UB 里的跨度(元素): 使 T 域与 int32 域的行起点同时落在 32B 边界
85- int64_t dstStrideT_ = 0; // T 域(var/updates/y)的 burst 间隔, 单位 32B 块91+ int64_t dstStrideT_ = 0; // T 域(var/updates/y)的 burst 间隔, 单位 32B 块
86- int64_t dstStrideI_ = 0; // int32 域(indices)的 burst 间隔, 单位 32B 块92+ int64_t dstStrideI_ = 0; // int32 域(indices)的 burst 间隔, 单位 32B 块
87- bool packedRows_ = false; // 行长是 32B 整数倍 → 走连续搬运 + UB 内倍增铺 idx/updates93+ bool directRows_ = false; // 逐行直算: 轴下标留在寄存器, indices/updates 原地重复读, 不复制操作数
94+ bool alignedFill_ = false; // 铺 tile 时行起点是否 32B 对齐(是则 UB→UB 拷贝, 否则非对齐流式写)
95+ bool compactRows_ = false; // 行在 UB 里紧排(一次连续搬运); false 为按行补齐(每行一个 burst)
88 int64_t rowsPerChunk_ = 1; // inner 较小时一次合并搬运多少行(摊薄每行的 GM 事务开销)96 int64_t rowsPerChunk_ = 1; // inner 较小时一次合并搬运多少行(摊薄每行的 GM 事务开销)
89 int64_t maxChunk_ = 0; // selBuf_ 每段(var/out/updates)的容量, 单位: 元素97 int64_t maxChunk_ = 0; // selBuf_ 每段(var/out/updates)的容量, 单位: 元素
90 int64_t startElem_ = 0;98 int64_t startElem_ = 0;
91 int64_t endElem_ = 0;99 int64_t endElem_ = 0;
92 100 
93protected:101protected:
102+ // 与 op_host 的 MergeMode 一一对应
103+ constexpr static int64_t MERGE_MODE_PAD = 0;
104+ constexpr static int64_t MERGE_MODE_COMPACT_DIRECT = 1;
105+ constexpr static int64_t MERGE_MODE_COMPACT_FILL = 2;
106+ constexpr static int64_t MERGE_MODE_COMPACT_STREAM = 3;
94 constexpr static uint32_t BLOCK_SIZE = platform::GetUbBlockSize();107 constexpr static uint32_t BLOCK_SIZE = platform::GetUbBlockSize();
95 constexpr static uint32_t VL_INT32 = platform::GetVRegSize() / sizeof(int32_t);108 constexpr static uint32_t VL_INT32 = platform::GetVRegSize() / sizeof(int32_t);
96 constexpr static uint32_t VL_T = platform::GetVRegSize() / sizeof(T); // T 域车道数(铺 tile 用)109 constexpr static uint32_t VL_T = platform::GetVRegSize() / sizeof(T); // T 域车道数(铺 tile 用)
@@ -98,7 +111,7 @@ protected:
98 // ① indices/updates 若不同样双缓冲会有 WAR 竞态; ② float16 走 Compare<int32> 产出的掩码 +111 // ① indices/updates 若不同样双缓冲会有 WAR 竞态; ② float16 走 Compare<int32> 产出的掩码 +
99 // Select<half> 时两侧 lane 宽度不一致, 深流水下才显形。112 // Select<half> 时两侧 lane 宽度不一致, 深流水下才显形。
100 // 两者都需要连带改造(见 02_design.md §9.1), 本次交付取确定正确的单缓冲。113 // 两者都需要连带改造(见 02_design.md §9.1), 本次交付取确定正确的单缓冲。
101- constexpr static uint32_t BUFFER_NUM = 1;114+ constexpr static uint32_t BUFFER_NUM = 2;
102};115};
103 116 
104template <typename T, bool INNER_IS_ONE>117template <typename T, bool INNER_IS_ONE>
@@ -132,40 +145,26 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::Init(GM_ADDR var, GM_ADDR
132 // 后续 buffer 偏移整体变化)时暴露。145 // 后续 buffer 偏移整体变化)时暴露。
133 int64_t chunkAlign = colsPerChunk_;146 int64_t chunkAlign = colsPerChunk_;
134 pipe_->InitBuffer(varQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->tBufBytes));147 pipe_->InitBuffer(varQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->tBufBytes));
135- pipe_->InitBuffer(assistBuf_, static_cast<uint32_t>(tiling->i32BufBytes));148+ // i32BufBytes==0 表示本形状的轴下标全部由寄存器生成(inner==1 或逐行直算档), 不开这块缓冲
136- pipe_->InitBuffer(outQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->tBufBytes));149+ if (tiling->i32BufBytes > 0) {
137- if constexpr (!INNER_IS_ONE) {150+ pipe_->InitBuffer(assistBuf_, static_cast<uint32_t>(tiling->i32BufBytes));
138- pipe_->InitBuffer(idxQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->i32BufBytes));
139- pipe_->InitBuffer(updQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->tBufBytes));
140 }151 }
152+ pipe_->InitBuffer(outQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->tBufBytes));
153+ // inner>1 时是整行/多行的 indices/updates; inner==1 时是一段内 m 个 outer 各自的标量
154+ pipe_->InitBuffer(idxQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->idxBufBytes));
155+ pipe_->InitBuffer(updQue_, BUFFER_NUM, static_cast<uint32_t>(tiling->updBufBytes));
141 pipe_->InitBuffer(maskBuf_, static_cast<uint32_t>(tiling->maskBufBytes));156 pipe_->InitBuffer(maskBuf_, static_cast<uint32_t>(tiling->maskBufBytes));
142 maxChunk_ = chunkAlign;157 maxChunk_ = chunkAlign;
143- // 多行合并: 每行在 UB 里按 32B 边界落位(DataCopyPad 的每个 burst 占整块), 行跨度取158+ // 多行合并的布局参数一律取 host 算好的值, 内核不再自行推导:
144- // T 域与 int32 域的公共对齐粒度, 使同一 lane 在两个域里指向同一个元素。159+ // 曾在内核侧按 inner_ 推导行跨度, 导致 assist 逐行写 Duplicate(assistL[r*inner_], ...)
145- // 这样对【任意 inner】都能合并, 不再要求 inner*sizeof(T) 是 32B 整数倍。160+ // 的目的地址在 inner_ 非 8 整数倍时偏离 32B 边界(真机 aivec errcode 340)。
146- if constexpr (!INNER_IS_ONE) {161+ rowElems_ = tiling->rowElems;
147- int64_t alignElems = BLOCK_SIZE / static_cast<int64_t>(sizeof(int32_t));162+ dstStrideT_ = tiling->dstStrideT;
148- int64_t alignElemsT = BLOCK_SIZE / static_cast<int64_t>(sizeof(T));163+ dstStrideI_ = tiling->dstStrideI;
149- if (alignElemsT > alignElems) {164+ rowsPerChunk_ = tiling->rowsPerChunk;
150- alignElems = alignElemsT;165+ directRows_ = (tiling->mergeMode == MERGE_MODE_COMPACT_DIRECT);
151- }166+ alignedFill_ = (tiling->mergeMode != MERGE_MODE_PAD) && (tiling->mergeMode != MERGE_MODE_COMPACT_STREAM);
152- rowElems_ = CeilAlign(inner_ > 0 ? inner_ : 1, alignElems);167+ compactRows_ = (tiling->mergeMode != MERGE_MODE_PAD);
153- int64_t occT = CeilAlign(inner_ * static_cast<int64_t>(sizeof(T)), BLOCK_SIZE);
154- int64_t occI = CeilAlign(inner_ * static_cast<int64_t>(sizeof(int32_t)), BLOCK_SIZE);
155- dstStrideT_ = (rowElems_ * static_cast<int64_t>(sizeof(T)) - occT) / BLOCK_SIZE;
156- dstStrideI_ = (rowElems_ * static_cast<int64_t>(sizeof(int32_t)) - occI) / BLOCK_SIZE;
157- // 多行合并只在 assist 的行跨度 32B 对齐时才启用。
158- // 根因: 合并路径下 GenAssist 逐行写 Duplicate(assistL[r * inner_], ...), 目的地址是
159- // r * inner_ * sizeof(int32_t) 字节; inner_ 不是 8 的整数倍时该地址不落在 32B 边界,
160- // 触发 aivec errcode 340 "The address for VEC to access UB is not aligned"(真机实测,
161- // 且只在 inner>1 出现 —— inner==1 走 CreateVecIndex, 偏移恒为 0)。
162- // 行跨度不对齐时退回单行路径(偏移恒为 0, 天然对齐), 功能正确优先。
163- const bool assistRowAligned = ((inner_ * static_cast<int64_t>(sizeof(int32_t))) % BLOCK_SIZE == 0);
164- rowsPerChunk_ = (assistRowAligned && rowElems_ > 0 && rowElems_ <= chunkAlign) ? (chunkAlign / rowElems_) : 1;
165- // 行长本身就是 32B 整数倍时, 补齐后与紧排一致 → 可以走"一次大块连续搬运"的打包路径;
166- // 否则只能一行一个 burst(多 burst 路径), 那是为覆盖任意 inner 付的代价。
167- packedRows_ = (rowElems_ == inner_);
168- }
169 if constexpr (NEED_CAST) {168 if constexpr (NEED_CAST) {
170 // var / out / updates 三段 half 暂存(字节数同样由 host 算好)169 // var / out / updates 三段 half 暂存(字节数同样由 host 算好)
171 pipe_->InitBuffer(selBuf_, static_cast<uint32_t>(tiling->selBufBytes));170 pipe_->InitBuffer(selBuf_, static_cast<uint32_t>(tiling->selBufBytes));
@@ -212,31 +211,44 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::CopyInIdxUpd(int64_t idxOf
212}211}
213 212 
214// 生成 A2 融合 pass 里那个 assist 张量的等价内容: assist[o,k,i] == k(沿 dimension 轴的下标)。213// 生成 A2 融合 pass 里那个 assist 张量的等价内容: assist[o,k,i] == k(沿 dimension 轴的下标)。
215-// INNER_IS_ONE 时一段覆盖连续的 k, 用一条 CreateVecIndex 出等差数列; 否则一行内 k 是常量,214+// **只服务"一个寄存器块跨多行"的两档**(紧排铺 tile / 按行补齐): 这两档里 k 在一个寄存器内
216-// 逐行 Duplicate。两者都只动 UB, 相比 A2 由图侧传入整块 int32 省掉一整条 GM→UB 流。215+// 逐行变, 算不出闭式, 只能物化到 UB。其余路径的 k 已由 VF 内的 Reg::Arange / Duplicate 直接
216+// 在寄存器里生成(见 AssistSrc), 既不占 UB 也不多走一趟写。
217template <typename T, bool INNER_IS_ONE>217template <typename T, bool INNER_IS_ONE>
218__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::GenAssist(int64_t kStart, int64_t rows, int64_t rowLen,218__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::GenAssist(int64_t kStart, int64_t rows, int64_t rowLen,
219 int64_t rowStride)219 int64_t rowStride)
220{220{
221 LocalTensor<int32_t> assistL = assistBuf_.template Get<int32_t>();221 LocalTensor<int32_t> assistL = assistBuf_.template Get<int32_t>();
222- if constexpr (INNER_IS_ONE) {222+ {
223- AscendC::CreateVecIndex(assistL, static_cast<int32_t>(kStart), static_cast<uint32_t>(rowLen));223+ // 多行: 第 0 行填 kStart, 其余行用倍增复制 + 一次常量偏移铺满 —— 指令数 O(log2 rows) 而非
224- } else {224+ // O(rows), 且每条都是成段满车道。逐行一条 Duplicate 时每行的指令发射固定成本会随行数线性
225- // 一行之内轴下标恒等于该行的 k, 逐行填常量即可。225+ // 累加, 在 inner 小、行数多的形状上压倒有效计算(实测占总时延 83%~96%)。
226- // 曾尝试两种"减少指令数"的优化, 实测都是负优化, 已回退(详见 02 设计文档):226+ // 行跨度 rowStride 由 host 取 32B 公倍数(见 tiling), 故 filled * rowStride 处的目的偏移恒对齐。
227- // ① 另开一块 chunk 大小的模板缓冲 + 每 chunk 一次 Adds —— 缓冲进 UB 预算后 colsPerChunk227+ AscendC::Duplicate(assistL, static_cast<int32_t>(kStart), static_cast<int32_t>(rowLen));
228- // 变小、chunk 数变多, 固定开销远超省下的指令(最差例 11.4us→33.4us);228+ for (int64_t filled = 1; filled < rows; filled *= 2) {
229- // ② 在 assistBuf_ 上做增量 Adds —— 模板被单行/INNER_IS_ONE 路径反复置无效, 每次重建要229+ int64_t copyRows = Min(filled, rows - filled);
230- // 循环 rowsPerChunk_ 次(多于 rows 次), 更慢(11.4us→41.5us)。230+ int64_t cur = filled * rowStride;
231- for (int64_t r = 0; r < rows; ++r) {231+ int64_t cnt = copyRows * rowStride;
232- AscendC::Duplicate(assistL[r * rowStride], static_cast<int32_t>(kStart + r), static_cast<int32_t>(rowLen));232+ AscendC::PipeBarrier<PIPE_V>();
233+ if ((cur * static_cast<int64_t>(sizeof(int32_t))) % BLOCK_SIZE == 0) {
234+ // 行跨度是 32B 整数倍: 目的偏移天然对齐, UB→UB 搬 + 一次整体加常量
235+ AscendC::DataCopy(assistL[cur], assistL[0], static_cast<uint32_t>(cnt));
236+ AscendC::PipeBarrier<PIPE_V>();
237+ AscendC::Adds(assistL[cur], assistL[cur], static_cast<int32_t>(filled), static_cast<int32_t>(cnt));
238+ } else {
239+ // 非对齐: 复制与加常量合成一条非对齐流式写
240+ uint16_t loops = static_cast<uint16_t>((cnt + VL_INT32 - 1) / VL_INT32);
241+ TileAppendAddsUnalignVF<int32_t>((__ubuf__ int32_t*)assistL.GetPhyAddr(), static_cast<uint32_t>(cur),
242+ static_cast<uint32_t>(cnt), static_cast<int32_t>(filled), VL_INT32,
243+ loops);
244+ }
233 }245 }
234 }246 }
235 AscendC::PipeBarrier<PIPE_V>();247 AscendC::PipeBarrier<PIPE_V>();
236}248}
237 249 
238-// 打包路径(行长是 32B 整数倍): var 由一次连续搬运载入, 轴下标在 UB 内生成, indices/updates 只搬一行,250+// 紧排但行长非 32B 整数倍时的路径: 行起点不落在寄存器块边界, 无法逐行直算, 只能先把 indices/updates
239-// 在 UB 内倍增铺满(log2(rows) 次 UB→UB, 走 PIPE_V), 然后整段一次比较 + 一次选择。251+// 倍增铺成多行(非对齐流式写), 配合同样铺出来的轴下标, 再对整段做一次比较 + 一次选择。
240// 与多 burst 路径的取舍: 这里省掉了 m 个小 burst 的 DMA 固定开销 —— inner 小时那是主要成本。252// 与多 burst 路径的取舍: 这里省掉了 m 个小 burst 的 DMA 固定开销 —— inner 小时那是主要成本。
241template <typename T, bool INNER_IS_ONE>253template <typename T, bool INNER_IS_ONE>
242__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputePacked(int64_t rows, int64_t totalLen)254__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputePacked(int64_t rows, int64_t totalLen)
@@ -251,7 +263,7 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputePacked(int64_t rows
251 int64_t copyRows = Min(filled, rows - filled);263 int64_t copyRows = Min(filled, rows - filled);
252 int64_t cur = filled * inner_;264 int64_t cur = filled * inner_;
253 int64_t cnt = copyRows * inner_;265 int64_t cnt = copyRows * inner_;
254- if (packedRows_) {266+ if (alignedFill_) {
255 // 行长是 32B 整数倍 → 目的偏移天然对齐, 直接 UB→UB 搬267 // 行长是 32B 整数倍 → 目的偏移天然对齐, 直接 UB→UB 搬
256 AscendC::DataCopy(idxL[cur], idxL[0], static_cast<uint32_t>(cnt));268 AscendC::DataCopy(idxL[cur], idxL[0], static_cast<uint32_t>(cnt));
257 AscendC::DataCopy(updL[cur], updL[0], static_cast<uint32_t>(cnt));269 AscendC::DataCopy(updL[cur], updL[0], static_cast<uint32_t>(cnt));
@@ -279,6 +291,78 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputePacked(int64_t rows
279 updQue_.FreeTensor(updL);291 updQue_.FreeTensor(updL);
280}292}
281 293 
294+// var 的 m 行一次搬入, UB 侧按 rowElems_ 跨度落位(每行占整数个 32B 块), 与 idx/updates/assist
295+// 的行起点对齐 —— 这是"按行补齐(pad)"合并路径的搬入端, 对任意 inner 都成立。
296+template <typename T, bool INNER_IS_ONE>
297+__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::CopyInRowsPad(int64_t varOffset, int64_t rows)
298+{
299+ LocalTensor<T> varL = varQue_.template AllocTensor<T>();
300+ DataCopyPadExtParams<T> padT{true, 0, 0, ScalarZero<T>()};
301+ DataCopyExtParams params;
302+ params.blockCount = static_cast<uint16_t>(rows);
303+ params.blockLen = static_cast<uint32_t>(inner_ * sizeof(T));
304+ params.srcStride = 0; // GM 侧 rows 行连续
305+ params.dstStride = dstStrideT_; // UB 侧按 32B 块补齐
306+ AscendC::DataCopyPad(varL, varGm_[varOffset], params, padT);
307+ varQue_.template EnQue<T>(varL);
308+}
309+ 
310+// inner==1 且一段能装下多个 outer 时的主路径: var 一次连续搬入 m 个 outer, indices/updates 各 m 个
311+// 标量也一次搬入, 轴下标 0..D-1 对每个 outer 都相同故只生成一次, 然后逐 outer 一条 VF(用各自的标量)。
312+// 相比"一段一个 outer", 每个 outer 省掉一整条 MTE2→V→MTE3 的段固定成本与两次 GM 标量读。
313+template <typename T, bool INNER_IS_ONE>
314+__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeOuters(int64_t outers)
315+{
316+ LocalTensor<T> varL = varQue_.template DeQue<T>();
317+ LocalTensor<int32_t> idxL = idxQue_.template DeQue<int32_t>();
318+ LocalTensor<T> updL = updQue_.template DeQue<T>();
319+ LocalTensor<T> outL = outQue_.template AllocTensor<T>();
320+ 
321+ // 轴下标不再物化到 UB: 每个 outer 沿被选轴是 0..D-1 的等差数列, VF 内一条 Reg::Arange 生成
322+ // 标量要从 UB 读: 队列的 DeQue 只给出 MTE2→V, 标量口径需另外补 MTE2→S
323+ AscendC::SetFlag<AscendC::HardEvent::MTE2_S>(EVENT_ID0);
324+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_S>(EVENT_ID0);
325+ 
326+ uint16_t repeatTimes = static_cast<uint16_t>((dimSize_ + VL_INT32 - 1) / VL_INT32);
327+ __ubuf__ T* outAddr = (__ubuf__ T*)outL.GetPhyAddr();
328+ __ubuf__ T* varAddr = (__ubuf__ T*)varL.GetPhyAddr();
329+ for (int64_t j = 0; j < outers; ++j) {
330+ int32_t idxValue = idxL.GetValue(j);
331+ T updValue = updL.GetValue(j);
332+ ArgMaxGradSelectVF<T, true, AssistSrc::ARANGE>(outAddr + j * dimSize_, varAddr + j * dimSize_, nullptr, nullptr,
333+ nullptr, idxValue, updValue, ScalarToFloat<T>(updValue),
334+ static_cast<uint32_t>(dimSize_), VL_INT32, repeatTimes, 0);
335+ }
336+ 
337+ outQue_.template EnQue<T>(outL);
338+ varQue_.FreeTensor(varL);
339+ idxQue_.FreeTensor(idxL);
340+ updQue_.FreeTensor(updL);
341+}
342+ 
343+// 行长是 32B 整数倍时的主路径: 不生成轴下标、不把 indices/updates 铺成多行。
344+// 每个向量寄存器块必然落在同一行内, 该行的轴下标是标量常量(留在寄存器里), indices/updates
345+// 每行都从同一份原地重复读 —— 相比"先复制操作数再整段选择", 省掉三遍对整段数据的 UB 读写。
346+template <typename T, bool INNER_IS_ONE>
347+__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeRowsDirect(int64_t kStart, int64_t rows)
348+{
349+ LocalTensor<T> varL = varQue_.template DeQue<T>();
350+ LocalTensor<int32_t> idxL = idxQue_.template DeQue<int32_t>();
351+ LocalTensor<T> updL = updQue_.template DeQue<T>();
352+ LocalTensor<T> outL = outQue_.template AllocTensor<T>();
353+ 
354+ uint16_t repeatsPerRow = static_cast<uint16_t>((inner_ + VL_INT32 - 1) / VL_INT32);
355+ ArgMaxGradSelectRowsVF<T>((__ubuf__ T*)outL.GetPhyAddr(), (__ubuf__ T*)varL.GetPhyAddr(),
356+ (__ubuf__ T*)updL.GetPhyAddr(), (__ubuf__ int32_t*)idxL.GetPhyAddr(),
357+ static_cast<int32_t>(kStart), ScalarZero<T>(), static_cast<uint32_t>(rows),
358+ static_cast<uint32_t>(inner_), VL_INT32, repeatsPerRow);
359+ 
360+ outQue_.template EnQue<T>(outL);
361+ varQue_.FreeTensor(varL);
362+ idxQue_.FreeTensor(idxL);
363+ updQue_.FreeTensor(updL);
364+}
365+ 
282// indices/updates 只有一行, 但要与 m 行 var 对齐: 令 GM 侧 srcStride = -blockLen,366// indices/updates 只有一行, 但要与 m 行 var 对齐: 令 GM 侧 srcStride = -blockLen,
283// 每个 burst 回到同一行 —— 相当于在搬运里就把这一行复制了 m 遍, 不需要 UB 内再铺。367// 每个 burst 回到同一行 —— 相当于在搬运里就把这一行复制了 m 遍, 不需要 UB 内再铺。
284// (arch35 的 DataCopyExtParams::srcStride 是 int64_t 有符号, 且约束为 >= -blockLen。)368// (arch35 的 DataCopyExtParams::srcStride 是 int64_t 有符号, 且约束为 >= -blockLen。)
@@ -344,11 +428,11 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputePad(int64_t rows)
344// rows 行 × rowLen 列: var 已按 rows*rowLen 连续搬入, 轴下标在 UB 内生成, indices/updates 只有 rowLen 个、各行复用。428// rows 行 × rowLen 列: var 已按 rows*rowLen 连续搬入, 轴下标在 UB 内生成, indices/updates 只有 rowLen 个、各行复用。
345// 逐行做 Compare/Select(向量开销远小于每行 4 次 GM 事务), GM 搬运在外层已合并。429// 逐行做 Compare/Select(向量开销远小于每行 4 次 GM 事务), GM 搬运在外层已合并。
346template <typename T, bool INNER_IS_ONE>430template <typename T, bool INNER_IS_ONE>
347-__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeRows(int64_t rows, int64_t rowLen, int64_t alignRowLen)431+__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeRows(int64_t rows, int64_t rowLen, int64_t alignRowLen,
432+ int64_t kStart)
348{433{
349 (void)rows;434 (void)rows;
350 LocalTensor<T> varL = varQue_.template DeQue<T>();435 LocalTensor<T> varL = varQue_.template DeQue<T>();
351- LocalTensor<int32_t> assistL = assistBuf_.template Get<int32_t>();
352 LocalTensor<int32_t> idxL = idxQue_.template DeQue<int32_t>();436 LocalTensor<int32_t> idxL = idxQue_.template DeQue<int32_t>();
353 LocalTensor<T> updL = updQue_.template DeQue<T>();437 LocalTensor<T> updL = updQue_.template DeQue<T>();
354 LocalTensor<T> outL = outQue_.template AllocTensor<T>();438 LocalTensor<T> outL = outQue_.template AllocTensor<T>();
@@ -356,11 +440,12 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeRows(int64_t rows,
356 // 元素数按 VL 取整: 每轮都是整寄存器读写, 不出现半个寄存器的尾轮(实测尾轮会触发440 // 元素数按 VL 取整: 每轮都是整寄存器读写, 不出现半个寄存器的尾轮(实测尾轮会触发
357 // errcode 340 "VEC 访问 UB 地址未对齐")。尾部算出来的是脏值, 但 CopyOut 只写回前441 // errcode 340 "VEC 访问 UB 地址未对齐")。尾部算出来的是脏值, 但 CopyOut 只写回前
358 // rowLen 个元素, 不影响结果; buffer 容量由 host 保证 >= 一个整寄存器。442 // rowLen 个元素, 不影响结果; buffer 容量由 host 保证 >= 一个整寄存器。
443+ // 单行段: 整段同一个 k, VF 内一条 Duplicate 进寄存器, 不落 UB
359 uint16_t repeatTimes = static_cast<uint16_t>((alignRowLen + VL_INT32 - 1) / VL_INT32);444 uint16_t repeatTimes = static_cast<uint16_t>((alignRowLen + VL_INT32 - 1) / VL_INT32);
360- ArgMaxGradSelectVF<T, false>((__ubuf__ T*)outL.GetPhyAddr(), (__ubuf__ T*)varL.GetPhyAddr(),445+ ArgMaxGradSelectVF<T, false, AssistSrc::SCALAR>(
361- (__ubuf__ T*)updL.GetPhyAddr(), (__ubuf__ int32_t*)assistL.GetPhyAddr(),446+ (__ubuf__ T*)outL.GetPhyAddr(), (__ubuf__ T*)varL.GetPhyAddr(), (__ubuf__ T*)updL.GetPhyAddr(), nullptr,
362- (__ubuf__ int32_t*)idxL.GetPhyAddr(), 0, ScalarZero<T>(), 0.0f,447+ (__ubuf__ int32_t*)idxL.GetPhyAddr(), 0, ScalarZero<T>(), 0.0f, static_cast<uint32_t>(alignRowLen), VL_INT32,
363- static_cast<uint32_t>(alignRowLen), VL_INT32, repeatTimes);448+ repeatTimes, static_cast<int32_t>(kStart));
364 449 
365 outQue_.template EnQue<T>(outL);450 outQue_.template EnQue<T>(outL);
366 varQue_.FreeTensor(varL);451 varQue_.FreeTensor(varL);
@@ -370,17 +455,17 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeRows(int64_t rows,
370 455 
371template <typename T, bool INNER_IS_ONE>456template <typename T, bool INNER_IS_ONE>
372__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeSeg(int64_t len, int64_t alignLen, int32_t idxValue,457__aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::ComputeSeg(int64_t len, int64_t alignLen, int32_t idxValue,
373- T updValue)458+ T updValue, int64_t kStart)
374{459{
375 (void)alignLen;460 (void)alignLen;
376 LocalTensor<T> varL = varQue_.template DeQue<T>();461 LocalTensor<T> varL = varQue_.template DeQue<T>();
377- LocalTensor<int32_t> assistL = assistBuf_.template Get<int32_t>();
378 LocalTensor<T> outL = outQue_.template AllocTensor<T>();462 LocalTensor<T> outL = outQue_.template AllocTensor<T>();
379 463 
464+ // 本段沿被选轴连续: k = kStart + 车道号, VF 内一条 Reg::Arange 生成, 不落 UB
380 uint16_t repeatTimes = static_cast<uint16_t>((len + VL_INT32 - 1) / VL_INT32);465 uint16_t repeatTimes = static_cast<uint16_t>((len + VL_INT32 - 1) / VL_INT32);
381- ArgMaxGradSelectVF<T, true>((__ubuf__ T*)outL.GetPhyAddr(), (__ubuf__ T*)varL.GetPhyAddr(), nullptr,466+ ArgMaxGradSelectVF<T, true, AssistSrc::ARANGE>(
382- (__ubuf__ int32_t*)assistL.GetPhyAddr(), nullptr, idxValue, updValue,467+ (__ubuf__ T*)outL.GetPhyAddr(), (__ubuf__ T*)varL.GetPhyAddr(), nullptr, nullptr, nullptr, idxValue, updValue,
383- ScalarToFloat<T>(updValue), static_cast<uint32_t>(len), VL_INT32, repeatTimes);468+ ScalarToFloat<T>(updValue), static_cast<uint32_t>(len), VL_INT32, repeatTimes, static_cast<int32_t>(kStart));
384 469 
385 outQue_.template EnQue<T>(outL);470 outQue_.template EnQue<T>(outL);
386 varQue_.FreeTensor(varL);471 varQue_.FreeTensor(varL);
@@ -406,56 +491,91 @@ __aicore__ inline void ArgMaxGradND<T, INNER_IS_ONE>::Process()
406 if (totalElems_ <= 0 || startElem_ >= endElem_) {491 if (totalElems_ <= 0 || startElem_ >= endElem_) {
407 return;492 return;
408 }493 }
409- int64_t oStride = dimSize_ * inner_; // 一个 outer 覆盖的元素数494+ const int64_t oStride = dimSize_ * inner_; // 一个 outer 覆盖的元素数
410 int64_t g = startElem_;495 int64_t g = startElem_;
411 while (g < endElem_) {496 while (g < endElem_) {
412- int64_t o = g / oStride;497+ const int64_t o = g / oStride;
413- int64_t rem = endElem_ - g;498+ const int64_t inO = g - o * oStride; // 段起点在本 outer 内的偏移
414- int64_t len = 0;499+ const int64_t rem = endElem_ - g;
415- int64_t idxOffset = 0;
416 if constexpr (INNER_IS_ONE) {500 if constexpr (INNER_IS_ONE) {
417- // inner==1: 同一个 outer 的元素沿被选轴连续, indices/updates 各只有一个值501+ g += ProcessOuterSeg(g, o, inO, rem);
418- int64_t k = g - o * oStride;
419- len = Min(Min(dimSize_ - k, rem), colsPerChunk_);
420- int32_t idxValue = indicesGm_.GetValue(o);
421- T updValue = updatesGm_.GetValue(o);
422- int64_t alignLen = CeilAlign(len, static_cast<int64_t>(VL_INT32));
423- CopyInCommon(g, len);
424- GenAssist(k, 1, alignLen, 0); // 本段覆盖连续的 k: 等差数列
425- ComputeSeg(len, alignLen, idxValue, updValue);
426- CopyOut(g, len);
427 } else {502 } else {
428- // inner>1: 落在整行起点时把多行并成一段(多 burst 搬运, 每行在 UB 里 32B 对齐落位);503+ g += ProcessRowSeg(g, o, inO % inner_, rem);
429- // 头/尾不完整的行单独走一次单段处理。
430- int64_t inRow = (g - o * oStride) % inner_;
431- int64_t rows = 1;
432- if (inRow == 0 && rowsPerChunk_ > 1 && rem >= inner_) {
433- int64_t rowsLeftInO = (oStride - (g - o * oStride)) / inner_;
434- rows = Min(Min(rem / inner_, rowsLeftInO), rowsPerChunk_);
435- }
436- if (rows > 1) {
437- // 无论行长是否 32B 对齐, 都走"一次连续搬运 + UB 内铺 tile + 整段一次选择";
438- // 两者只差铺 tile 的手段(对齐用 UB→UB 搬, 非对齐用非对齐流式写)。
439- len = rows * inner_;
440- CopyInCommon(g, len);
441- CopyInIdxUpd(o * inner_, inner_);
442- GenAssist((g - o * oStride) / inner_, rows, inner_, inner_); // 每行一个常量 k
443- ComputePacked(rows, len);
444- CopyOut(g, len);
445- } else {
446- len = Min(Min(inner_ - inRow, rem), colsPerChunk_);
447- int64_t alignLen = CeilAlign(len, static_cast<int64_t>(VL_INT32));
448- CopyInCommon(g, len);
449- CopyInIdxUpd(o * inner_ + inRow, len);
450- GenAssist((g - o * oStride) / inner_, 1, alignLen, 0); // 单行: 整段同一个 k
451- ComputeRows(1, len, alignLen);
452- CopyOut(g, len);
453- }
454 }504 }
455- g += len;
456 }505 }
457}506}
458 507 
508+// inner==1: 同一个 outer 的元素沿被选轴连续, indices/updates 各只有一个值
509+template <typename T, bool INNER_IS_ONE>
510+__aicore__ inline int64_t ArgMaxGradND<T, INNER_IS_ONE>::ProcessOuterSeg(int64_t g, int64_t o, int64_t k, int64_t rem)
511+{
512+ if (k == 0 && rowsPerChunk_ > 1 && rem >= dimSize_) {
513+ // 一段装下多个 outer: 合并搬运 + 逐 outer 直算, 摊薄每段的固定成本
514+ const int64_t outers = Min(rem / dimSize_, rowsPerChunk_);
515+ const int64_t len = outers * dimSize_;
516+ CopyInCommon(g, len);
517+ CopyInIdxUpd(o, outers);
518+ ComputeOuters(outers);
519+ CopyOut(g, len);
520+ return len;
521+ }
522+ const int64_t len = Min(Min(dimSize_ - k, rem), colsPerChunk_);
523+ const int32_t idxValue = indicesGm_.GetValue(o);
524+ const T updValue = updatesGm_.GetValue(o);
525+ const int64_t alignLen = CeilAlign(len, static_cast<int64_t>(VL_INT32));
526+ CopyInCommon(g, len);
527+ ComputeSeg(len, alignLen, idxValue, updValue, k); // 轴下标在 VF 内按 k 起的等差数列生成
528+ CopyOut(g, len);
529+ return len;
530+}
531+ 
532+// inner>1: 落在整行起点时把多行并成一段(合并形态由 host 的 mergeMode 定);
533+// 头/尾不完整的行单独走一次单段处理。
534+template <typename T, bool INNER_IS_ONE>
535+__aicore__ inline int64_t ArgMaxGradND<T, INNER_IS_ONE>::ProcessRowSeg(int64_t g, int64_t o, int64_t inRow, int64_t rem)
536+{
537+ const int64_t oStride = dimSize_ * inner_;
538+ const int64_t kStart = (g - o * oStride) / inner_; // 段起点所在行在被选轴上的下标
539+ int64_t rows = 1;
540+ if (inRow == 0 && rowsPerChunk_ > 1 && rem >= inner_) {
541+ const int64_t rowsLeftInO = (oStride - (g - o * oStride)) / inner_;
542+ rows = Min(Min(rem / inner_, rowsLeftInO), rowsPerChunk_);
543+ }
544+ if (rows <= 1) {
545+ const int64_t len = Min(Min(inner_ - inRow, rem), colsPerChunk_);
546+ const int64_t alignLen = CeilAlign(len, static_cast<int64_t>(VL_INT32));
547+ CopyInCommon(g, len);
548+ CopyInIdxUpd(o * inner_ + inRow, len);
549+ ComputeRows(1, len, alignLen, kStart); // 单行: 整段同一个 k, 寄存器内 Duplicate
550+ CopyOut(g, len);
551+ return len;
552+ }
553+ const int64_t len = rows * inner_;
554+ if (directRows_) {
555+ // 一次连续搬运 + 逐行直算(轴下标留在寄存器, 操作数不复制)
556+ CopyInCommon(g, len);
557+ CopyInIdxUpd(o * inner_, inner_);
558+ ComputeRowsDirect(kStart, rows);
559+ CopyOut(g, len);
560+ } else if (compactRows_) {
561+ // 一行占不满寄存器(或行起点不落在块边界): 先把操作数铺成多行, 再对整段做一次选择。
562+ CopyInCommon(g, len);
563+ CopyInIdxUpd(o * inner_, inner_);
564+ GenAssist(kStart, rows, inner_, inner_); // 每行一个常量 k
565+ ComputePacked(rows, len);
566+ CopyOut(g, len);
567+ } else {
568+ // 按行补齐(pad)合并: 每行一个 burst, UB 内按 rowElems_ 落位。行长够肥时 burst 的
569+ // 固定开销已被摊薄, 比紧排少付一遍 UB 内铺 tile 的向量开销。
570+ CopyInRowsPad(g, rows);
571+ CopyInIdxUpdRepeat(o * inner_, rows);
572+ GenAssist(kStart, rows, inner_, rowElems_);
573+ ComputePad(rows);
574+ CopyOutRowsPad(g, rows);
575+ }
576+ return len;
577+}
578+ 
459} // namespace ArgMaxGrad579} // namespace ArgMaxGrad
460 580 
461#endif // ARG_MAX_GRAD_ND_H581#endif // ARG_MAX_GRAD_ND_H
Mindex/arg_max_grad/op_kernel/arch35/arg_max_grad_tiling_data.h+9-1
@@ -27,9 +27,17 @@ struct ArgMaxGradArch35TilingData {
27 // InitBuffer, 不得再做二次对齐/补齐 —— 内核侧二次对齐会让实际占用偏离 host 的预算,27 // InitBuffer, 不得再做二次对齐/补齐 —— 内核侧二次对齐会让实际占用偏离 host 的预算,
28 // 且各 buffer 起点会偏离向量寄存器整宽边界(实测 errcode 340: VEC 访问 UB 地址未对齐)。28 // 且各 buffer 起点会偏离向量寄存器整宽边界(实测 errcode 340: VEC 访问 UB 地址未对齐)。
29 int64_t tBufBytes = 0; // T 域单块字节数(var / out / updates 各一块)29 int64_t tBufBytes = 0; // T 域单块字节数(var / out / updates 各一块)
30- int64_t i32BufBytes = 0; // int32 域单块字节数(assist / indices 各一块)30+ int64_t i32BufBytes = 0; // 轴下标(assist)缓冲字节数; 逐行直算档不生成轴下标, 只留最小块
31+ int64_t idxBufBytes = 0; // indices 缓冲字节数: 逐行直算只需一行, 铺 tile 档需整段
32+ int64_t updBufBytes = 0; // updates 缓冲字节数, 口径同上
31 int64_t maskBufBytes = 0; // 掩码缓冲字节数33 int64_t maskBufBytes = 0; // 掩码缓冲字节数
32 int64_t selBufBytes = 0; // int8 借道 half 的三块暂存合计字节数; 非 int8 为 034 int64_t selBufBytes = 0; // int8 借道 half 的三块暂存合计字节数; 非 int8 为 0
35+ // 多行合并参数由 host 一次算准, 内核直接使用, 不自行推导(内核侧重算会与 host 的 UB 预算脱节)
36+ int64_t rowsPerChunk = 1; // 一次最多合并多少行; 1 表示不合并
37+ int64_t rowElems = 0; // MERGE_PAD 下一行在 UB 里的跨度(元素), 取 T 域与 int32 域的公共 32B 粒度
38+ int64_t dstStrideT = 0; // MERGE_PAD 下 T 域(var/updates/y)的 burst 间隔, 单位 32B 块
39+ int64_t dstStrideI = 0; // MERGE_PAD 下 int32 域(indices)的 burst 间隔, 单位 32B 块
40+ int64_t mergeMode = 0; // 见 op_host 同名常量: 0=按行补齐, 1=紧排逐行直算, 2=紧排铺 tile, 3=紧排流式铺设
33};41};
34 42 
35#endif // ARG_MAX_GRAD_TILING_DATA_H43#endif // ARG_MAX_GRAD_TILING_DATA_H
Mindex/arg_max_grad/op_kernel/arch35/arg_max_grad_vf.h+74-3
@@ -156,16 +156,45 @@ __simd_vf__ inline void TileAppendUnalignVF(__ubuf__ T* base, uint32_t curElems,
156 AscendC::Reg::StoreUnAlignPost<T>(dst, ureg, 0);156 AscendC::Reg::StoreUnAlignPost<T>(dst, ureg, 0);
157}157}
158 158 
159+// 与 TileAppendUnalignVF 同构, 但追加时给每个元素加一个常量: 轴下标 assist 的第 r 行等于第 0 行加 r,
160+// 倍增铺设时正好是"复制 filled 行 + 整体加 filled"。合成一条指令省掉一次独立的 Adds 遍历。
161+template <typename T>
162+__simd_vf__ inline void TileAppendAddsUnalignVF(__ubuf__ T* base, uint32_t curElems, uint32_t copyElems, T addend,
163+ uint32_t lane, uint16_t loops)
164+{
165+ AscendC::Reg::UnalignRegForStore ureg;
166+ __ubuf__ T* dst = base + curElems;
167+ AscendC::Reg::MaskReg mask;
168+ uint32_t remain = copyElems;
169+ for (uint16_t t = 0; t < loops; ++t) {
170+ uint32_t before = remain;
171+ mask = AscendC::Reg::UpdateMask<T>(remain);
172+ uint32_t step = before - remain;
173+ AscendC::Reg::RegTensor<T> reg;
174+ AscendC::Reg::LoadAlign(reg, base + t * lane);
175+ AscendC::Reg::Adds(reg, reg, addend, mask);
176+ AscendC::Reg::StoreUnAlign(dst, reg, ureg, step);
177+ }
178+ AscendC::Reg::StoreUnAlignPost<T>(dst, ureg, 0);
179+}
180+ 
159// ── regbase VF: mask = (轴下标 == idx), out = mask ? updates : var ──────────────181// ── regbase VF: mask = (轴下标 == idx), out = mask ? updates : var ──────────────
160// 掩码全程留在 MaskReg 里(不落 UB), 比较恒在 int32 的 64 车道上做; T 宽度不等于 4 字节时182// 掩码全程留在 MaskReg 里(不落 UB), 比较恒在 int32 的 64 车道上做; T 宽度不等于 4 字节时
161// 随路拆包到 32bit 车道、选完再打包回去。范式照 loss/chamfer_distance 的 ChamferDistVF。183// 随路拆包到 32bit 车道、选完再打包回去。范式照 loss/chamfer_distance 的 ChamferDistVF。
184+// 轴下标(assist)的来源。**能在寄存器里算出来的就不落 UB**:
185+// ARANGE: 本段沿被选轴连续(inner==1), k = kStart + 车道号 —— 一条 vci(Reg::Arange)搞定;
186+// SCALAR: 整段同一个 k(inner>1 的单行段) —— 一条 Duplicate 进寄存器;
187+// UB : 一个寄存器块里跨多行、k 逐行变(紧排/补齐合并档) —— 只有这一档还需要物化到 UB。
188+// 前两档省掉的不只是指令: assist 的整块 UB 缓冲(每元素 4B)一并消失, 段长随之变大。
189+enum class AssistSrc : int { ARANGE = 0, SCALAR = 1, UB = 2 };
190+ 
162// 写法要点: 必须是 __simd_vf__ 自由函数 + 裸 __ubuf__ 指针推进; 塞进成员函数会触发后端191// 写法要点: 必须是 __simd_vf__ 自由函数 + 裸 __ubuf__ 指针推进; 塞进成员函数会触发后端
163// "Unsupported Inst must be hoisted"。192// "Unsupported Inst must be hoisted"。
164-template <typename T, bool IDX_IS_SCALAR>193+template <typename T, bool IDX_IS_SCALAR, AssistSrc ASSIST = AssistSrc::UB>
165__simd_vf__ inline void ArgMaxGradSelectVF(__ubuf__ T* outAddr, __ubuf__ T* varAddr, __ubuf__ T* updAddr,194__simd_vf__ inline void ArgMaxGradSelectVF(__ubuf__ T* outAddr, __ubuf__ T* varAddr, __ubuf__ T* updAddr,
166 __ubuf__ int32_t* assistAddr, __ubuf__ int32_t* idxAddr, int32_t idxScalar,195 __ubuf__ int32_t* assistAddr, __ubuf__ int32_t* idxAddr, int32_t idxScalar,
167 T updScalar, float updScalarF, uint32_t count, uint32_t lane,196 T updScalar, float updScalarF, uint32_t count, uint32_t lane,
168- uint16_t repeatTimes)197+ uint16_t repeatTimes, int32_t kStart = 0)
169{198{
170 AscendC::Reg::RegTensor<int32_t> assistReg;199 AscendC::Reg::RegTensor<int32_t> assistReg;
171 AscendC::Reg::RegTensor<int32_t> idxReg;200 AscendC::Reg::RegTensor<int32_t> idxReg;
@@ -176,9 +205,16 @@ __simd_vf__ inline void ArgMaxGradSelectVF(__ubuf__ T* outAddr, __ubuf__ T* varA
176 if constexpr (IDX_IS_SCALAR) {205 if constexpr (IDX_IS_SCALAR) {
177 AscendC::Reg::Duplicate(idxReg, idxScalar);206 AscendC::Reg::Duplicate(idxReg, idxScalar);
178 }207 }
208+ if constexpr (ASSIST == AssistSrc::SCALAR) {
209+ AscendC::Reg::Duplicate(assistReg, kStart); // 整段同一个 k, 循环外一次
210+ }
179 for (uint16_t i = 0; i < repeatTimes; ++i) {211 for (uint16_t i = 0; i < repeatTimes; ++i) {
180 mask = AscendC::Reg::UpdateMask<int32_t>(remain);212 mask = AscendC::Reg::UpdateMask<int32_t>(remain);
181- AscendC::Reg::LoadAlign(assistReg, assistAddr + i * lane);213+ if constexpr (ASSIST == AssistSrc::ARANGE) {
214+ AscendC::Reg::Arange(assistReg, kStart + static_cast<int32_t>(i * lane));
215+ } else if constexpr (ASSIST == AssistSrc::UB) {
216+ AscendC::Reg::LoadAlign(assistReg, assistAddr + i * lane);
217+ }
182 if constexpr (!IDX_IS_SCALAR) {218 if constexpr (!IDX_IS_SCALAR) {
183 AscendC::Reg::LoadAlign(idxReg, idxAddr + i * lane);219 AscendC::Reg::LoadAlign(idxReg, idxAddr + i * lane);
184 }220 }
@@ -197,6 +233,41 @@ __simd_vf__ inline void ArgMaxGradSelectVF(__ubuf__ T* outAddr, __ubuf__ T* varA
197 }233 }
198}234}
199 235 
236+// ── 多行直算: 不物化轴下标, 也不把 indices/updates 复制成多行 ────────────────────
237+// 行长是 32B 整数倍时, 每个向量寄存器块必然落在同一行内 —— 该行的轴下标就是个标量常量,
238+// Duplicate 进寄存器即可(不落 UB); indices/updates 只有一行, 每行都从同一份原地重复读。
239+// 于是"复制操作数"这件事整个消失, 只剩一遍比较 + 一遍选择。
240+template <typename T>
241+__simd_vf__ inline void ArgMaxGradSelectRowsVF(__ubuf__ T* outAddr, __ubuf__ T* varAddr, __ubuf__ T* updAddr,
242+ __ubuf__ int32_t* idxAddr, int32_t kStart, T zeroScalar, uint32_t rows,
243+ uint32_t innerElems, uint32_t lane, uint16_t repeatsPerRow)
244+{
245+ AscendC::Reg::RegTensor<int32_t> assistReg;
246+ AscendC::Reg::RegTensor<int32_t> idxReg;
247+ AscendC::Reg::MaskReg mask;
248+ AscendC::Reg::MaskReg hit;
249+ for (uint32_t m = 0; m < rows; ++m) {
250+ AscendC::Reg::Duplicate(assistReg, static_cast<int32_t>(kStart) + static_cast<int32_t>(m));
251+ uint32_t remain = innerElems; // UpdateMask 每轮自减, 尾块自动收窄
252+ uint32_t rowBase = m * innerElems;
253+ for (uint16_t t = 0; t < repeatsPerRow; ++t) {
254+ mask = AscendC::Reg::UpdateMask<int32_t>(remain);
255+ AscendC::Reg::LoadAlign(idxReg, idxAddr + t * lane);
256+ AscendC::Reg::Compare<int32_t, AscendC::CMPMODE::EQ>(hit, assistReg, idxReg, mask);
257+ if constexpr (IsSameType<T, float>::value || IsSameType<T, int32_t>::value) {
258+ SelectLane32<T, false>(outAddr + rowBase + t * lane, varAddr + rowBase + t * lane, updAddr + t * lane,
259+ zeroScalar, hit, mask);
260+ } else if constexpr (sizeof(T) == B16_BYTES) {
261+ SelectLaneB16<T, false>(outAddr + rowBase + t * lane, varAddr + rowBase + t * lane, updAddr + t * lane,
262+ 0.0f, hit, mask);
263+ } else {
264+ SelectLaneB8<T, false>(outAddr + rowBase + t * lane, varAddr + rowBase + t * lane, updAddr + t * lane,
265+ zeroScalar, hit, mask);
266+ }
267+ }
268+ }
269+}
270+ 
200} // namespace ArgMaxGrad271} // namespace ArgMaxGrad
201 272 
202#endif // ARG_MAX_GRAD_VF_H273#endif // ARG_MAX_GRAD_VF_H
Mindex/arg_max_grad/tests/ut/op_host/arch35/test_arg_max_grad_tiling.cpp+59-4
@@ -35,10 +35,13 @@ protected:
35};35};
36 36 
37static void InitPlatForm(fe::PlatFormInfos& platFormInfo, map<string, string>& socInfos,37static void InitPlatForm(fe::PlatFormInfos& platFormInfo, map<string, string>& socInfos,
38- map<string, string>& aicoreSpec, map<string, string>& intrinsics)38+ map<string, string>& aicoreSpec, map<string, string>& intrinsics, int64_t ubSize = 253952,
39+ int64_t coreNum = 64)
39{40{
41+ // 平台参数可覆盖: 用于验证 tiling 对异常平台信息的拒收(见 test_tiling_platform_* 三例)
40 string hardwareInfo = R"({42 string hardwareInfo = R"({
41- "hardware_info": {"UB_SIZE": 253952, "CORE_NUM": 64}43+ "hardware_info": {"UB_SIZE": )" +
44+ std::to_string(ubSize) + R"(, "CORE_NUM": )" + std::to_string(coreNum) + R"(}
42 })";45 })";
43 GetPlatFormInfos(hardwareInfo.c_str(), socInfos, aicoreSpec, intrinsics);46 GetPlatFormInfos(hardwareInfo.c_str(), socInfos, aicoreSpec, intrinsics);
44 platFormInfo.Init();47 platFormInfo.Init();
@@ -162,13 +165,14 @@ static ge::graphStatus RunTiling(TilingFunc tilingFunc, HolderT& holder, map<str
162}165}
163 166 
164static ge::graphStatus DoArgMaxGradTilingCase(const ArgMaxGradCase& c, uint64_t& tilingKey, uint32_t& blockDim,167static ge::graphStatus DoArgMaxGradTilingCase(const ArgMaxGradCase& c, uint64_t& tilingKey, uint32_t& blockDim,
165- ArgMaxGradArch35TilingData* tdOut = nullptr)168+ ArgMaxGradArch35TilingData* tdOut = nullptr, int64_t ubSize = 253952,
169+ int64_t coreNum = 64)
166{170{
167 fe::PlatFormInfos platFormInfo;171 fe::PlatFormInfos platFormInfo;
168 map<string, string> socInfos;172 map<string, string> socInfos;
169 map<string, string> aicoreSpec;173 map<string, string> aicoreSpec;
170 map<string, string> intrinsics;174 map<string, string> intrinsics;
171- InitPlatForm(platFormInfo, socInfos, aicoreSpec, intrinsics);175+ InitPlatForm(platFormInfo, socInfos, aicoreSpec, intrinsics, ubSize, coreNum);
172 176 
173 struct ArgMaxGradCompileInfo {};177 struct ArgMaxGradCompileInfo {};
174 ArgMaxGradCompileInfo compileInfo;178 ArgMaxGradCompileInfo compileInfo;
@@ -426,6 +430,57 @@ TEST_F(ArgMaxGradTiling, test_tiling_rank_full_range)
426}430}
427 431 
428// 多核切分: elemsPerCore 必须 32B 对齐(跨核不共享搬运块), 且能覆盖全部元素432// 多核切分: elemsPerCore 必须 32B 对齐(跨核不共享搬运块), 且能覆盖全部元素
433+// inner==1 且 D 很小、outer 很大: 每段能装 colsPerChunk/D 个 outer, indices/updates 各要
434+// 一个与之等长的标量缓冲, 但 bytesPerPoint 对 inner==1 把这两块记成 0 —— 段长按不含它们的
435+// 口径算满 UB 后再分配, 总量必然超。这类形状在支持面之内, 不该被 tiling 拒收。
436+TEST_F(ArgMaxGradTiling, test_tiling_inner_one_small_axis_large_outer)
437+{
438+ uint64_t tilingKey = 0;
439+ uint32_t blockDim = 0;
440+ ArgMaxGradArch35TilingData td{};
441+ for (int64_t d : {2, 4, 8}) {
442+ EXPECT_EQ(DoArgMaxGradTilingCase({{1000000, d, 1}, 1, ge::DT_FLOAT}, tilingKey, blockDim, &td),
443+ ge::GRAPH_SUCCESS)
444+ << "d=" << d;
445+ }
446+}
447+ 
448+// 平台信息异常时必须干净拒收(而不是拿非法值继续算)。这三条对应 tiling 里三处拒收分支:
449+// 核数为 0 / UB 不大于 SIMD-SIMT dcache 预留 / UB 扣掉预留后装不下一个向量整宽的元素
450+// rank 为 0(标量 var)必须拒收: 被选轴不存在, 语义不成立
451+TEST_F(ArgMaxGradTiling, test_tiling_rank0_rejected)
452+{
453+ uint64_t tilingKey = 0;
454+ uint32_t blockDim = 0;
455+ EXPECT_EQ(DoArgMaxGradTilingCase({{}, 0, ge::DT_FLOAT}, tilingKey, blockDim), ge::GRAPH_FAILED);
456+}
457+ 
458+TEST_F(ArgMaxGradTiling, test_tiling_platform_zero_core_rejected)
459+{
460+ uint64_t tilingKey = 0;
461+ uint32_t blockDim = 0;
462+ EXPECT_EQ(DoArgMaxGradTilingCase({{4, 16, 8}, 1, ge::DT_FLOAT}, tilingKey, blockDim, nullptr, 253952, 0),
463+ ge::GRAPH_FAILED);
464+}
465+ 
466+TEST_F(ArgMaxGradTiling, test_tiling_platform_ub_not_above_dcache_rejected)
467+{
468+ uint64_t tilingKey = 0;
469+ uint32_t blockDim = 0;
470+ // 32 KB 恰好等于 SIMD/SIMT 共用 dcache 的预留量, 扣完一个字节都不剩
471+ EXPECT_EQ(DoArgMaxGradTilingCase({{4, 16, 8}, 1, ge::DT_FLOAT}, tilingKey, blockDim, nullptr, 32 * 1024, 64),
472+ ge::GRAPH_FAILED);
473+}
474+ 
475+TEST_F(ArgMaxGradTiling, test_tiling_platform_ub_too_small_for_one_vector_rejected)
476+{
477+ uint64_t tilingKey = 0;
478+ uint32_t blockDim = 0;
479+ // 比 dcache 多 512 B: 过得了 Init 的门槛, 但一个向量整宽(64 个 fp32 元素)的各路 buffer 放不下
480+ EXPECT_EQ(DoArgMaxGradTilingCase({{4, 16, 8}, 1, ge::DT_FLOAT}, tilingKey, blockDim, nullptr, 32 * 1024 + 512, 64),
481+ ge::GRAPH_FAILED);
482+}
483+ 
429TEST_F(ArgMaxGradTiling, test_tiling_core_split_is_block_aligned)484TEST_F(ArgMaxGradTiling, test_tiling_core_split_is_block_aligned)
430{485{
431 uint64_t tilingKey = 0;486 uint64_t tilingKey = 0;