Star1012
1937
代码
代码
Issues194
Pull Requests332
Actions
项目讨论
Wiki
分析
Star1012
1937
已合并
feat(optim): lamb 七算子 arch35 适配，系数输入由标量放宽为可广播 Tensor #10786
zl_hw创建于 7 天前
feat(optim): lamb 七算子 arch35 适配，系数输入由标量放宽为可广播 Tensor #10786
已合并
zl_hw创建于 7 天前
从zl_hw/ops-nnfix/lamb-scalar-inputs-broadcast合入到cann/ops-nnmaster
共 66 个文件变更+2499-1105
Aoptim/lamb_apply_common/CMakeLists.txt+15-0
@@ -0,0 +1,15 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ----------------------------------------------------------------------------
10+ 
11+# 共用目录(非独立算子): 注册后 op_kernel/ 才会被拷入 build/tbe/ascendc/lamb_apply_common,
12+# 供各 lamb 算子以 "../lamb_apply_common/arch35/xxx.h" 引用(同 index/scatter_reduce_common)。
13+set(SUPPORT_COMPUTE_UNIT "ascend950")
14+set(SUPPORT_TILING_DIR "arch35")
15+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
Roptim/lamb_apply_common/lamb_apply_check_util.h→optim/lamb_apply_common/op_host/arch35/lamb_apply_check_util.h+45-14
@@ -10,19 +10,21 @@
10 10 
11/*!11/*!
12 * \file lamb_apply_check_util.h12 * \file lamb_apply_check_util.h
13- * \brief lamb_apply_optimizer_assign / lamb_apply_weight_assign 复用的输入校验(dtype 一致性、标量非空)。13+ * \brief lamb_apply_optimizer_assign / lamb_apply_weight_assign 复用的输入校验(dtype 一致性、ref 输出形状)。
14 */14 */
15#ifndef OPS_OPTIM_LAMB_APPLY_COMMON_CHECK_UTIL_H15#ifndef OPS_OPTIM_LAMB_APPLY_COMMON_CHECK_UTIL_H
16#define OPS_OPTIM_LAMB_APPLY_COMMON_CHECK_UTIL_H16#define OPS_OPTIM_LAMB_APPLY_COMMON_CHECK_UTIL_H
17 17 
18#include <cstddef>18#include <cstddef>
19#include <string>19#include <string>
20+#include <vector>
20#include "exe_graph/runtime/tiling_context.h"21#include "exe_graph/runtime/tiling_context.h"
22+#include "atvoss/broadcast/broadcast_tiling.h"
23+#include "infershape_broadcast_util.h"
21#include "log/log.h"24#include "log/log.h"
22 25 
23namespace optiling {26namespace optiling {
24 27 
25-// 所有输入(1..inputNum-1)、输出(0..outputNum-1)的 dtype 必须与 input0 一致。
26inline ge::graphStatus CheckLambApplyDtypeConsistency(gert::TilingContext* context, int32_t inputNum,28inline ge::graphStatus CheckLambApplyDtypeConsistency(gert::TilingContext* context, int32_t inputNum,
27 const char* const inputNames[], int32_t outputNum,29 const char* const inputNames[], int32_t outputNum,
28 const char* const outputNames[])30 const char* const outputNames[])
@@ -57,20 +59,49 @@ inline ge::graphStatus CheckLambApplyDtypeConsistency(gert::TilingContext* conte
57 return ge::GRAPH_SUCCESS;59 return ge::GRAPH_SUCCESS;
58}60}
59 61 
60-// 标量类输入(系数)为每元素计算所必需, 空Tensor 视为缺失必选值(畸形输入), 不支持。62+// ref(原地写回)输出所绑定的输入, 其形状必须恰好等于全部输入广播的结果:
61-inline ge::graphStatus CheckLambApplyScalarNotEmpty(gert::TilingContext* context, const int32_t* scalarIdx,63+// 内核按广播后的完整网格计算并写回该输入的 buffer, 形状不等就会越过它的显存边界。
62- size_t scalarCount, const char* const inputNames[])64+inline ge::graphStatus CheckLambApplyBroadcastIntoRef(gert::TilingContext* context, int32_t inputNum, int32_t refIdx,
65+ const char* refName)
63{66{
64- for (size_t i = 0; i < scalarCount; i++) {67+ std::vector<const gert::Shape*> inShapes;
65- int32_t idx = scalarIdx[i];68+ inShapes.reserve(static_cast<size_t>(inputNum));
66- auto scalarShape = context->GetInputShape(idx);69+ for (int32_t i = 0; i < inputNum; i++) {
67- OP_CHECK_NULL_WITH_CONTEXT(context, scalarShape);70+ auto inShape = context->GetInputShape(i);
68- if (scalarShape->GetStorageShape().GetShapeSize() == 0) {71+ OP_CHECK_NULL_WITH_CONTEXT(context, inShape);
69- OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context->GetNodeName(), inputNames[idx],72+ inShapes.push_back(&Ops::Base::EnsureNotScalar(inShape->GetStorageShape()));
70- Ops::Base::ToString(scalarShape->GetStorageShape()).c_str(),73+ }
71- "scalar input does not support empty tensor");74+ // 本地折叠广播, 不依赖 Ops::Base::BroadcastShape —— 该符号由 libops_base.so 导出, 但自定义
72- return ge::GRAPH_FAILED;75+ // vendor 包的 tiling 库不链它, 装包后 so 带未解析符号(ldd -r 可见)。带未解析符号的 so 早期
76+ // dlopen 会失败、被退到内置之后加载, 导致本算子的 tiling 模板抢不到注册槽位(实测 tilingKey
77+ // 拿到的是内置 ATVOSS 的值)。改为本地实现后 so 无未解析符号。
78+ gert::Shape bcShape = *inShapes[0];
79+ for (size_t i = 1; i < inShapes.size(); i++) {
80+ const gert::Shape& rhs = *inShapes[i];
81+ size_t lr = bcShape.GetDimNum();
82+ size_t rr = rhs.GetDimNum();
83+ size_t rank = (lr > rr) ? lr : rr;
84+ gert::Shape tmpShape;
85+ tmpShape.SetDimNum(rank);
86+ for (size_t d = 0; d < rank; d++) {
87+ // 右对齐: 缺的高位维按 1 处理
88+ int64_t a = (d + lr >= rank) ? bcShape.GetDim(d + lr - rank) : 1;
89+ int64_t b = (d + rr >= rank) ? rhs.GetDim(d + rr - rank) : 1;
90+ if (a != b && a != 1 && b != 1) {
91+ OP_LOGE(context->GetNodeName(), "input shapes cannot broadcast together");
92+ return ge::GRAPH_FAILED;
93+ }
94+ // 取"非 1 的那个", 不能取 max —— 空 Tensor 场景下 0 与 1 广播应得 0。
95+ tmpShape.SetDim(d, (a == 1) ? b : a);
73 }96 }
97+ bcShape = tmpShape;
98+ }
99+ const gert::Shape& refShape = *inShapes[refIdx];
100+ if (!(bcShape == refShape)) {
101+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
102+ context->GetNodeName(), refName, Ops::Base::ToString(refShape).c_str(),
103+ "it is an in-place(ref) output, so the broadcast shape of all inputs must equal it");
104+ return ge::GRAPH_FAILED;
74 }105 }
75 return ge::GRAPH_SUCCESS;106 return ge::GRAPH_SUCCESS;
76}107}
Aoptim/lamb_apply_common/op_host/arch35/lamb_brc_tiling_plan.h+320-0
@@ -0,0 +1,320 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_brc_tiling_plan.h
13+ * \brief LAMB 族「多入多出 + 任意 numpy 广播」逐元素算子的共用形状规划(host 侧)
14+ *
15+ * 折叠 -> 分类 -> 由 ubSize 反解分片 -> 选平铺/分块。全程无经验阈值:
16+ * 分片长度是 ubSize 除以槽位总字节解出来的; 行块轴是"能装下"这个条件由内向外推出来的,
17+ * 退化到最内层必然成立, 因此不存在因尺寸拒收的分支。
18+ */
19+ 
20+#ifndef LAMB_BRC_TILING_PLAN_H
21+#define LAMB_BRC_TILING_PLAN_H
22+ 
23+#include <algorithm>
24+#include <vector>
25+#include "../../op_kernel/arch35/lamb_brc_tiling_data.h"
26+#include "exe_graph/runtime/tiling_context.h"
27+#include "log/log.h"
28+ 
29+namespace optiling {
30+ 
31+// ---------------------------------------------------------------------------------------------
32+// BuildLambBrcPlan 的分步实现: 读形状 -> 折叠轴 -> 分类输入 -> 解分片 -> 平铺分核 /
33+// 分块(选行块轴 -> 分核 -> 算步长)。拆开是为了每段都短到能一眼读完、各自可单独复核。
34+// ---------------------------------------------------------------------------------------------
35+ 
36+// 读出各输入形状, 右对齐补维, 并校验可广播性。
37+template <uint32_t NIN>
38+ge::graphStatus LambBrcReadShapes(gert::TilingContext* context, size_t rank, const std::vector<int64_t>& rawOut,
39+ std::vector<std::vector<int64_t>>& rawIn)
40+{
41+ for (uint32_t i = 0; i < NIN; i++) {
42+ auto sp = context->GetInputShape(i);
43+ OP_CHECK_NULL_WITH_CONTEXT(context, sp);
44+ const gert::Shape& s = sp->GetStorageShape();
45+ size_t r = s.GetDimNum();
46+ OP_CHECK_IF(r > rank, OP_LOGE(context->GetNodeName(), "input %u rank %zu exceeds output rank %zu", i, r, rank),
47+ return ge::GRAPH_FAILED);
48+ for (size_t d = 0; d < r; d++) { // 右对齐补维
49+ rawIn[i][rank - r + d] = s.GetDim(d);
50+ }
51+ for (size_t d = 0; d < rank; d++) {
52+ OP_CHECK_IF(rawIn[i][d] != 1 && rawIn[i][d] != rawOut[d],
53+ OP_LOGE(context->GetNodeName(), "input %u dim %zu (%ld) not broadcastable to %ld", i, d,
54+ rawIn[i][d], rawOut[d]),
55+ return ge::GRAPH_FAILED);
56+ }
57+ }
58+ return ge::GRAPH_SUCCESS;
59+}
60+ 
61+// 折叠: 丢掉输出为 1 的轴, 再合并广播状态一致的相邻轴。
62+template <uint32_t NIN>
63+void LambBrcCollapse(const std::vector<int64_t>& rawOut, const std::vector<std::vector<int64_t>>& rawIn, size_t rank,
64+ std::vector<int64_t>& co, std::vector<std::vector<int64_t>>& ci)
65+{
66+ // 折叠: 丢掉输出为 1 的轴, 再合并广播状态一致的相邻轴。
67+ for (size_t d = 0; d < rank; d++) {
68+ if (rawOut[d] == 1) {
69+ continue;
70+ }
71+ bool merge = !co.empty();
72+ for (uint32_t i = 0; i < NIN && merge; i++) {
73+ merge = ((ci[i].back() == 1) == (rawIn[i][d] == 1));
74+ }
75+ if (merge) {
76+ co.back() *= rawOut[d];
77+ for (uint32_t i = 0; i < NIN; i++) {
78+ ci[i].back() *= rawIn[i][d];
79+ }
80+ } else {
81+ co.push_back(rawOut[d]);
82+ for (uint32_t i = 0; i < NIN; i++) {
83+ ci[i].push_back(rawIn[i][d]);
84+ }
85+ }
86+ }
87+ if (co.empty()) { // 全 1 形状
88+ co.push_back(1);
89+ for (uint32_t i = 0; i < NIN; i++) {
90+ ci[i].push_back(1);
91+ }
92+ }
93+}
94+ 
95+// 给每个输入定性: 标量 / 与输出同形 / 需要广播。返回是否存在需要广播的输入。
96+template <uint32_t NIN, uint32_t NOUT>
97+bool LambBrcClassify(const std::vector<int64_t>& co, const std::vector<std::vector<int64_t>>& ci, uint32_t rc,
98+ LambBrcTilingData<NIN, NOUT>& td)
99+{
100+ bool anyBrc = false;
101+ for (uint32_t i = 0; i < NIN; i++) {
102+ uint64_t numel = 1;
103+ bool brc = false;
104+ for (uint32_t d = 0; d < rc; d++) {
105+ td.inShape[i * LAMB_BRC_MAX_DIM + d] = static_cast<uint32_t>(ci[i][d]);
106+ numel *= static_cast<uint64_t>(ci[i][d]);
107+ if (ci[i][d] != co[d]) {
108+ brc = true;
109+ }
110+ }
111+ if (numel == 1) {
112+ td.inKind[i] = LAMB_BRC_KIND_SCALAR;
113+ } else if (!brc) {
114+ td.inKind[i] = LAMB_BRC_KIND_SAME;
115+ } else {
116+ td.inKind[i] = LAMB_BRC_KIND_BRC;
117+ anyBrc = true;
118+ }
119+ }
120+ return anyBrc;
121+}
122+ 
123+// 由 ubSize 反解分片长度。
124+template <uint32_t NIN, uint32_t NOUT>
125+uint32_t LambBrcSolveTileLen(uint64_t ubSize, uint32_t dtSize)
126+{
127+ // 分片长度: UB 总字节 / 槽位数 / 元素字节, 向下对齐到 256B(保证每槽首址对齐)。
128+ constexpr uint32_t slotNum = LambBrcTilingData<NIN, NOUT>::SlotNum();
129+ uint32_t alignElems = LAMB_BRC_VREG_BYTES / dtSize;
130+ // 从 ubSize 反解, 再向下对齐到一个向量寄存器; 解不出整数倍时兜底取一个对齐单位 ——
131+ // 不做"装不下就拒收"的判断, 装不下就继续切, 这是切分问题不是能力问题。
132+ uint64_t bytesPerElem = static_cast<uint64_t>(slotNum) * dtSize + 4U * sizeof(float);
133+ uint32_t tileLen = static_cast<uint32_t>((ubSize / bytesPerElem / alignElems) * alignElems);
134+ if (tileLen == 0) {
135+ tileLen = alignElems;
136+ }
137+ return tileLen;
138+}
139+ 
140+// 无广播输入: 按元素平铺分核。
141+template <uint32_t NIN, uint32_t NOUT>
142+void LambBrcPlanFlat(uint64_t coreNum, uint32_t dtSize, uint64_t total, uint32_t tileLen,
143+ LambBrcTilingData<NIN, NOUT>& td)
144+{
145+ // 分核粒度必须让每个核的输出起始落在 32B 边界上: 否则相邻核的 DataCopyPad 会写到
146+ // 同一个 32 字节块里互相覆盖。
147+ const uint64_t gmAlign = LAMB_BRC_GM_BLOCK_BYTES / dtSize;
148+ 
149+ td.tilingKey = LAMB_BRC_KEY_FLAT;
150+ uint64_t tiles = (total + tileLen - 1) / tileLen;
151+ uint64_t cores = std::min<uint64_t>(coreNum, std::max<uint64_t>(tiles, 1));
152+ uint64_t perCore = (total + cores - 1) / cores;
153+ perCore = ((perCore + gmAlign - 1) / gmAlign) * gmAlign; // 对齐到 32B
154+ td.perCoreElems = perCore;
155+ td.usedCoreNum = static_cast<uint32_t>((total + perCore - 1) / perCore);
156+}
157+ 
158+// 分块: 选行块轴(splitAxis)与块长(blockLen)。
159+template <uint32_t NIN, uint32_t NOUT>
160+void LambBrcChooseSplit(const std::vector<int64_t>& co, const std::vector<std::vector<int64_t>>& ci, uint32_t rc,
161+ uint32_t tileLen, LambBrcTilingData<NIN, NOUT>& td, uint32_t& splitAxis, uint64_t& blockLen)
162+{
163+ splitAxis = rc;
164+ blockLen = 1;
165+ // 块内只允许"尾轴广播"这一种形态(那条 2D Broadcast 路径是验证过的)。任何非尾轴的广播轴
166+ // 都把 splitAxis 推到它之后, 使其变成行维、由 effStride=0 处理 —— 否则块内要在任意元素偏移
167+ // 上做 UB->UB 复制, 而向量指令对起始地址有对齐要求。
168+ uint32_t splitFloor = 0;
169+ for (uint32_t i = 0; i < NIN; i++) {
170+ if (td.inKind[i] != LAMB_BRC_KIND_BRC) {
171+ continue;
172+ }
173+ for (uint32_t d = 0; d + 1 < rc; d++) { // 只看非尾轴
174+ if (ci[i][d] == 1 && co[d] != 1 && d + 1 > splitFloor) {
175+ splitFloor = d + 1;
176+ }
177+ }
178+ }
179+ 
180+ // 分块: 由内向外找最小的前导行维个数(不低于 splitFloor), 使一个行块装得下分片。
181+ for (int32_t k = static_cast<int32_t>(rc); k >= static_cast<int32_t>(splitFloor); k--) {
182+ uint64_t len = 1;
183+ for (uint32_t d = static_cast<uint32_t>(k); d < rc; d++) {
184+ len *= static_cast<uint64_t>(co[d]);
185+ }
186+ if (len > tileLen) {
187+ break;
188+ }
189+ splitAxis = static_cast<uint32_t>(k);
190+ blockLen = len;
191+ }
192+ // 最内轴自己就装不下一个分片时(blockLen 被逼到 1), 不新增轴, 而是在最内轴内部切段:
193+ // 行索引的最低位变成"段号", 块长取最内轴不超过分片长度的最大因子。轴数不变, 不设上限。
194+ if (blockLen == 1 && rc > 0 && static_cast<uint64_t>(co[rc - 1]) > 1) {
195+ int64_t d = co[rc - 1];
196+ int64_t chunk = 1;
197+ for (int64_t c = static_cast<int64_t>(tileLen); c >= 1; c--) {
198+ if (d % c == 0) {
199+ chunk = c;
200+ break;
201+ }
202+ }
203+ // 最内轴为大质数时取不到 >1 的因子, 此时保持逐元素: 结果仍正确(行批处理保证搬出连续),
204+ // 只是 DMA 次数多, 属性能退化, 不拒收。
205+ if (chunk > 1) {
206+ splitAxis = rc - 1;
207+ blockLen = static_cast<uint64_t>(chunk);
208+ td.innerChunk = static_cast<uint32_t>(chunk);
209+ td.innerChunkCnt = static_cast<uint32_t>(d / chunk);
210+ }
211+ }
212+}
213+ 
214+// 分块模式的分核: 行批大小与每核行数。
215+template <uint32_t NIN, uint32_t NOUT>
216+void LambBrcCoreSplit(uint64_t coreNum, uint32_t dtSize, uint64_t total, uint32_t tileLen, uint32_t splitAxis,
217+ uint64_t blockLen, LambBrcTilingData<NIN, NOUT>& td)
218+{
219+ const uint64_t gmAlign = LAMB_BRC_GM_BLOCK_BYTES / dtSize;
220+ td.tilingKey = LAMB_BRC_KEY_BLOCK;
221+ td.splitAxis = splitAxis;
222+ td.blockLen = static_cast<uint32_t>(blockLen);
223+ // 一个分片内批量攒多行再一次性搬出: 逐行搬出会退化成 4/2 字节的 DataCopyPad,
224+ // 同一个 32 字节块内的多次写互相覆盖。
225+ // 行批处理要求每行在 UB 里的起始偏移 j*blockLen 落在 32 字节边界上, 否则 DataCopyPad 的
226+ // UB 侧地址不对齐。blockLen 不是 32 字节整数倍时退回每次一行(偏移恒为 0)。
227+ // 跨核不冲突由 rowsPerCore 的 32 字节对齐保证, 与本项无关。
228+ bool blockAligned = ((blockLen * dtSize) % LAMB_BRC_GM_BLOCK_BYTES) == 0;
229+ td.rowsPerTile = blockAligned ? static_cast<uint32_t>(std::max<uint64_t>(tileLen / blockLen, 1)) : 1U;
230+ td.totalRows = total / blockLen;
231+ // 每核行数取到 rowsPerGroup 的整数倍, 使 rowsPerCore*blockLen 是 32B 的整数倍。
232+ uint64_t g = blockLen % gmAlign;
233+ uint64_t a = gmAlign;
234+ while (g != 0) {
235+ uint64_t t2 = a % g;
236+ a = g;
237+ g = t2;
238+ } // gcd(blockLen, gmAlign)
239+ uint64_t rowsPerGroup = gmAlign / a;
240+ uint64_t cores = std::min<uint64_t>(coreNum, std::max<uint64_t>(td.totalRows, 1));
241+ uint64_t rowsPerCore = (td.totalRows + cores - 1) / cores;
242+ rowsPerCore = ((rowsPerCore + rowsPerGroup - 1) / rowsPerGroup) * rowsPerGroup;
243+ td.rowsPerCore = rowsPerCore;
244+ td.usedCoreNum = static_cast<uint32_t>((td.totalRows + rowsPerCore - 1) / rowsPerCore);
245+}
246+ 
247+// 每个输入的块内长度与行维步长(广播轴步长为 0)。
248+template <uint32_t NIN, uint32_t NOUT>
249+void LambBrcStrides(const std::vector<std::vector<int64_t>>& ci, uint32_t rc, uint32_t splitAxis,
250+ LambBrcTilingData<NIN, NOUT>& td)
251+{
252+ for (uint32_t i = 0; i < NIN; i++) {
253+ uint64_t srcBlock = 1;
254+ for (uint32_t d = splitAxis; d < rc; d++) {
255+ srcBlock *= static_cast<uint64_t>(ci[i][d]);
256+ }
257+ if (td.innerChunkCnt > 1) {
258+ // 块只覆盖最内轴的一段: 该输入在这一轴上要么是 1(广播), 要么与输出等长。
259+ srcBlock = (ci[i][rc - 1] == 1) ? 1 : static_cast<uint64_t>(td.innerChunk);
260+ }
261+ td.srcBlockLen[i] = srcBlock;
262+ for (uint32_t j = 0; j < splitAxis; j++) {
263+ if (ci[i][j] == 1) {
264+ td.effStride[i * LAMB_BRC_MAX_DIM + j] = 0; // 广播轴不推进源地址
265+ continue;
266+ }
267+ uint64_t st = 1;
268+ for (uint32_t d = j + 1; d < rc; d++) {
269+ st *= static_cast<uint64_t>(ci[i][d]);
270+ }
271+ td.effStride[i * LAMB_BRC_MAX_DIM + j] = st;
272+ }
273+ }
274+}
275+ 
276+template <uint32_t NIN, uint32_t NOUT>
277+ge::graphStatus BuildLambBrcPlan(gert::TilingContext* context, uint64_t coreNum, uint64_t ubSize, uint32_t dtSize,
278+ LambBrcTilingData<NIN, NOUT>& td)
279+{
280+ auto outShapePtr = context->GetOutputShape(0);
281+ OP_CHECK_NULL_WITH_CONTEXT(context, outShapePtr);
282+ const gert::Shape& outShape = outShapePtr->GetStorageShape();
283+ size_t rank = outShape.GetDimNum();
284+ std::vector<int64_t> rawOut(rank);
285+ for (size_t d = 0; d < rank; d++) {
286+ rawOut[d] = outShape.GetDim(d);
287+ }
288+ std::vector<std::vector<int64_t>> rawIn(NIN, std::vector<int64_t>(rank, 1));
289+ OP_CHECK_IF(LambBrcReadShapes<NIN>(context, rank, rawOut, rawIn) != ge::GRAPH_SUCCESS,
290+ OP_LOGE(context->GetNodeName(), "read input shapes failed"), return ge::GRAPH_FAILED);
291+ 
292+ std::vector<int64_t> co;
293+ std::vector<std::vector<int64_t>> ci(NIN);
294+ LambBrcCollapse<NIN>(rawOut, rawIn, rank, co, ci);
295+ uint32_t rc = static_cast<uint32_t>(co.size());
296+ td.collapsedRank = rc;
297+ uint64_t total = 1;
298+ for (uint32_t d = 0; d < rc; d++) {
299+ td.outShape[d] = static_cast<uint32_t>(co[d]);
300+ total *= static_cast<uint64_t>(co[d]);
301+ }
302+ td.totalNum = total;
303+ 
304+ bool anyBrc = LambBrcClassify<NIN, NOUT>(co, ci, rc, td);
305+ uint32_t tileLen = LambBrcSolveTileLen<NIN, NOUT>(ubSize, dtSize);
306+ td.tileLen = tileLen;
307+ if (!anyBrc) {
308+ LambBrcPlanFlat<NIN, NOUT>(coreNum, dtSize, total, tileLen, td);
309+ return ge::GRAPH_SUCCESS;
310+ }
311+ 
312+ uint32_t splitAxis = rc;
313+ uint64_t blockLen = 1;
314+ LambBrcChooseSplit<NIN, NOUT>(co, ci, rc, tileLen, td, splitAxis, blockLen);
315+ LambBrcCoreSplit<NIN, NOUT>(coreNum, dtSize, total, tileLen, splitAxis, blockLen, td);
316+ LambBrcStrides<NIN, NOUT>(ci, rc, splitAxis, td);
317+ return ge::GRAPH_SUCCESS;
318+}
319+} // namespace optiling
320+#endif // LAMB_BRC_TILING_PLAN_H
Aoptim/lamb_apply_common/op_kernel/arch35/lamb_brc_kernel.h+346-0
@@ -0,0 +1,346 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_brc_kernel.h
13+ * \brief LAMB 族「多入多出 + 任意 numpy 广播」逐元素算子的共用搬运/分核/广播骨架
14+ *
15+ * 为什么不走 ATVOSS: 本族算子入参 12~13 个, DAGSch 在
16+ * (mte2Count + mte3Count) * BUF_PING_PONG + tempBufCount <= 32 上编不过;
17+ * 且 ATVOSS 的 Vec::Brc 未实现, UB 广播档不可用, 只剩 NDDMA 档, 铺不平
18+ * 「尾轴 1->n」与「低秩补维」两类广播。
19+ *
20+ * 本骨架:
21+ * - 广播在搬入阶段铺平, 计算段只看逐元素数据; 铺平只用平台已验证的 2D
22+ * AscendC::Broadcast 由内向外迭代展开, 因此支持任意 numpy 广播形态。
23+ * - 分片长度/行块轴由 host 从 ubSize 反解下发, 内核不做"算完再跟 UB 比大小"的判断,
24+ * 也没有经验阈值; 行块轴退化到最内层必然成立, 故不存在因尺寸拒收的分支。
25+ * - 计算段由算子自己以 Compute::Run(in[], out[], count) 提供。
26+ */
27+ 
28+#ifndef LAMB_BRC_KERNEL_H
29+#define LAMB_BRC_KERNEL_H
30+ 
31+#include "kernel_operator.h"
32+#include "op_kernel/platform_util.h"
33+#include "lamb_brc_tiling_data.h"
34+ 
35+namespace LambBrc {
36+using namespace AscendC;
37+ 
38+static constexpr Reg::CastTrait LAMB_CAST_UP = {Reg::RegLayout::ZERO, Reg::SatMode::UNKNOWN,
39+ Reg::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
40+static constexpr Reg::CastTrait LAMB_CAST_DOWN = {Reg::RegLayout::ZERO, Reg::SatMode::NO_SAT,
41+ Reg::MaskMergeMode::ZEROING, RoundMode::CAST_RINT};
42+ 
43+constexpr uint32_t LAMB_VF_LEN = Ops::Base::GetVRegSize() / sizeof(float);
44+ 
45+// host 侧无法调用 __aicore__ 的平台接口, 只能在共用头里写常量; 这里回连校验, 防止两侧走偏。
46+static_assert(Ops::Base::GetUbBlockSize() == LAMB_BRC_GM_BLOCK_BYTES, "UB block size mismatch");
47+static_assert(Ops::Base::GetVRegSize() == LAMB_BRC_VREG_BYTES, "vector register size mismatch");
48+ 
49+// 精度档: 默认的 DivAlgo/SqrtAlgo::INTRINSIC 是硬件近似实现, 且按 FTZ(非规格数归零)处理 ——
50+// 非规格中间量被冲刷成 ±0 后, sqrt(-tiny) 应得的 NaN 会退化成 sqrt(-0)=-0, 再除即成 inf,
51+// 与 CPU 参照定性不符。取仓内通行的 0ULP + FTZ_FALSE 档(见 index/scatter_reduce_common、
52+// loss/soft_margin_loss_grad 等), 由库完成非规格数的完整处理。
53+static constexpr AscendC::Reg::DivSpecificMode LAMB_PRECISE_DIV = {AscendC::Reg::MaskMergeMode::ZEROING, false,
54+ AscendC::DivAlgo::PRECISION_0ULP_FTZ_FALSE};
55+static constexpr AscendC::Reg::SqrtSpecificMode LAMB_PRECISE_SQRT = {AscendC::Reg::MaskMergeMode::ZEROING, false,
56+ AscendC::SqrtAlgo::PRECISION_0ULP_FTZ_FALSE};
57+static constexpr AscendC::Reg::ExpSpecificMode LAMB_PRECISE_EXP = {AscendC::Reg::MaskMergeMode::ZEROING,
58+ AscendC::ExpAlgo::PRECISION_1ULP_FTZ_FALSE};
59+static constexpr AscendC::Reg::LnSpecificMode LAMB_PRECISE_LN = {AscendC::Reg::MaskMergeMode::ZEROING,
60+ AscendC::LnAlgo::PRECISION_1ULP_FTZ_FALSE};
61+// FTZ_TRUE 档: exp 的下溢结果按冲刷到 ±0 处理。ln(b)*steps 小于 fp32 exp 的下溢边界时,
62+// b^steps 的数学真值本就小于最小非规格数, 冲刷到 0 即是正确结果(b_corr=1)。
63+static constexpr AscendC::Reg::ExpSpecificMode LAMB_PRECISE_EXP_FTZT = {AscendC::Reg::MaskMergeMode::ZEROING,
64+ AscendC::ExpAlgo::PRECISION_1ULP_FTZ_TRUE};
65+static constexpr AscendC::Reg::LnSpecificMode LAMB_PRECISE_LN_FTZT = {AscendC::Reg::MaskMergeMode::ZEROING,
66+ AscendC::LnAlgo::PRECISION_1ULP_FTZ_TRUE};
67+ 
68+// corr = 1 - exp(x)。x -> 0 时 exp(x) -> 1, 直写 1-exp(x) 是灾难性抵消(fp32 在 1.0 附近
69+// 间距 6e-8), 与 exp 实现多准无关。故 |x| < 0.1 走 5 项 Horner 多项式直算 expm1(x) 绕开该减法,
70+// 否则直算 exp(x)-1。同式同阈值见 activation/elu。
71+__aicore__ inline void OneMinusExp(Reg::RegTensor<float>& dst, Reg::RegTensor<float>& x, Reg::RegTensor<float>& t1,
72+ Reg::RegTensor<float>& t2, Reg::MaskReg& cmp, Reg::MaskReg& preg)
73+{
74+ Reg::Muls(t1, x, 1.0f / 120.0f, preg);
75+ Reg::Adds(t1, t1, 1.0f / 24.0f, preg);
76+ Reg::Mul(t1, t1, x, preg);
77+ Reg::Adds(t1, t1, 1.0f / 6.0f, preg);
78+ Reg::Mul(t1, t1, x, preg);
79+ Reg::Adds(t1, t1, 0.5f, preg);
80+ Reg::Mul(t1, t1, x, preg);
81+ Reg::Adds(t1, t1, 1.0f, preg);
82+ Reg::Mul(t1, t1, x, preg); // t1 = expm1(x), 多项式支
83+ Reg::Exp<float, &LAMB_PRECISE_EXP_FTZT>(t2, x, preg);
84+ Reg::Adds(t2, t2, -1.0f, preg); // t2 = expm1(x), exp 支
85+ Reg::Abs(dst, x, preg); // x 之后不再用, dst 可与 x 同寄存器
86+ Reg::Compares<float, CMPMODE::LT>(cmp, dst, 0.1f, preg);
87+ Reg::Select<float>(dst, t1, t2, cmp);
88+ Reg::Muls(dst, dst, -1.0f, preg); // 1 - exp(x) = -expm1(x)
89+}
90+ 
91+template <typename T>
92+__aicore__ inline void Load(__local_mem__ T* base, Reg::RegTensor<float>& dst, Reg::MaskReg& preg, uint32_t offset)
93+{
94+ if constexpr (std::is_same_v<T, float>) {
95+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(dst, base + offset);
96+ } else {
97+ Reg::RegTensor<T> narrow;
98+ Reg::LoadAlign<T, Reg::LoadDist::DIST_UNPACK_B16>(narrow, base + offset);
99+ Reg::Cast<float, T, LAMB_CAST_UP>(dst, narrow, preg);
100+ }
101+}
102+ 
103+template <typename T>
104+__aicore__ inline void Store(__local_mem__ T* base, Reg::RegTensor<float>& src, Reg::MaskReg& preg, uint32_t offset)
105+{
106+ if constexpr (std::is_same_v<T, float>) {
107+ Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM>(base + offset, src, preg);
108+ } else {
109+ Reg::RegTensor<T> narrow;
110+ Reg::Cast<T, float, LAMB_CAST_DOWN>(narrow, src, preg);
111+ Reg::StoreAlign<T, Reg::StoreDist::DIST_PACK_B32>(base + offset, narrow, preg);
112+ }
113+}
114+ 
115+template <typename T, uint32_t NIN, uint32_t NOUT, class Compute>
116+class BrcElementwiseKernel {
117+public:
118+ using Tiling = LambBrcTilingData<NIN, NOUT>;
119+ static constexpr uint32_t SLOT_NUM = NIN + NOUT + 1;
120+ static constexpr uint32_t TMP_SLOT = SLOT_NUM - 1;
121+ 
122+ __aicore__ inline BrcElementwiseKernel() {}
123+ 
124+ __aicore__ inline void Init(GM_ADDR (&inAddr)[NIN], GM_ADDR (&outAddr)[NOUT], const Tiling* tiling, TPipe* pipe)
125+ {
126+ tiling_ = tiling;
127+ blockIdx_ = GetBlockIdx();
128+ for (uint32_t i = 0; i < NIN; i++) {
129+ inGm_[i].SetGlobalBuffer((__gm__ T*)inAddr[i]);
130+ }
131+ for (uint32_t o = 0; o < NOUT; o++) {
132+ outGm_[o].SetGlobalBuffer((__gm__ T*)outAddr[o]);
133+ }
134+ // 槽位 + fp32 暂存区(4*tileLen): 计算段过长时可拆成两个 __VEC_SCOPE__,
135+ // 中间量用 fp32 暂存传递, 与单段版数值完全一致(不经 T 的舍入)。
136+ pipe->InitBuffer(ubBuf_, SLOT_NUM * tiling_->tileLen * sizeof(T) + 4U * tiling_->tileLen * sizeof(float));
137+ LocalTensor<T> all = ubBuf_.Get<T>();
138+ for (uint32_t s = 0; s < SLOT_NUM; s++) {
139+ slot_[s] = all[s * tiling_->tileLen];
140+ }
141+ scratch_ = (__local_mem__ float*)((__local_mem__ uint8_t*)all.GetPhyAddr() +
142+ SLOT_NUM * tiling_->tileLen * sizeof(T));
143+ FillScalarSlots();
144+ }
145+ 
146+ __aicore__ inline void Process()
147+ {
148+ if (blockIdx_ >= tiling_->usedCoreNum) {
149+ return;
150+ }
151+ if (tiling_->tilingKey == LAMB_BRC_KEY_FLAT) {
152+ ProcessFlat();
153+ } else {
154+ ProcessBlock();
155+ }
156+ }
157+ 
158+private:
159+ // 单元素输入是常量, 只铺一次, 后续分片不再重复搬运。
160+ __aicore__ inline void FillScalarSlots()
161+ {
162+ event_t e2s = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE2_S>());
163+ event_t s2v = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::S_V>());
164+ event_t v2m2 = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>());
165+ for (uint32_t i = 0; i < NIN; i++) {
166+ if (tiling_->inKind[i] != LAMB_BRC_KIND_SCALAR) {
167+ continue;
168+ }
169+ DataCopyExtParams cp{1, static_cast<uint32_t>(sizeof(T)), 0, 0, 0};
170+ DataCopyPadExtParams<T> pad{false, 0, 0, 0};
171+ DataCopyPad(slot_[TMP_SLOT], inGm_[i], cp, pad);
172+ SetFlag<HardEvent::MTE2_S>(e2s);
173+ WaitFlag<HardEvent::MTE2_S>(e2s);
174+ T v = slot_[TMP_SLOT].GetValue(0);
175+ SetFlag<HardEvent::S_V>(s2v);
176+ WaitFlag<HardEvent::S_V>(s2v);
177+ Duplicate(slot_[i], v, static_cast<int32_t>(tiling_->tileLen));
178+ // 暂存槽下一轮会被 MTE2 覆写, 需等本轮标量读取(S)与 Duplicate(V) 完成。
179+ SetFlag<HardEvent::V_MTE2>(v2m2);
180+ WaitFlag<HardEvent::V_MTE2>(v2m2);
181+ }
182+ GetTPipePtr()->ReleaseEventID<HardEvent::MTE2_S>(e2s);
183+ GetTPipePtr()->ReleaseEventID<HardEvent::S_V>(s2v);
184+ GetTPipePtr()->ReleaseEventID<HardEvent::V_MTE2>(v2m2);
185+ PipeBarrier<PIPE_V>();
186+ }
187+ 
188+ __aicore__ inline void ProcessFlat()
189+ {
190+ uint64_t total = tiling_->totalNum;
191+ uint64_t perCore = tiling_->perCoreElems;
192+ uint64_t begin = perCore * blockIdx_;
193+ uint64_t end = (begin + perCore < total) ? (begin + perCore) : total;
194+ for (uint64_t off = begin; off < end; off += tiling_->tileLen) {
195+ uint32_t len = static_cast<uint32_t>(((end - off) < tiling_->tileLen) ? (end - off) : tiling_->tileLen);
196+ for (uint32_t i = 0; i < NIN; i++) {
197+ if (tiling_->inKind[i] == LAMB_BRC_KIND_SAME) {
198+ CopyInRun(i, off, len);
199+ }
200+ }
201+ ComputeAndOut(off, len);
202+ }
203+ }
204+ 
205+ __aicore__ inline void ProcessBlock()
206+ {
207+ uint64_t rows = tiling_->totalRows;
208+ uint64_t perCore = tiling_->rowsPerCore;
209+ uint64_t begin = perCore * blockIdx_;
210+ uint64_t end = (begin + perCore < rows) ? (begin + perCore) : rows;
211+ uint32_t blockLen = tiling_->blockLen;
212+ uint32_t rowsPerTile = tiling_->rowsPerTile;
213+ 
214+ for (uint64_t r0 = begin; r0 < end; r0 += rowsPerTile) {
215+ uint64_t batch = ((end - r0) < rowsPerTile) ? (end - r0) : rowsPerTile;
216+ for (uint64_t j = 0; j < batch; j++) {
217+ uint32_t dstOff = static_cast<uint32_t>(j) * blockLen;
218+ for (uint32_t i = 0; i < NIN; i++) {
219+ uint32_t kind = tiling_->inKind[i];
220+ if (kind == LAMB_BRC_KIND_SCALAR) {
221+ continue; // 标量槽已整片铺好, 各行共用
222+ }
223+ uint64_t srcOff = RowSrcOffset(i, r0 + j);
224+ if (kind == LAMB_BRC_KIND_SAME) {
225+ CopyInRun(i, srcOff, blockLen, dstOff);
226+ } else {
227+ CopyInBrc(i, srcOff, dstOff);
228+ }
229+ }
230+ }
231+ ComputeAndOut(r0 * blockLen, static_cast<uint32_t>(batch) * blockLen);
232+ }
233+ }
234+ 
235+ // 行索引按输出前导维分解, 再乘各输入的有效步长(广播轴步长为 0)。
236+ __aicore__ inline uint64_t RowSrcOffset(uint32_t i, uint64_t r)
237+ {
238+ uint64_t off = 0;
239+ uint64_t rem = r;
240+ if (tiling_->innerChunkCnt > 1) { // 行索引最低位是最内轴的段号
241+ uint64_t seg = rem % tiling_->innerChunkCnt;
242+ rem /= tiling_->innerChunkCnt;
243+ if (tiling_->inShape[i * LAMB_BRC_MAX_DIM + tiling_->collapsedRank - 1] != 1) {
244+ off += seg * tiling_->innerChunk;
245+ }
246+ }
247+ for (int32_t j = static_cast<int32_t>(tiling_->splitAxis) - 1; j >= 0; j--) {
248+ uint64_t d = tiling_->outShape[j];
249+ off += (rem % d) * tiling_->effStride[i * LAMB_BRC_MAX_DIM + j];
250+ rem /= d;
251+ }
252+ return off;
253+ }
254+ 
255+ __aicore__ inline void CopyInRun(uint32_t i, uint64_t srcOff, uint32_t len, uint32_t dstOff = 0)
256+ {
257+ DataCopyExtParams cp{1, static_cast<uint32_t>(len * sizeof(T)), 0, 0, 0};
258+ DataCopyPadExtParams<T> pad{false, 0, 0, 0};
259+ DataCopyPad(slot_[i][dstOff], inGm_[i][srcOff], cp, pad);
260+ }
261+ 
262+ // 块内只可能是"尾轴 1->n"广播(host 的 splitFloor 保证), 或完全不需要展开。
263+ __aicore__ inline void CopyInBrc(uint32_t i, uint64_t srcOff, uint32_t dstOff = 0)
264+ {
265+ uint32_t blockLen = tiling_->blockLen;
266+ uint32_t srcLen = static_cast<uint32_t>(tiling_->srcBlockLen[i]);
267+ if (srcLen == blockLen) { // 块内无广播: 直接搬入, 不经暂存
268+ CopyInRun(i, srcOff, blockLen, dstOff);
269+ return;
270+ }
271+ // 上一次(同一行的前一个广播输入, 或前一行)的 Broadcast 还在读 TMP_SLOT, 必须等它读完
272+ // 再往里搬下一份 —— 否则 MTE2 的写越过 V 的读(WAR), 结果随机少量出错。
273+ event_t v2m2 = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>());
274+ SetFlag<HardEvent::V_MTE2>(v2m2);
275+ WaitFlag<HardEvent::V_MTE2>(v2m2);
276+ GetTPipePtr()->ReleaseEventID<HardEvent::V_MTE2>(v2m2);
277+ 
278+ DataCopyExtParams cp{1, static_cast<uint32_t>(srcLen * sizeof(T)), 0, 0, 0};
279+ DataCopyPadExtParams<T> pad{false, 0, 0, 0};
280+ DataCopyPad(slot_[TMP_SLOT], inGm_[i][srcOff], cp, pad);
281+ 
282+ event_t e2v = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE2_V>());
283+ SetFlag<HardEvent::MTE2_V>(e2v);
284+ WaitFlag<HardEvent::MTE2_V>(e2v);
285+ GetTPipePtr()->ReleaseEventID<HardEvent::MTE2_V>(e2v);
286+ 
287+ // [srcLen, 1] -> [srcLen, d]: 尾轴放大 d 倍
288+ uint32_t d = blockLen / srcLen;
289+ uint32_t dstShape[2] = {srcLen, d};
290+ uint32_t srcShape[2] = {srcLen, 1};
291+ Broadcast<T, 2, 1>(slot_[i][dstOff], slot_[TMP_SLOT], dstShape, srcShape);
292+ PipeBarrier<PIPE_V>();
293+ }
294+ 
295+ __aicore__ inline void ComputeAndOut(uint64_t dstOff, uint32_t len)
296+ {
297+ event_t e2v = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE2_V>());
298+ SetFlag<HardEvent::MTE2_V>(e2v);
299+ WaitFlag<HardEvent::MTE2_V>(e2v);
300+ GetTPipePtr()->ReleaseEventID<HardEvent::MTE2_V>(e2v);
301+ 
302+ __local_mem__ T* inPtr[NIN];
303+ __local_mem__ T* outPtr[NOUT];
304+ for (uint32_t i = 0; i < NIN; i++) {
305+ inPtr[i] = (__local_mem__ T*)slot_[i].GetPhyAddr();
306+ }
307+ for (uint32_t o = 0; o < NOUT; o++) {
308+ outPtr[o] = (__local_mem__ T*)slot_[NIN + o].GetPhyAddr();
309+ }
310+ // 暂存区按 tileLen 跨距切分: tileLen 已对齐到向量寄存器长度, 保证 Load/StoreAlign 的
311+ // 地址对齐要求; 若按当前分片长度 len 切分, len=1 这类用例会得到 4 字节偏移而触发对齐错误。
312+ __local_mem__ float* sc[4] = {scratch_, scratch_ + tiling_->tileLen, scratch_ + 2U * tiling_->tileLen,
313+ scratch_ + 3U * tiling_->tileLen};
314+ Compute::template Run<T>(inPtr, outPtr, sc, len);
315+ 
316+ event_t v2m3 = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE3>());
317+ SetFlag<HardEvent::V_MTE3>(v2m3);
318+ WaitFlag<HardEvent::V_MTE3>(v2m3);
319+ GetTPipePtr()->ReleaseEventID<HardEvent::V_MTE3>(v2m3);
320+ 
321+ DataCopyExtParams cp{1, static_cast<uint32_t>(len * sizeof(T)), 0, 0, 0};
322+ for (uint32_t o = 0; o < NOUT; o++) {
323+ DataCopyPad(outGm_[o][dstOff], slot_[NIN + o], cp);
324+ }
325+ 
326+ // 搬出未完成前不得让下轮的 V/S 覆写输出槽(MTE3->MTE2 挡不住 V), 也不得让 MTE2 覆写输入槽。
327+ event_t m32v = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>());
328+ SetFlag<HardEvent::MTE3_V>(m32v);
329+ WaitFlag<HardEvent::MTE3_V>(m32v);
330+ GetTPipePtr()->ReleaseEventID<HardEvent::MTE3_V>(m32v);
331+ event_t m32m2 = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_MTE2>());
332+ SetFlag<HardEvent::MTE3_MTE2>(m32m2);
333+ WaitFlag<HardEvent::MTE3_MTE2>(m32m2);
334+ GetTPipePtr()->ReleaseEventID<HardEvent::MTE3_MTE2>(m32m2);
335+ }
336+ 
337+ const Tiling* tiling_ = nullptr;
338+ uint32_t blockIdx_ = 0;
339+ GlobalTensor<T> inGm_[NIN];
340+ GlobalTensor<T> outGm_[NOUT];
341+ TBuf<TPosition::VECCALC> ubBuf_;
342+ LocalTensor<T> slot_[SLOT_NUM];
343+ __local_mem__ float* scratch_ = nullptr;
344+};
345+} // namespace LambBrc
346+#endif // LAMB_BRC_KERNEL_H
Aoptim/lamb_apply_common/op_kernel/arch35/lamb_brc_tiling_data.h+65-0
@@ -0,0 +1,65 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_brc_tiling_data.h
13+ * \brief LAMB 族「多入多出 + 任意 numpy 广播」逐元素算子的共用 tiling 数据
14+ *
15+ * host 与 kernel 引用同一个模板实例, 布局天然一致(host 侧 memcpy 整个 POD 下发)。
16+ */
17+ 
18+#ifndef LAMB_BRC_TILING_DATA_H
19+#define LAMB_BRC_TILING_DATA_H
20+ 
21+#include <cstdint>
22+ 
23+// GE 原型最大轴数(与 silu_grad / reverse_v2 等同仓算子一致)。规划过程只做轴合并与块内切分,
24+// 不会增加轴数, 因此这个尺寸对任何合法输入都够用 —— 不存在"超了就拒收"的分支。
25+constexpr uint32_t LAMB_BRC_MAX_DIM = 8;
26+ 
27+// 硬件常量(非经验阈值): UB 搬运的最小块 32B、向量寄存器 256B。kernel 侧用平台接口
28+// GetUbBlockSize()/GetVRegSize() 静态校验二者一致, 避免两侧取值走偏。
29+constexpr uint32_t LAMB_BRC_GM_BLOCK_BYTES = 32;
30+constexpr uint32_t LAMB_BRC_VREG_BYTES = 256;
31+ 
32+// inKind 取值
33+constexpr uint32_t LAMB_BRC_KIND_SCALAR = 0; // 单元素, UB 槽位铺一次后常驻
34+constexpr uint32_t LAMB_BRC_KIND_SAME = 1; // 与输出同形, 连续搬入
35+constexpr uint32_t LAMB_BRC_KIND_BRC = 2; // 行块内需展开广播
36+ 
37+// tilingKey(数据内的分支, 与 binary 的 tilingKey 无关)
38+constexpr uint32_t LAMB_BRC_KEY_FLAT = 0; // 无需广播的输入 -> 纯线性分片, 不受行块大小约束
39+constexpr uint32_t LAMB_BRC_KEY_BLOCK = 1; // 有需广播的输入 -> 按输出行块处理
40+ 
41+template <uint32_t NIN, uint32_t NOUT>
42+struct LambBrcTilingData {
43+ uint64_t totalNum = 0; // 输出元素总数
44+ uint64_t totalRows = 0; // 分块模式下的行块总数
45+ uint64_t effStride[NIN * LAMB_BRC_MAX_DIM] = {}; // [i][j] 输入i在外层轴j的源步长(广播轴为0)
46+ uint64_t srcBlockLen[NIN] = {}; // 输入i在一个行块内的源元素数
47+ uint64_t perCoreElems = 0; // 平铺模式每核元素数(已对齐 32B)
48+ uint64_t rowsPerCore = 0; // 分块模式每核行块数(已对齐 32B)
49+ uint32_t tileLen = 0; // 由 ubSize 反解的分片元素数
50+ uint32_t blockLen = 0; // 一个输出行块的元素数
51+ uint32_t rowsPerTile = 0; // 一个分片内批量处理的行块数
52+ uint32_t innerChunk = 0; // 最内轴装不下时的块内切分长度(0=不切)
53+ uint32_t innerChunkCnt = 0; // 最内轴被切成几段
54+ uint32_t usedCoreNum = 0;
55+ uint32_t splitAxis = 0; // 前导行维个数
56+ uint32_t collapsedRank = 0;
57+ uint32_t tilingKey = 0;
58+ uint32_t inKind[NIN] = {};
59+ uint32_t outShape[LAMB_BRC_MAX_DIM] = {};
60+ uint32_t inShape[NIN * LAMB_BRC_MAX_DIM] = {};
61+ 
62+ static constexpr uint32_t SlotNum() { return NIN + NOUT + 1; } // 入 + 出 + 1 广播暂存
63+};
64+ 
65+#endif // LAMB_BRC_TILING_DATA_H
Aoptim/lamb_apply_common/op_kernel/arch35/lamb_next_mv_vf.h+187-0
@@ -0,0 +1,187 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_next_mv_vf.h
13+ * \brief LambNextMV 的 regbase 计算段(13 入 4 出), 搬运/广播/分核见 lamb_apply_common/lamb_brc_kernel.h
14+ *
15+ * 与 golden 同运算序(先乘后加两步, 不用 FMA 融合形式):
16+ * next_v = v*b2 + g2*(1-b2) -> y3
17+ * next_m = m*b1 + g *(1-b1) -> y2
18+ * v_unb = next_v/rd1 ; m_unb = next_m/rd0
19+ * y1 = param*wd + m_unb / sqrt(v_unb + eps)
20+ * y4 = m_unb / (sqrt(v_unb) + eps)
21+ * 中间一律 fp32(A2 的 TBE compute dtype='float32'), 出口窄回 T。
22+ */
23+ 
24+#ifndef LAMB_NEXT_MV_VF_H
25+#define LAMB_NEXT_MV_VF_H
26+ 
27+#include "lamb_brc_kernel.h"
28+ 
29+namespace LambNextMVOp {
30+using namespace AscendC;
31+ 
32+// 输入下标(与原型入参顺序一致)
33+enum : uint32_t {
34+ IN_G2 = 0, // input_mul3 = g^2
35+ IN_V = 1, // input_mul2 = v
36+ IN_RD1 = 2, // input_realdiv1 = 1 - b2^t
37+ IN_G = 3, // input_mul1 = g
38+ IN_M = 4, // input_mul0 = m
39+ IN_RD0 = 5, // input_realdiv0 = 1 - b1^t
40+ IN_PARAM = 6, // input_mul4 = param
41+ IN_B1 = 7,
42+ IN_OMB1 = 8,
43+ IN_B2 = 9,
44+ IN_OMB2 = 10,
45+ IN_WD = 11,
46+ IN_EPS = 12
47+};
48+ 
49+template <bool WITH_DECAY>
50+struct LambNextMVVf {
51+ template <typename T>
52+ static __aicore__ inline void Run(__local_mem__ T** in, __local_mem__ T** out, __local_mem__ float** scratch,
53+ uint32_t count)
54+ {
55+ uint16_t vfTimes = static_cast<uint16_t>(count / LambBrc::LAMB_VF_LEN);
56+ uint32_t tail = count % LambBrc::LAMB_VF_LEN;
57+ uint16_t tailTimes = (tail > 0) ? 1 : 0;
58+ 
59+ __local_mem__ T* pG2 = in[IN_G2];
60+ __local_mem__ T* pV = in[IN_V];
61+ __local_mem__ T* pRd1 = in[IN_RD1];
62+ __local_mem__ T* pG = in[IN_G];
63+ __local_mem__ T* pM = in[IN_M];
64+ __local_mem__ T* pRd0 = in[IN_RD0];
65+ __local_mem__ T* pParam = in[IN_PARAM];
66+ __local_mem__ T* pB1 = in[IN_B1];
67+ __local_mem__ T* pOmB1 = in[IN_OMB1];
68+ __local_mem__ T* pB2 = in[IN_B2];
69+ __local_mem__ T* pOmB2 = in[IN_OMB2];
70+ __local_mem__ T* pWd = in[IN_WD];
71+ __local_mem__ T* pEps = in[IN_EPS];
72+ __local_mem__ T* y1 = out[0];
73+ __local_mem__ T* y2 = out[1];
74+ __local_mem__ T* y3 = out[2];
75+ __local_mem__ T* y4 = out[3];
76+ // 中间量以 fp32 暂存, 不经 T 的舍入 —— 与单段实现数值完全一致。
77+ __local_mem__ float* sNextV = scratch[0];
78+ __local_mem__ float* sNextM = scratch[1];
79+ 
80+ __local_mem__ float* sVUnb = scratch[2];
81+ __local_mem__ float* sMUnb = scratch[3];
82+ 
83+ // 分 4 段: 全量精确档(0ULP/FTZ_FALSE)展开后指令很多, 单段的循环回边偏移会超出
84+ // scbzi 立即数范围 [-512,511](--cce-long-scbz=true 压不住)。每段最多 2 个精确档调用,
85+ // 中间量一律以 fp32 暂存传递, 不经 T 的舍入, 与单段实现数值完全一致。
86+ 
87+ // ---- 段1: next_v(y3) / next_m(y2) ----
88+ __VEC_SCOPE__
89+ {
90+ Reg::RegTensor<float> vA, vB, vNextV, vNextM;
91+ Reg::MaskReg maskAll = Reg::CreateMask<float, Reg::MaskPattern::ALL>();
92+ Reg::MaskReg maskT = Reg::UpdateMask<float>(tail);
93+ for (uint16_t vfIdx = 0; vfIdx < vfTimes + tailTimes; vfIdx++) {
94+ uint32_t off = vfIdx * LambBrc::LAMB_VF_LEN;
95+ Reg::MaskReg preg = (vfIdx < vfTimes) ? maskAll : maskT;
96+ LambBrc::Load<T>(pV, vA, preg, off);
97+ LambBrc::Load<T>(pB2, vB, preg, off);
98+ Reg::Mul(vNextV, vA, vB, preg);
99+ LambBrc::Load<T>(pG2, vA, preg, off);
100+ LambBrc::Load<T>(pOmB2, vB, preg, off);
101+ Reg::Mul(vA, vA, vB, preg);
102+ Reg::Add(vNextV, vNextV, vA, preg);
103+ LambBrc::Store<T>(y3, vNextV, preg, off);
104+ Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM>(sNextV + off, vNextV, preg);
105+ LambBrc::Load<T>(pM, vA, preg, off);
106+ LambBrc::Load<T>(pB1, vB, preg, off);
107+ Reg::Mul(vNextM, vA, vB, preg);
108+ LambBrc::Load<T>(pG, vA, preg, off);
109+ LambBrc::Load<T>(pOmB1, vB, preg, off);
110+ Reg::Mul(vA, vA, vB, preg);
111+ Reg::Add(vNextM, vNextM, vA, preg);
112+ LambBrc::Store<T>(y2, vNextM, preg, off);
113+ Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM>(sNextM + off, vNextM, preg);
114+ }
115+ }
116+ 
117+ // ---- 段2: 去偏 (2 次精确除法) ----
118+ __VEC_SCOPE__
119+ {
120+ Reg::RegTensor<float> vA, vT, vU;
121+ Reg::MaskReg maskAll = Reg::CreateMask<float, Reg::MaskPattern::ALL>();
122+ Reg::MaskReg maskT = Reg::UpdateMask<float>(tail);
123+ for (uint16_t vfIdx = 0; vfIdx < vfTimes + tailTimes; vfIdx++) {
124+ uint32_t off = vfIdx * LambBrc::LAMB_VF_LEN;
125+ Reg::MaskReg preg = (vfIdx < vfTimes) ? maskAll : maskT;
126+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vT, sNextV + off);
127+ LambBrc::Load<T>(pRd1, vA, preg, off);
128+ Reg::Div<float, &LambBrc::LAMB_PRECISE_DIV>(vU, vT, vA, preg);
129+ Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM>(sVUnb + off, vU, preg);
130+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vT, sNextM + off);
131+ LambBrc::Load<T>(pRd0, vA, preg, off);
132+ Reg::Div<float, &LambBrc::LAMB_PRECISE_DIV>(vU, vT, vA, preg);
133+ Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM>(sMUnb + off, vU, preg);
134+ }
135+ }
136+ 
137+ // ---- 段3: y1 = param*wd + m_unb / sqrt(v_unb + eps) ----
138+ __VEC_SCOPE__
139+ {
140+ Reg::RegTensor<float> vA, vB, vPw, vVU, vMU;
141+ Reg::MaskReg maskAll = Reg::CreateMask<float, Reg::MaskPattern::ALL>();
142+ Reg::MaskReg maskT = Reg::UpdateMask<float>(tail);
143+ for (uint16_t vfIdx = 0; vfIdx < vfTimes + tailTimes; vfIdx++) {
144+ uint32_t off = vfIdx * LambBrc::LAMB_VF_LEN;
145+ Reg::MaskReg preg = (vfIdx < vfTimes) ? maskAll : maskT;
146+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vVU, sVUnb + off);
147+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vMU, sMUnb + off);
148+ LambBrc::Load<T>(pParam, vA, preg, off);
149+ LambBrc::Load<T>(pWd, vB, preg, off);
150+ Reg::Mul(vPw, vA, vB, preg);
151+ LambBrc::Load<T>(pEps, vB, preg, off);
152+ Reg::Add(vA, vVU, vB, preg);
153+ Reg::Sqrt<float, &LambBrc::LAMB_PRECISE_SQRT>(vA, vA, preg);
154+ Reg::Div<float, &LambBrc::LAMB_PRECISE_DIV>(vA, vMU, vA, preg);
155+ Reg::Add(vA, vPw, vA, preg);
156+ LambBrc::Store<T>(y1, vA, preg, off);
157+ }
158+ }
159+ 
160+ // ---- 段4: y4 = [param*wd +] m_unb / (sqrt(v_unb) + eps) ----
161+ __VEC_SCOPE__
162+ {
163+ Reg::RegTensor<float> vA, vB, vPw, vVU, vMU;
164+ Reg::MaskReg maskAll = Reg::CreateMask<float, Reg::MaskPattern::ALL>();
165+ Reg::MaskReg maskT = Reg::UpdateMask<float>(tail);
166+ for (uint16_t vfIdx = 0; vfIdx < vfTimes + tailTimes; vfIdx++) {
167+ uint32_t off = vfIdx * LambBrc::LAMB_VF_LEN;
168+ Reg::MaskReg preg = (vfIdx < vfTimes) ? maskAll : maskT;
169+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vVU, sVUnb + off);
170+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vMU, sMUnb + off);
171+ Reg::Sqrt<float, &LambBrc::LAMB_PRECISE_SQRT>(vA, vVU, preg);
172+ LambBrc::Load<T>(pEps, vB, preg, off);
173+ Reg::Add(vA, vA, vB, preg);
174+ Reg::Div<float, &LambBrc::LAMB_PRECISE_DIV>(vA, vMU, vA, preg);
175+ if constexpr (WITH_DECAY) {
176+ LambBrc::Load<T>(pParam, vPw, preg, off);
177+ LambBrc::Load<T>(pWd, vB, preg, off);
178+ Reg::Mul(vPw, vPw, vB, preg);
179+ Reg::Add(vA, vPw, vA, preg);
180+ }
181+ LambBrc::Store<T>(y4, vA, preg, off);
182+ }
183+ }
184+ }
185+};
186+} // namespace LambNextMVOp
187+#endif // LAMB_NEXT_MV_VF_H
Moptim/lamb_apply_optimizer_assign/CMakeLists.txt+3-3
@@ -1,9 +1,9 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4# CANN Open Software License Agreement Version 2.0 (the "License").4# CANN Open Software License Agreement Version 2.0 (the "License").
5# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
6-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
@@ -13,4 +13,4 @@ set(SUPPORT_COMPUTE_UNIT "ascend950")
13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14set(SUPPORT_TILING_DIR "arch35")14set(SUPPORT_TILING_DIR "arch35")
15 15 
16-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_apply_optimizer_assign ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)16+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_apply_optimizer_assign ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DEPENDENCIES lamb_apply_common DISABLE_IN_OPP TRUE)
Moptim/lamb_apply_optimizer_assign/README.md+11-11
@@ -67,63 +67,63 @@
67 <tr>67 <tr>
68 <td>input3</td>68 <td>input3</td>
69 <td>输入</td>69 <td>输入</td>
70- <td>支持空Tensor。公式中的input3（参与权重衰减的参数）。<b>唯一参与广播的输入</b>：按右对齐broadcast规则向inputv对齐，维度数可少于inputv（不可多于），对应维需相等或为1。</td>70+ <td>支持空Tensor。公式中的input3（参与权重衰减的参数），shape需与其他输入满足broadcast关系。</td>
71 <td>FLOAT16、FLOAT</td>71 <td>FLOAT16、FLOAT</td>
72 <td>ND</td>72 <td>ND</td>
73 </tr>73 </tr>
74 <tr>74 <tr>
75 <td>mul0_x</td>75 <td>mul0_x</td>
76 <td>输入</td>76 <td>输入</td>
77- <td>不支持空Tensor。公式中的mul0_x（beta1），标量。</td>77+ <td>支持空Tensor。公式中的mul0_x（beta1），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
78 <td>FLOAT16、FLOAT</td>78 <td>FLOAT16、FLOAT</td>
79 <td>ND</td>79 <td>ND</td>
80 </tr>80 </tr>
81 <tr>81 <tr>
82 <td>mul1_x</td>82 <td>mul1_x</td>
83 <td>输入</td>83 <td>输入</td>
84- <td>不支持空Tensor。公式中的mul1_x（1-beta1），标量。</td>84+ <td>支持空Tensor。公式中的mul1_x（1-beta1），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
85 <td>FLOAT16、FLOAT</td>85 <td>FLOAT16、FLOAT</td>
86 <td>ND</td>86 <td>ND</td>
87 </tr>87 </tr>
88 <tr>88 <tr>
89 <td>mul2_x</td>89 <td>mul2_x</td>
90 <td>输入</td>90 <td>输入</td>
91- <td>不支持空Tensor。公式中的mul2_x（beta2），标量。</td>91+ <td>支持空Tensor。公式中的mul2_x（beta2），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
92 <td>FLOAT16、FLOAT</td>92 <td>FLOAT16、FLOAT</td>
93 <td>ND</td>93 <td>ND</td>
94 </tr>94 </tr>
95 <tr>95 <tr>
96 <td>mul3_x</td>96 <td>mul3_x</td>
97 <td>输入</td>97 <td>输入</td>
98- <td>不支持空Tensor。公式中的mul3_x（1-beta2），标量。</td>98+ <td>支持空Tensor。公式中的mul3_x（1-beta2），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
99 <td>FLOAT16、FLOAT</td>99 <td>FLOAT16、FLOAT</td>
100 <td>ND</td>100 <td>ND</td>
101 </tr>101 </tr>
102 <tr>102 <tr>
103 <td>add2_y</td>103 <td>add2_y</td>
104 <td>输入</td>104 <td>输入</td>
105- <td>不支持空Tensor。公式中的add2_y（epsilon），标量。</td>105+ <td>支持空Tensor。公式中的add2_y（epsilon），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
106 <td>FLOAT16、FLOAT</td>106 <td>FLOAT16、FLOAT</td>
107 <td>ND</td>107 <td>ND</td>
108 </tr>108 </tr>
109 <tr>109 <tr>
110 <td>steps</td>110 <td>steps</td>
111 <td>输入</td>111 <td>输入</td>
112- <td>不支持空Tensor。公式中的steps（步数），标量。</td>112+ <td>支持空Tensor。公式中的steps（步数），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
113 <td>FLOAT16、FLOAT</td>113 <td>FLOAT16、FLOAT</td>
114 <td>ND</td>114 <td>ND</td>
115 </tr>115 </tr>
116 <tr>116 <tr>
117 <td>do_use_weight</td>117 <td>do_use_weight</td>
118 <td>输入</td>118 <td>输入</td>
119- <td>不支持空Tensor。公式中的do_use_weight（是否使用权重衰减），标量。</td>119+ <td>支持空Tensor。公式中的do_use_weight（是否使用权重衰减），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
120 <td>FLOAT16、FLOAT</td>120 <td>FLOAT16、FLOAT</td>
121 <td>ND</td>121 <td>ND</td>
122 </tr>122 </tr>
123 <tr>123 <tr>
124 <td>weight_decay_rate</td>124 <td>weight_decay_rate</td>
125 <td>输入</td>125 <td>输入</td>
126- <td>不支持空Tensor。公式中的weight_decay_rate（权重衰减率），标量。</td>126+ <td>支持空Tensor。公式中的weight_decay_rate（权重衰减率），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
127 <td>FLOAT16、FLOAT</td>127 <td>FLOAT16、FLOAT</td>
128 <td>ND</td>128 <td>ND</td>
129 </tr>129 </tr>
@@ -152,9 +152,9 @@
152 152 
153## 约束说明153## 约束说明
154 154 
155-- shape约束：`grad`、`inputv`、`inputm`三者shape必须**完全相同**，并直接决定三个输出的shape。`inputv`与`inputm`是原地更新的动量输出，不会被广播放大；`grad`同样不参与广播。**广播只发生在`input3`上**：`input3`按右对齐broadcast规则向`inputv`对齐，维度数可少于`inputv`（含标量），对应维需与`inputv`相等或为1；`input3`的维度数大于`inputv`、或某一维大于`inputv`对应维，均不被支持。155+- shape约束：所有输入的shape需两两满足broadcast规则（右对齐，对应维相等或为1）。`inputv`与`inputm`是原地(in-place)更新的动量输出，两者shape必须**完全相同**，且全部输入的broadcast结果必须恰好等于该shape（否则原地写回会越界）；三个输出的shape均取该shape。其余输入（含`grad`）可向其广播。
156+- 所有输入及输出的维度数不超过8。
156 157 
157-- `mul0_x`、`mul1_x`、`mul2_x`、`mul3_x`、`add2_y`、`steps`、`do_use_weight`、`weight_decay_rate`这8个标量输入不支持空Tensor；`grad`、`inputv`、`inputm`、`input3`支持空Tensor。
158 158 
159- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。159- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。
160 160 
Moptim/lamb_apply_optimizer_assign/op_host/arch35/lamb_apply_optimizer_assign_tiling_arch35.cpp+53-77
@@ -14,26 +14,29 @@
14 */14 */
15 15 
16#include "lamb_apply_optimizer_assign_tiling_arch35.h"16#include "lamb_apply_optimizer_assign_tiling_arch35.h"
17-#include "../../../lamb_apply_common/lamb_apply_check_util.h"17+#include "../../../lamb_apply_common/op_host/arch35/lamb_apply_check_util.h"
18#include <graph/utils/type_utils.h>18#include <graph/utils/type_utils.h>
19+#include <securec.h>
20+#include <algorithm>
19#include <string>21#include <string>
20#include "infershape_broadcast_util.h"22#include "infershape_broadcast_util.h"
21-#include "../../op_kernel/arch35/lamb_apply_optimizer_assign_dag.h"
22-#include "atvoss/broadcast/broadcast_tiling.h"
23#include "log/log.h"23#include "log/log.h"
24#include "platform/platform_info.h"24#include "platform/platform_info.h"
25#include "register/op_impl_registry.h"25#include "register/op_impl_registry.h"
26#include "register/tilingdata_base.h"26#include "register/tilingdata_base.h"
27#include "op_host/tiling_templates_registry.h"27#include "op_host/tiling_templates_registry.h"
28 28 
29-using namespace AscendC;
30using namespace ge;29using namespace ge;
31 30 
32namespace optiling {31namespace optiling {
33 32 
34constexpr static uint64_t LAMB_APPLY_OPTIMIZER_ASSIGN_TILING_PRIORITY = 0;33constexpr static uint64_t LAMB_APPLY_OPTIMIZER_ASSIGN_TILING_PRIORITY = 0;
34+constexpr static uint64_t TILING_KEY_FP32 = 100;
35+constexpr static uint64_t TILING_KEY_FP16 = 200;
35constexpr static int32_t INPUT_NUM = 12;36constexpr static int32_t INPUT_NUM = 12;
36constexpr static int32_t OUTPUT_NUM = 3;37constexpr static int32_t OUTPUT_NUM = 3;
38+constexpr static int32_t INPUTV_IDX = 1; // inputv: ref(原地)输出
39+constexpr static int32_t INPUTM_IDX = 2; // inputm: ref(原地)输出
37static const char* const kInputNames[] = {"grad", "inputv", "inputm", "input3", "mul0_x", "mul1_x",40static const char* const kInputNames[] = {"grad", "inputv", "inputm", "input3", "mul0_x", "mul1_x",
38 "mul2_x", "mul3_x", "add2_y", "steps", "do_use_weight", "weight_decay_rate"};41 "mul2_x", "mul3_x", "add2_y", "steps", "do_use_weight", "weight_decay_rate"};
39static const char* const kOutputNames[] = {"output0", "inputv", "inputm"};42static const char* const kOutputNames[] = {"output0", "inputv", "inputm"};
@@ -52,112 +55,75 @@ static ge::graphStatus TilingPrepareForLambApplyOptimizerAssign(gert::TilingPars
52 55 
53ge::graphStatus LambApplyOptimizerAssignTiling::GetShapeAttrsInfo()56ge::graphStatus LambApplyOptimizerAssignTiling::GetShapeAttrsInfo()
54{57{
55- static const int32_t kScalarInputIdx[] = {4, 5, 6, 7, 8, 9, 10, 11};
56 if (CheckLambApplyDtypeConsistency(context_, INPUT_NUM, kInputNames, OUTPUT_NUM, kOutputNames) !=58 if (CheckLambApplyDtypeConsistency(context_, INPUT_NUM, kInputNames, OUTPUT_NUM, kOutputNames) !=
57 ge::GRAPH_SUCCESS) {59 ge::GRAPH_SUCCESS) {
58 return ge::GRAPH_FAILED;60 return ge::GRAPH_FAILED;
59 }61 }
60- if (CheckLambApplyScalarNotEmpty(context_, kScalarInputIdx, sizeof(kScalarInputIdx) / sizeof(kScalarInputIdx[0]),
61- kInputNames) != ge::GRAPH_SUCCESS) {
62- return ge::GRAPH_FAILED;
63- }
64 return CheckInplaceShapeConstraint();62 return CheckInplaceShapeConstraint();
65}63}
66 64 
67// inputv、inputm 是 in-place 更新的动量输出(next_v/next_m 原地写回它们的输入 buffer,见 proto "(in-place)"),65// inputv、inputm 是 in-place 更新的动量输出(next_v/next_m 原地写回它们的输入 buffer,见 proto "(in-place)"),
68-// 输出形状由它们决定,故 inputv、inputm 必须同形状。66+// 两者形状必须相同, 且必须 == 全部输入广播的完整网格。其余输入(含绑在 In0 的 grad)可广播进这个网格。
69-// grad 绑定在广播 DAG 的 In0 位,底层 Ops::Base 广播模板(DoDimensionCollapse)不支持对 In0 做广播:
70-// grad 为标量时 EnsureNotScalar 只抬到 {1}、不会左补 1 对齐输出 rank,直接撞 "dim num is not same";
71-// grad 与输出同 rank 但某维为 1 时同样被拒("dim index is not same with out")。故 grad 必须与 inputv 等形。
72-// 仅 input3 参与广播(右对齐,维度数可少于 inputv,含标量),这是实测支持的形态。
73-// 若后续 ops-base 放开 In0 广播,此处与 infershape 需同步放宽。
74ge::graphStatus LambApplyOptimizerAssignTiling::CheckInplaceShapeConstraint()67ge::graphStatus LambApplyOptimizerAssignTiling::CheckInplaceShapeConstraint()
75{68{
76- auto gradShape = context_->GetInputShape(0);69+ auto inputvShape = context_->GetInputShape(INPUTV_IDX);
77- auto inputvShape = context_->GetInputShape(1);70+ auto inputmShape = context_->GetInputShape(INPUTM_IDX);
78- auto inputmShape = context_->GetInputShape(2);
79- auto input3Shape = context_->GetInputShape(3);
80- OP_CHECK_NULL_WITH_CONTEXT(context_, gradShape);
81 OP_CHECK_NULL_WITH_CONTEXT(context_, inputvShape);71 OP_CHECK_NULL_WITH_CONTEXT(context_, inputvShape);
82 OP_CHECK_NULL_WITH_CONTEXT(context_, inputmShape);72 OP_CHECK_NULL_WITH_CONTEXT(context_, inputmShape);
83- OP_CHECK_NULL_WITH_CONTEXT(context_, input3Shape);
84- const auto& gs = gradShape->GetStorageShape();
85 const auto& vs = inputvShape->GetStorageShape();73 const auto& vs = inputvShape->GetStorageShape();
86 const auto& ms = inputmShape->GetStorageShape();74 const auto& ms = inputmShape->GetStorageShape();
87- const auto& ps = input3Shape->GetStorageShape();
88- gert::Shape bcShape;
89 if (!(vs == ms)) {75 if (!(vs == ms)) {
90 OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(76 OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(
91 context_->GetNodeName(), "inputv and inputm",77 context_->GetNodeName(), "inputv and inputm",
92 (Ops::Base::ToString(vs) + " and " + Ops::Base::ToString(ms)).c_str(),78 (Ops::Base::ToString(vs) + " and " + Ops::Base::ToString(ms)).c_str(),
93- "inputv and inputm are in-place updated moments and must have the same shape equal to the broadcast "79+ "inputv and inputm are in-place updated moments and must have the same shape");
94- "output shape");
95 return ge::GRAPH_FAILED;80 return ge::GRAPH_FAILED;
96 }81 }
97- // grad 不参与广播,必须与 inputv/inputm 等形82+ return CheckLambApplyBroadcastIntoRef(context_, INPUT_NUM, INPUTV_IDX, "inputv");
98- if (!(gs == vs)) {
99- OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
100- context_->GetNodeName(), "grad", Ops::Base::ToString(gs).c_str(),
101- "grad does not support broadcast and must have exactly the same shape as inputv/inputm");
102- return ge::GRAPH_FAILED;
103- }
104- // input3 能广播进 inputv <=> broadcast(input3, inputv) == inputv
105- if (!Ops::Base::BroadcastShape(&ps, &vs, &bcShape) || !(bcShape == vs)) {
106- OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
107- context_->GetNodeName(), "input3", Ops::Base::ToString(ps).c_str(),
108- "input3 must be broadcastable into the in-place moment shape inputv/inputm");
109- return ge::GRAPH_FAILED;
110- }
111- return ge::GRAPH_SUCCESS;
112}83}
113 84 
114bool LambApplyOptimizerAssignTiling::IsCapable() { return true; }85bool LambApplyOptimizerAssignTiling::IsCapable() { return true; }
115 86 
116ge::graphStatus LambApplyOptimizerAssignTiling::DoOpTiling()87ge::graphStatus LambApplyOptimizerAssignTiling::DoOpTiling()
117{88{
118- // 空 tensor 应对(空进空出): 输出为空(0元素)时设 1 核(空转), 配合全0 tiling 数据(blockFormer=0)使 kernel 空转退出,89+ auto rawTilingData = context_->GetRawTilingData();
119- // 直接成功。90+ OP_CHECK_NULL_WITH_CONTEXT(context_, rawTilingData);
120- auto emptyTensorOutShape0 = context_->GetOutputShape(0);
121- if (emptyTensorOutShape0 != nullptr && emptyTensorOutShape0->GetStorageShape().GetShapeSize() == 0) {
122- auto emptyRawTiling = context_->GetRawTilingData();
123- if (emptyRawTiling != nullptr && emptyRawTiling->GetData() != nullptr) {
124- size_t emptyCap = emptyRawTiling->GetCapacity();
125- uint8_t* emptyPtr = static_cast<uint8_t*>(emptyRawTiling->GetData());
126- for (size_t emptyIdx = 0; emptyIdx < emptyCap; ++emptyIdx) {
127- emptyPtr[emptyIdx] = 0;
128- }
129- emptyRawTiling->SetDataSize(emptyCap);
130- }
131- size_t* emptyWs = context_->GetWorkspaceSizes(1);
132- if (emptyWs != nullptr) {
133- emptyWs[0] = 0;
134- }
135- context_->SetBlockDim(1);
136- tilingKey = GET_TPL_TILING_KEY(1); // schMode=1(已编译), 配合全0 tiling(blockFormer=0)空转
137- return ge::GRAPH_SUCCESS;
138- }
139 auto input0Desc = context_->GetInputDesc(0);91 auto input0Desc = context_->GetInputDesc(0);
140 OP_CHECK_NULL_WITH_CONTEXT(context_, input0Desc);92 OP_CHECK_NULL_WITH_CONTEXT(context_, input0Desc);
93+ 
141 ge::DataType input0DType = input0Desc->GetDataType();94 ge::DataType input0DType = input0Desc->GetDataType();
142- if (input0DType == ge::DT_FLOAT16) {95+ uint32_t dtSize = 0;
143- BroadcastBaseTiling<LambApplyOptimizerAssignOp::LambApplyOptimizerAssignCompute<half, float>::OpDag>96+ if (input0DType == ge::DT_FLOAT) {
144- brcBaseTiling(context_, static_cast<uint32_t>(BROADCAST_KERNEL_TYPE::KERNEL_TYPE_NDDMA));97+ tilingKey = TILING_KEY_FP32;
145- OP_CHECK_IF(brcBaseTiling.DoTiling() == ge::GRAPH_FAILED,98+ dtSize = sizeof(float);
146- OP_LOGE(context_->GetNodeName(), "Do tiling failed. Please check the detailed log."),99+ } else if (input0DType == ge::DT_FLOAT16) {
147- return ge::GRAPH_FAILED);100+ tilingKey = TILING_KEY_FP16;
148- tilingKey = GET_TPL_TILING_KEY(brcBaseTiling.GetSchMode());101+ dtSize = sizeof(uint16_t);
149- } else if (input0DType == ge::DT_FLOAT) {
150- BroadcastBaseTiling<LambApplyOptimizerAssignOp::LambApplyOptimizerAssignCompute<float, float>::OpDag>
151- brcBaseTiling(context_, static_cast<uint32_t>(BROADCAST_KERNEL_TYPE::KERNEL_TYPE_NDDMA));
152- OP_CHECK_IF(brcBaseTiling.DoTiling() == ge::GRAPH_FAILED,
153- OP_LOGE(context_->GetNodeName(), "Do tiling failed. Please check the detailed log."),
154- return ge::GRAPH_FAILED);
155- tilingKey = GET_TPL_TILING_KEY(brcBaseTiling.GetSchMode());
156 } else {102 } else {
157 OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "grad", Ops::Base::ToString(input0DType).c_str(),103 OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "grad", Ops::Base::ToString(input0DType).c_str(),
158 "fp16 or fp32");104 "fp16 or fp32");
159 return ge::GRAPH_FAILED;105 return ge::GRAPH_FAILED;
160 }106 }
107+ 
108+ using PlanTiling = LambBrcTilingData<12, 3>;
109+ td_ = PlanTiling{};
110+ // 空进空出: 输出 0 元素时 tiling 全 0, kernel 按 usedCoreNum=0 直接退出。
111+ auto outShape0 = context_->GetOutputShape(0);
112+ if (outShape0 == nullptr || outShape0->GetStorageShape().GetShapeSize() != 0) {
113+ // 先取返回值再判: 模板实参里的逗号会被预处理器当成 OP_CHECK_IF 的参数分隔符。
114+ ge::graphStatus planRet = BuildLambBrcPlan<12, 3>(context_, coreNum_, ubSize_, dtSize, td_);
115+ OP_CHECK_IF(planRet != ge::GRAPH_SUCCESS, OP_LOGE(context_->GetNodeName(), "build broadcast plan failed"),
116+ return ge::GRAPH_FAILED);
117+ }
118+ 
119+ auto ret = memcpy_s(rawTilingData->GetData(), rawTilingData->GetCapacity(), &td_, sizeof(td_));
120+ OP_CHECK_IF(ret != EOK, OP_LOGE(context_->GetNodeName(), "copy tiling data failed, ret %d", ret),
121+ return ge::GRAPH_FAILED);
122+ rawTilingData->SetDataSize(sizeof(td_));
123+ context_->SetBlockDim(std::max<uint32_t>(td_.usedCoreNum, 1));
124+ size_t* ws = context_->GetWorkspaceSizes(1);
125+ OP_CHECK_NULL_WITH_CONTEXT(context_, ws);
126+ ws[0] = 0U;
161 return ge::GRAPH_SUCCESS;127 return ge::GRAPH_SUCCESS;
162}128}
163 129 
@@ -169,7 +135,17 @@ ge::graphStatus LambApplyOptimizerAssignTiling::GetWorkspaceSize() { return ge::
169 135 
170ge::graphStatus LambApplyOptimizerAssignTiling::PostTiling() { return ge::GRAPH_SUCCESS; }136ge::graphStatus LambApplyOptimizerAssignTiling::PostTiling() { return ge::GRAPH_SUCCESS; }
171 137 
172-ge::graphStatus LambApplyOptimizerAssignTiling::GetPlatformInfo() { return ge::GRAPH_SUCCESS; }138+ge::graphStatus LambApplyOptimizerAssignTiling::GetPlatformInfo()
139+{
140+ auto compileInfo = static_cast<const LambApplyOptimizerAssignCompileInfo*>(context_->GetCompileInfo());
141+ OP_CHECK_NULL_WITH_CONTEXT(context_, compileInfo);
142+ coreNum_ = compileInfo->coreNum;
143+ ubSize_ = compileInfo->ubSize;
144+ OP_CHECK_IF(coreNum_ == 0 || ubSize_ == 0,
145+ OP_LOGE(context_->GetNodeName(), "invalid platform info: coreNum %lu ubSize %lu", coreNum_, ubSize_),
146+ return ge::GRAPH_FAILED);
147+ return ge::GRAPH_SUCCESS;
148+}
173 149 
174static ge::graphStatus TilingForLambApplyOptimizerAssign(gert::TilingContext* context)150static ge::graphStatus TilingForLambApplyOptimizerAssign(gert::TilingContext* context)
175{151{
Moptim/lamb_apply_optimizer_assign/op_host/arch35/lamb_apply_optimizer_assign_tiling_arch35.h+6-2
@@ -16,7 +16,8 @@
16#ifndef OPS_OPTIM_ADAM_APPLY_ONE_OP_HOST_ADAM_APPLY_ONE_TILING_ARCH35_H16#ifndef OPS_OPTIM_ADAM_APPLY_ONE_OP_HOST_ADAM_APPLY_ONE_TILING_ARCH35_H
17#define OPS_OPTIM_ADAM_APPLY_ONE_OP_HOST_ADAM_APPLY_ONE_TILING_ARCH35_H17#define OPS_OPTIM_ADAM_APPLY_ONE_OP_HOST_ADAM_APPLY_ONE_TILING_ARCH35_H
18 18 
19-#include "../../op_kernel/arch35/lamb_apply_optimizer_assign_tiling_key.h"19+#include "../lamb_apply_optimizer_assign_tiling_def.h"
20+#include "../../../lamb_apply_common/op_host/arch35/lamb_brc_tiling_plan.h"
20#include "op_host/tiling_base.h"21#include "op_host/tiling_base.h"
21 22 
22using namespace Ops::NN::Optiling;23using namespace Ops::NN::Optiling;
@@ -44,9 +45,12 @@ protected:
44 45 
45private:46private:
46 // 校验 in-place 更新的动量输入(inputv/inputm)形状 == 全广播输出网格(grad/input3 可向上广播)。47 // 校验 in-place 更新的动量输入(inputv/inputm)形状 == 全广播输出网格(grad/input3 可向上广播)。
47- // dtype 一致性、标量非空的通用校验见 lamb_apply_common/lamb_apply_check_util.h。48+ // dtype 一致性、标量非空的通用校验见 lamb_apply_common/op_host/arch35/lamb_apply_check_util.h。
48 ge::graphStatus CheckInplaceShapeConstraint();49 ge::graphStatus CheckInplaceShapeConstraint();
49 uint64_t tilingKey = 0;50 uint64_t tilingKey = 0;
51+ uint64_t coreNum_ = 0;
52+ uint64_t ubSize_ = 0;
53+ LambBrcTilingData<12, 3> td_;
50};54};
51 55 
52} // namespace optiling56} // namespace optiling
Moptim/lamb_apply_optimizer_assign/op_host/lamb_apply_optimizer_assign_infershape.cpp+57-50
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include <vector>
16#include "register/op_impl_registry.h"17#include "register/op_impl_registry.h"
17#include "log/log.h"18#include "log/log.h"
18#include "infershape_broadcast_util.h"19#include "infershape_broadcast_util.h"
@@ -20,65 +21,71 @@
20using namespace Ops::Base;21using namespace Ops::Base;
21using namespace ge;22using namespace ge;
22namespace ops {23namespace ops {
23-constexpr size_t GRAD_IDX = 0;24+// A2 语义: 本族算子的所有输入都是可广播的 ND Tensor(见 canndev
25+// ops/built-in/tbe/impl/lamb_*.py, 每一步 mul/sub/div 都先 shape_util.broadcast_shapes
26+// 再 tbe.broadcast), 输出形状为全部输入广播的结果。A2 的 op_proto 只声明了其中两个输入,
27+// 属声明宽松, 不作为支持面依据。
28+// inputv/inputm 同时是 ref 输出(原地更新的动量), 故广播结果必须恰好等于它们的形状,
29+// 否则原地写回会越过其显存边界。其余输入(含 In0 的 grad)可广播进这个形状。
30+constexpr size_t IN_NUM = 12;
24constexpr size_t INPUTV_IDX = 1;31constexpr size_t INPUTV_IDX = 1;
25constexpr size_t INPUTM_IDX = 2;32constexpr size_t INPUTM_IDX = 2;
26-constexpr size_t INPUT3_IDX = 3;33+constexpr size_t OUT_NUM = 3;
27-constexpr size_t OUTPUT0_IDX = 0;
28-constexpr size_t OUTPUTV_IDX = 1;
29-constexpr size_t OUTPUTM_IDX = 2;
30 34 
31static ge::graphStatus InferShape4LambApplyOptimizerAssign(gert::InferShapeContext* context)35static ge::graphStatus InferShape4LambApplyOptimizerAssign(gert::InferShapeContext* context)
32{36{
33- auto grad_shape = context->GetInputShape(GRAD_IDX);37+ std::vector<const gert::Shape*> inShapes;
34- OP_CHECK_NULL_WITH_CONTEXT(context, grad_shape);38+ inShapes.reserve(IN_NUM);
35- auto inputv_shape = context->GetInputShape(INPUTV_IDX);39+ for (size_t i = 0; i < IN_NUM; i++) {
36- OP_CHECK_NULL_WITH_CONTEXT(context, inputv_shape);40+ auto in = context->GetInputShape(i);
37- auto inputm_shape = context->GetInputShape(INPUTM_IDX);41+ OP_CHECK_NULL_WITH_CONTEXT(context, in);
38- OP_CHECK_NULL_WITH_CONTEXT(context, inputm_shape);42+ inShapes.push_back(in);
39- auto input3_shape = context->GetInputShape(INPUT3_IDX);43+ }
40- OP_CHECK_NULL_WITH_CONTEXT(context, input3_shape);44+ gert::Shape bcShape;
41- auto output0_shape = context->GetOutputShape(OUTPUT0_IDX);45+ // 逐对折叠广播: 只用两参数重载。vector 重载虽在 op_common/op_host/infershape_broadcast_util.h
42- OP_CHECK_NULL_WITH_CONTEXT(context, output0_shape);46+ // 以 Ops::Base 声明, 但 libops_base.so 只导出两参数版, vector 版的实现在 libop_common.so 的
43- auto outputv_shape = context->GetOutputShape(OUTPUTV_IDX);47+ // 小写 ops 命名空间下 —— 声明与实现命名空间不一致, 用它会编译期通过、加载期
44- OP_CHECK_NULL_WITH_CONTEXT(context, outputv_shape);48+ // undefined symbol 而装不上包。
45- auto outputm_shape = context->GetOutputShape(OUTPUTM_IDX);49+ bcShape = *inShapes[0];
46- OP_CHECK_NULL_WITH_CONTEXT(context, outputm_shape);50+ for (size_t i = 1; i < inShapes.size(); i++) {
47- 51+ gert::Shape tmp;
48- // grad、inputv、inputm 三者形状必须完全相同：inputv/inputm 是原地更新的动量输出，52+ OP_CHECK_IF(!BroadcastShape(&bcShape, inShapes[i], &tmp),
49- // 输出形状由它们决定；grad 绑定在广播 DAG 的 In0 位，底层 Ops::Base 广播模板53+ OP_LOGE(context->GetNodeName(), "input shapes cannot broadcast together"), return ge::GRAPH_FAILED);
50- // (DoDimensionCollapse) 不支持对 In0 做广播——实测 grad 为标量、或任一维为 1 时54+ bcShape = tmp;
51- // 均在 tiling 阶段被拒(“dim num is not same”/“dim index is not same with out”)。55+ }
52- // 故此处直接按等形拒绝，避免放行后到 tiling 才抛 E90003。56+ // 标量归一: 全标量输入 broadcast 得 0 维空 shape (), 与 A2 的 shape_util.scalar2tensor_one
53- // 仅 input3 参与广播(右对齐，维度数可少于 inputv)。57+ // 对齐, 归一为 (1,)。否则动态 shape 编译期 DFX 生成会对空 shape 做 reduce 连乘(无初值)而报
54- // 此处的判定与 tiling 的 CheckInplaceShapeConstraint 保持一致，避免两个 host58+ // TypeError 编译失败。
55- // 阶段对同一组合给出不同结论。59+ if (bcShape.GetDimNum() == 0) {
56- OP_CHECK_IF(!(*inputv_shape == *inputm_shape),60+ bcShape.SetDimNum(1);
57- OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(61+ bcShape.SetDim(0, 1);
58- context->GetNodeName(), "inputv and inputm",62+ }
59- (ToString(*inputv_shape) + " and " + ToString(*inputm_shape)).c_str(),63+ const gert::Shape& vShape = *inShapes[INPUTV_IDX];
60- "inputv and inputm are in-place updated moments and must have the same shape"),64+ const gert::Shape& mShape = *inShapes[INPUTM_IDX];
61- return ge::GRAPH_FAILED);65+ OP_CHECK_IF(
62- 66+ !(vShape == mShape),
63- gert::Shape broadcast_shape;67+ OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(
64- OP_CHECK_IF(!(*grad_shape == *inputv_shape),68+ context->GetNodeName(), "inputv and inputm", (ToString(vShape) + " and " + ToString(mShape)).c_str(),
69+ "inputv and inputm are in-place updated moments and must have the same shape"),
70+ return ge::GRAPH_FAILED);
71+ gert::Shape normV = vShape;
72+ if (normV.GetDimNum() == 0) {
73+ normV.SetDimNum(1);
74+ normV.SetDim(0, 1);
75+ }
76+ OP_CHECK_IF(!(bcShape == normV),
65 OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(77 OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
66- context->GetNodeName(), "grad", ToString(*grad_shape).c_str(),78+ context->GetNodeName(), "inputv", ToString(vShape).c_str(),
67- "grad does not support broadcast and must have exactly the same shape as inputv/inputm"),79+ "inputv/inputm are in-place(ref) outputs, so the broadcast shape of all inputs must equal them"),
68 return ge::GRAPH_FAILED);80 return ge::GRAPH_FAILED);
69- 81+ for (size_t i = 0; i < OUT_NUM; i++) {
70- OP_CHECK_IF(!BroadcastShape(input3_shape, inputv_shape, &broadcast_shape) || !(broadcast_shape == *inputv_shape),82+ auto out = context->GetOutputShape(i);
71- OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(83+ OP_CHECK_NULL_WITH_CONTEXT(context, out);
72- context->GetNodeName(), "input3", ToString(*input3_shape).c_str(),84+ *out = normV;
73- "input3 must be broadcastable into the in-place moment shape inputv/inputm"),85+ }
74- return ge::GRAPH_FAILED);
75- 
76- *output0_shape = *inputv_shape;
77- *outputv_shape = *inputv_shape;
78- *outputm_shape = *inputm_shape;
79- 
80 return GRAPH_SUCCESS;86 return GRAPH_SUCCESS;
81}87}
88+ 
82static ge::graphStatus InferDataType4LambApplyOptimizerAssign(gert::InferDataTypeContext* context)89static ge::graphStatus InferDataType4LambApplyOptimizerAssign(gert::InferDataTypeContext* context)
83{90{
84 if (context == nullptr) {91 if (context == nullptr) {
Aoptim/lamb_apply_optimizer_assign/op_host/lamb_apply_optimizer_assign_tiling_def.h+52-0
@@ -0,0 +1,52 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_apply_optimizer_assign_tiling_def.h
13+ * \brief LambApplyOptimizerAssign tiling data definition (arch35 手写内核)
14+ *
15+ * 只用于给 framework 定 raw tiling buffer 的容量; 实际下发由 host 侧
16+ * memcpy 整个 LambBrcTilingData<12, 3> POD 完成, 故字段与 POD 严格一一对应。
17+ */
18+ 
19+#ifndef LAMB_APPLY_OPTIMIZER_ASSIGN_TILING_DEF_H
20+#define LAMB_APPLY_OPTIMIZER_ASSIGN_TILING_DEF_H
21+ 
22+#include "register/tilingdata_base.h"
23+#include "register/op_impl_registry.h"
24+ 
25+namespace optiling {
26+constexpr int32_t LAMB_APPLY_OPTIMIZER_ASSIGN_IN_NUM = 12;
27+constexpr int32_t LAMB_APPLY_OPTIMIZER_ASSIGN_MAX_DIM = 8;
28+constexpr int32_t LAMB_APPLY_OPTIMIZER_ASSIGN_STRIDE_NUM = LAMB_APPLY_OPTIMIZER_ASSIGN_IN_NUM *
29+ LAMB_APPLY_OPTIMIZER_ASSIGN_MAX_DIM;
30+ 
31+BEGIN_TILING_DATA_DEF(LambApplyOptimizerAssignTilingData)
32+TILING_DATA_FIELD_DEF(uint64_t, totalNum);
33+TILING_DATA_FIELD_DEF(uint64_t, totalRows);
34+TILING_DATA_FIELD_DEF_ARR(uint64_t, LAMB_APPLY_OPTIMIZER_ASSIGN_STRIDE_NUM, effStride);
35+TILING_DATA_FIELD_DEF_ARR(uint64_t, LAMB_APPLY_OPTIMIZER_ASSIGN_IN_NUM, srcBlockLen);
36+TILING_DATA_FIELD_DEF(uint64_t, perCoreElems);
37+TILING_DATA_FIELD_DEF(uint64_t, rowsPerCore);
38+TILING_DATA_FIELD_DEF(uint32_t, tileLen);
39+TILING_DATA_FIELD_DEF(uint32_t, blockLen);
40+TILING_DATA_FIELD_DEF(uint32_t, rowsPerTile);
41+TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum);
42+TILING_DATA_FIELD_DEF(uint32_t, splitAxis);
43+TILING_DATA_FIELD_DEF(uint32_t, collapsedRank);
44+TILING_DATA_FIELD_DEF(uint32_t, tilingKey);
45+TILING_DATA_FIELD_DEF_ARR(uint32_t, LAMB_APPLY_OPTIMIZER_ASSIGN_IN_NUM, inKind);
46+TILING_DATA_FIELD_DEF_ARR(uint32_t, LAMB_APPLY_OPTIMIZER_ASSIGN_MAX_DIM, outShape);
47+TILING_DATA_FIELD_DEF_ARR(uint32_t, LAMB_APPLY_OPTIMIZER_ASSIGN_STRIDE_NUM, inShape);
48+END_TILING_DATA_DEF;
49+ 
50+REGISTER_TILING_DATA_CLASS(LambApplyOptimizerAssign, LambApplyOptimizerAssignTilingData)
51+} // namespace optiling
52+#endif // LAMB_APPLY_OPTIMIZER_ASSIGN_TILING_DEF_H
Aoptim/lamb_apply_optimizer_assign/op_kernel/CMakeLists.txt+16-0
@@ -0,0 +1,16 @@
1+# ----------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# ----------------------------------------------------------------------------
10+ 
11+# 手写 regbase 内核自行下发 SetFlag/WaitFlag, 关闭自动插同步。
12+add_kernel_sources(
13+ KERNEL_SRC arch35/lamb_apply_optimizer_assign.cpp
14+ COMPUTE_UNITS ascend950
15+ AUTO_SYNC false
16+)
Aoptim/lamb_apply_optimizer_assign/op_kernel/arch35/lamb_apply_optimizer_assign.cpp+48-0
@@ -0,0 +1,48 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_apply_optimizer_assign.cpp
13+ * \brief LambApplyOptimizerAssign arch35 (Ascend950) kernel entry
14+ *
15+ * 手写 regbase 实现, 不走 ATVOSS: 12 入 3 出在 DAGSch 的 32 buffer 预算上编不过,
16+ * 且 ATVOSS 的 Vec::Brc 未实现导致 UB 广播档不可用(铺不平尾轴 1->n 与低秩补维)。
17+ * 搬运/广播/分核骨架见 lamb_apply_common/op_kernel/arch35/lamb_brc_kernel.h。
18+ * TilingKey: fp32=100 / fp16=200。
19+ */
20+ 
21+#include "kernel_operator.h"
22+#include "lamb_apply_optimizer_assign_vf.h"
23+ 
24+// 宏展开处需要一个文件作用域可见的具体类型名, 不能用函数内的局部别名。
25+using LambApplyOptimizerAssignTilingType = LambBrcTilingData<12, 3>;
26+ 
27+extern "C" __global__ __aicore__ void lamb_apply_optimizer_assign(
28+ GM_ADDR grad, GM_ADDR inputv, GM_ADDR inputm, GM_ADDR input4, GM_ADDR mul0_x, GM_ADDR mul1_x, GM_ADDR mul2_x,
29+ GM_ADDR mul3_x, GM_ADDR add2_y, GM_ADDR steps, GM_ADDR do_use_weight, GM_ADDR weight_decay_rate, GM_ADDR output0,
30+ GM_ADDR inputv_ref, GM_ADDR inputm_ref, GM_ADDR workspace, GM_ADDR tiling)
31+{
32+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
33+ REGISTER_TILING_DEFAULT(LambApplyOptimizerAssignTilingType);
34+ GET_TILING_DATA_WITH_STRUCT(LambApplyOptimizerAssignTilingType, tilingData, tiling);
35+ GM_ADDR inAddr[12] = {grad, inputv, inputm, input4, mul0_x, mul1_x,
36+ mul2_x, mul3_x, add2_y, steps, do_use_weight, weight_decay_rate};
37+ GM_ADDR outAddr[3] = {output0, inputv_ref, inputm_ref};
38+ AscendC::TPipe pipe;
39+ if (TILING_KEY_IS(100)) {
40+ LambBrc::BrcElementwiseKernel<float, 12, 3, LambApplyOptimizerAssignOp::LambApplyOptimizerAssignVf> op;
41+ op.Init(inAddr, outAddr, &tilingData, &pipe);
42+ op.Process();
43+ } else if (TILING_KEY_IS(200)) {
44+ LambBrc::BrcElementwiseKernel<half, 12, 3, LambApplyOptimizerAssignOp::LambApplyOptimizerAssignVf> op;
45+ op.Init(inAddr, outAddr, &tilingData, &pipe);
46+ op.Process();
47+ }
48+}
Doptim/lamb_apply_optimizer_assign/op_kernel/arch35/lamb_apply_optimizer_assign_dag.h+0-108
@@ -1,108 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_apply_optimizer_assign_dag.h
13- * \brief lamb_apply_optimizer_assign_dag head file
14- */
15- 
16-#ifndef LAMB_APPLY_OPTIMIZER_ASSIGN_DAG_H
17-#define LAMB_APPLY_OPTIMIZER_ASSIGN_DAG_H
18-#include "atvoss/util/dag.h"
19-#include "atvoss/util/vec.h"
20-#include "atvoss/util/placeholder.h"
21- 
22-namespace LambApplyOptimizerAssignOp {
23-using namespace AscendC;
24-using namespace Ops::Base;
25-// In : 0 grad, 1 inputv(v), 2 inputm(m), 3 input3(param), 4 mul0_x(b1), 5 mul1_x(1-b1),
26-// 6 mul2_x(b2), 7 mul3_x(1-b2), 8 add2_y(eps), 9 steps, 10 do_use_weight, 11 weight_decay_rate
27-// Out: 0 update, 1 inputv(next_v, in-place), 2 inputm(next_m, in-place)
28-// Compute in U (=float): half casts to float (div/sqrt/log/exp need float precision); fp32 native.
29-template <typename T, typename U>
30-struct LambApplyOptimizerAssignCompute {
31- using CNegOne = MAKE_CONST(U, -1.0);
32- using COne = MAKE_CONST(U, 1.0);
33- 
34- using InGrad = Bind<Vec::CopyInBrc<T>, Placeholder::In0<T>>;
35- using InV = Bind<Vec::CopyInBrc<T>, Placeholder::In1<T>>;
36- using InM = Bind<Vec::CopyInBrc<T>, Placeholder::In2<T>>;
37- using InParam = Bind<Vec::CopyInBrc<T>, Placeholder::In3<T>>;
38- // 8 scalar [1] inputs: ScalarAttr + Duplicate (broadcast scalar to a tensor without a
39- // double-buffered mte2 buffer) so downstream tensor ops match and we stay in buffer budget.
40- using InB1 = Bind<Vec::Duplicate<T>, Placeholder::In4<T, Placeholder::ScalarAttr<true>>>;
41- using InOmB1 = Bind<Vec::Duplicate<T>, Placeholder::In5<T, Placeholder::ScalarAttr<true>>>;
42- using InB2 = Bind<Vec::Duplicate<T>, Placeholder::In6<T, Placeholder::ScalarAttr<true>>>;
43- using InOmB2 = Bind<Vec::Duplicate<T>, Placeholder::In7<T, Placeholder::ScalarAttr<true>>>;
44- using InEps = Bind<Vec::Duplicate<T>, Placeholder::In8<T, Placeholder::ScalarAttr<true>>>;
45- using InSteps = Bind<Vec::Duplicate<T>, Placeholder::In9<T, Placeholder::ScalarAttr<true>>>;
46- using InDoUse = Bind<Vec::Duplicate<T>, Placeholder::In10<T, Placeholder::ScalarAttr<true>>>;
47- using InWd = Bind<Vec::Duplicate<T>, Placeholder::In11<T, Placeholder::ScalarAttr<true>>>;
48- 
49- using Grad = Bind<Vec::Cast<U, T, 0>, InGrad>;
50- using V = Bind<Vec::Cast<U, T, 0>, InV>;
51- using M = Bind<Vec::Cast<U, T, 0>, InM>;
52- using Param = Bind<Vec::Cast<U, T, 0>, InParam>;
53- using B1 = Bind<Vec::Cast<U, T, 0>, InB1>;
54- using OmB1 = Bind<Vec::Cast<U, T, 0>, InOmB1>;
55- using B2 = Bind<Vec::Cast<U, T, 0>, InB2>;
56- using OmB2 = Bind<Vec::Cast<U, T, 0>, InOmB2>;
57- using Eps = Bind<Vec::Cast<U, T, 0>, InEps>;
58- using Steps = Bind<Vec::Cast<U, T, 0>, InSteps>;
59- using DoUse = Bind<Vec::Cast<U, T, 0>, InDoUse>;
60- using Wd = Bind<Vec::Cast<U, T, 0>, InWd>;
61- 
62- // next_v = g^2 * (1 - b2) + v * b2
63- using G2 = Bind<Vec::Mul<U>, Grad, Grad>;
64- using NV1 = Bind<Vec::Mul<U>, G2, OmB2>;
65- using NV2 = Bind<Vec::Mul<U>, V, B2>;
66- using NextV = Bind<Vec::Add<U>, NV1, NV2>;
67- using NextVCast = Bind<Vec::Cast<T, U, 1>, NextV>;
68- 
69- // next_m = m * b1 + g * (1 - b1)
70- using NM1 = Bind<Vec::Mul<U>, M, B1>;
71- using NM2 = Bind<Vec::Mul<U>, Grad, OmB1>;
72- using NextM = Bind<Vec::Add<U>, NM1, NM2>;
73- using NextMCast = Bind<Vec::Cast<T, U, 1>, NextM>;
74- 
75- // bias correction: b_corr = 1 - exp(steps * ln(b)) (b1/b2/steps are scalar [1])
76- using LnB1 = Bind<Vec::Log<U>, B1>;
77- using ExpArg1 = Bind<Vec::Mul<U>, LnB1, Steps>;
78- using B1Steps = Bind<Vec::Exp<U>, ExpArg1>;
79- using NegB1Steps = Bind<Vec::Muls<U>, B1Steps, CNegOne>;
80- using B1corr = Bind<Vec::Adds<U>, NegB1Steps, COne>;
81- using LnB2 = Bind<Vec::Log<U>, B2>;
82- using ExpArg2 = Bind<Vec::Mul<U>, LnB2, Steps>;
83- using B2Steps = Bind<Vec::Exp<U>, ExpArg2>;
84- using NegB2Steps = Bind<Vec::Muls<U>, B2Steps, CNegOne>;
85- using B2corr = Bind<Vec::Adds<U>, NegB2Steps, COne>;
86- 
87- // update = (next_m / b1corr) / (sqrt(next_v / b2corr) + eps) + param * wd * do_use_weight
88- using MUnb = Bind<Vec::Div<U>, NextM, B1corr>;
89- using VUnb = Bind<Vec::Div<U>, NextV, B2corr>;
90- using SqrtVUnb = Bind<Vec::Sqrt<U>, VUnb>;
91- using Den = Bind<Vec::Add<U>, SqrtVUnb, Eps>;
92- using Upd0 = Bind<Vec::Div<U>, MUnb, Den>;
93- using WdMul = Bind<Vec::Mul<U>, Param, Wd>;
94- using Wdec = Bind<Vec::Mul<U>, WdMul, DoUse>;
95- using Update = Bind<Vec::Add<U>, Wdec, Upd0>;
96- using UpdateCast = Bind<Vec::Cast<T, U, 1>, Update>;
97- 
98- using OpCopyOut0 = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, UpdateCast>;
99- using OpCopyOut1 = Bind<Vec::CopyOut<T>, Placeholder::Out1<T>, NextVCast>;
100- using OpCopyOut2 = Bind<Vec::CopyOut<T>, Placeholder::Out2<T>, NextMCast>;
101- 
102- using Outputs = Elems<OpCopyOut0, OpCopyOut1, OpCopyOut2>;
103- // 12 inputs + 3 outputs exceed the LEVEL_2 (double-buffered) 32-buffer budget; LEVEL_1 fits.
104- using MemCfg = MemOptCfg<MemLevel::LEVEL_1>;
105- using OpDag = DAGSch<Outputs, void, MemCfg>;
106-};
107-} // namespace LambApplyOptimizerAssignOp
108-#endif // LAMB_APPLY_OPTIMIZER_ASSIGN_DAG_H
Doptim/lamb_apply_optimizer_assign/op_kernel/arch35/lamb_apply_optimizer_assign_tiling_key.h+0-27
@@ -1,27 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_apply_optimizer_assign_tiling_key.h
13- * \brief lamb_apply_optimizer_assign_tiling_key head file
14- */
15- 
16-#ifndef ADAM_APPLY_ONE_STRUCT_H
17-#define ADAM_APPLY_ONE_STRUCT_H
18- 
19-#include "atvoss/broadcast/broadcast_base_struct.h"
20- 
21-using namespace Ops::Base;
22-// 算子自定义的tiling key字段
23-ASCENDC_TPL_ARGS_DECL(LambApplyOptimizerAssign, BRC_NDDMA_SCH_MODE_KEY_DECL(schMode));
24- 
25-ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(BRC_NDDMA_SCH_MODE_KEY_SEL(schMode)));
26- 
27-#endif // ADAM_APPLY_ONE_STRUCT_H
Aoptim/lamb_apply_optimizer_assign/op_kernel/arch35/lamb_apply_optimizer_assign_vf.h+152-0
@@ -0,0 +1,152 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_apply_optimizer_assign_vf.h
13+ * \brief LambApplyOptimizerAssign 的 regbase 计算段(12 入 3 出)
14+ *
15+ * 与 golden 同运算序:
16+ * next_v = v*b2 + (g*g)*(1-b2) -> out1
17+ * next_m = m*b1 + g*(1-b1) -> out2
18+ * b_corr = 1 + (-1)*exp(ln(b)*steps) (Log/Mul/Exp/Muls/Adds 五步, 不用幂运算:
19+ * b 接近 1 时 ln(b) 有相消损失, 直接幂得不到同一中间量)
20+ * update = (next_m/b1corr)/(sqrt(next_v/b2corr)+eps) + param*wd*do_use -> out0
21+ */
22+ 
23+#ifndef LAMB_APPLY_OPTIMIZER_ASSIGN_VF_H
24+#define LAMB_APPLY_OPTIMIZER_ASSIGN_VF_H
25+ 
26+#include "../lamb_apply_common/arch35/lamb_brc_kernel.h"
27+ 
28+namespace LambApplyOptimizerAssignOp {
29+using namespace AscendC;
30+ 
31+enum : uint32_t {
32+ IN_GRAD = 0,
33+ IN_V = 1,
34+ IN_M = 2,
35+ IN_PARAM = 3,
36+ IN_B1 = 4,
37+ IN_OMB1 = 5,
38+ IN_B2 = 6,
39+ IN_OMB2 = 7,
40+ IN_EPS = 8,
41+ IN_STEPS = 9,
42+ IN_DOUSE = 10,
43+ IN_WD = 11
44+};
45+ 
46+struct LambApplyOptimizerAssignVf {
47+ template <typename T>
48+ static __aicore__ inline void Run(__local_mem__ T** in, __local_mem__ T** out, __local_mem__ float** scratch,
49+ uint32_t count)
50+ {
51+ uint16_t vfTimes = static_cast<uint16_t>(count / LambBrc::LAMB_VF_LEN);
52+ uint32_t tail = count % LambBrc::LAMB_VF_LEN;
53+ uint16_t tailTimes = (tail > 0) ? 1 : 0;
54+ 
55+ __local_mem__ T* pG = in[IN_GRAD];
56+ __local_mem__ T* pV = in[IN_V];
57+ __local_mem__ T* pM = in[IN_M];
58+ __local_mem__ T* pParam = in[IN_PARAM];
59+ __local_mem__ T* pB1 = in[IN_B1];
60+ __local_mem__ T* pOmB1 = in[IN_OMB1];
61+ __local_mem__ T* pB2 = in[IN_B2];
62+ __local_mem__ T* pOmB2 = in[IN_OMB2];
63+ __local_mem__ T* pEps = in[IN_EPS];
64+ __local_mem__ T* pSteps = in[IN_STEPS];
65+ __local_mem__ T* pDoUse = in[IN_DOUSE];
66+ __local_mem__ T* pWd = in[IN_WD];
67+ __local_mem__ T* update = out[0];
68+ __local_mem__ T* nextVOut = out[1];
69+ __local_mem__ T* nextMOut = out[2];
70+ // 中间量以 fp32 暂存, 不经 T 的舍入 —— 与单段实现数值完全一致。
71+ __local_mem__ float* sNextV = scratch[0];
72+ __local_mem__ float* sNextM = scratch[1];
73+ 
74+ // ---- 段1: next_v / next_m ----
75+ // 拆成两段是因为全量精确档(0ULP/FTZ_FALSE)展开后单段指令过多, 循环回边偏移超出
76+ // scbzi 立即数范围 [-512,511]。拆段后每段都在范围内, 且不必降低任何一处的精度档。
77+ __VEC_SCOPE__
78+ {
79+ Reg::RegTensor<float> vA, vB, vC, vNextV, vNextM;
80+ Reg::MaskReg maskAll = Reg::CreateMask<float, Reg::MaskPattern::ALL>();
81+ Reg::MaskReg maskT = Reg::UpdateMask<float>(tail);
82+ for (uint16_t vfIdx = 0; vfIdx < vfTimes + tailTimes; vfIdx++) {
83+ uint32_t off = vfIdx * LambBrc::LAMB_VF_LEN;
84+ Reg::MaskReg preg = (vfIdx < vfTimes) ? maskAll : maskT;
85+ // next_v = v*b2 + (g*g)*(1-b2)
86+ LambBrc::Load<T>(pV, vA, preg, off);
87+ LambBrc::Load<T>(pB2, vB, preg, off);
88+ Reg::Mul(vNextV, vA, vB, preg);
89+ LambBrc::Load<T>(pG, vA, preg, off);
90+ Reg::Mul(vC, vA, vA, preg);
91+ LambBrc::Load<T>(pOmB2, vB, preg, off);
92+ Reg::Mul(vC, vC, vB, preg);
93+ Reg::Add(vNextV, vNextV, vC, preg);
94+ LambBrc::Store<T>(nextVOut, vNextV, preg, off);
95+ Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM>(sNextV + off, vNextV, preg);
96+ // next_m = m*b1 + g*(1-b1) (vA 仍是 g)
97+ LambBrc::Load<T>(pOmB1, vB, preg, off);
98+ Reg::Mul(vC, vA, vB, preg);
99+ LambBrc::Load<T>(pM, vA, preg, off);
100+ LambBrc::Load<T>(pB1, vB, preg, off);
101+ Reg::Mul(vNextM, vA, vB, preg);
102+ Reg::Add(vNextM, vNextM, vC, preg);
103+ LambBrc::Store<T>(nextMOut, vNextM, preg, off);
104+ Reg::StoreAlign<float, Reg::StoreDist::DIST_NORM>(sNextM + off, vNextM, preg);
105+ }
106+ }
107+ 
108+ // ---- 段2: 偏差校正 + update ----
109+ __VEC_SCOPE__
110+ {
111+ Reg::RegTensor<float> vA, vB, vC, vNextV, vNextM, vB1c, vB2c, vT1, vT2;
112+ Reg::MaskReg cmpExpm1;
113+ Reg::MaskReg maskAll = Reg::CreateMask<float, Reg::MaskPattern::ALL>();
114+ Reg::MaskReg maskT = Reg::UpdateMask<float>(tail);
115+ for (uint16_t vfIdx = 0; vfIdx < vfTimes + tailTimes; vfIdx++) {
116+ uint32_t off = vfIdx * LambBrc::LAMB_VF_LEN;
117+ Reg::MaskReg preg = (vfIdx < vfTimes) ? maskAll : maskT;
118+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vNextV, sNextV + off);
119+ Reg::LoadAlign<float, Reg::LoadDist::DIST_NORM>(vNextM, sNextM + off);
120+ // Ln/Exp 取 1ULP + FTZ_TRUE: 默认 INTRINSIC 是硬件近似档, 而 b_corr 的抵消会把 exp 的误差
121+ // 放大 1/(1-exp) 倍; 取 FTZ_TRUE 是因为 ln(b)*steps 低于 exp 下溢边界时真值本就应为 0。
122+ // 1 - b^steps 由 OneMinusExp 计算(见 lamb_brc_kernel.h)。
123+ // b1corr = 1 - exp(ln(b1)*steps)
124+ LambBrc::Load<T>(pSteps, vC, preg, off);
125+ LambBrc::Load<T>(pB1, vB, preg, off);
126+ Ln<float, &LambBrc::LAMB_PRECISE_LN_FTZT>(vB1c, vB, preg);
127+ Reg::Mul(vB1c, vB1c, vC, preg);
128+ LambBrc::OneMinusExp(vB1c, vB1c, vT1, vT2, cmpExpm1, preg);
129+ LambBrc::Load<T>(pB2, vB, preg, off);
130+ Ln<float, &LambBrc::LAMB_PRECISE_LN_FTZT>(vB2c, vB, preg);
131+ Reg::Mul(vB2c, vB2c, vC, preg);
132+ LambBrc::OneMinusExp(vB2c, vB2c, vT1, vT2, cmpExpm1, preg);
133+ // update = (next_m/b1corr)/(sqrt(next_v/b2corr)+eps) + param*wd*do_use
134+ Reg::Div<float, &LambBrc::LAMB_PRECISE_DIV>(vA, vNextM, vB1c, preg);
135+ Reg::Div<float, &LambBrc::LAMB_PRECISE_DIV>(vB, vNextV, vB2c, preg);
136+ Reg::Sqrt<float, &LambBrc::LAMB_PRECISE_SQRT>(vB, vB, preg);
137+ LambBrc::Load<T>(pEps, vC, preg, off);
138+ Reg::Add(vB, vB, vC, preg);
139+ Reg::Div<float, &LambBrc::LAMB_PRECISE_DIV>(vA, vA, vB, preg);
140+ LambBrc::Load<T>(pParam, vB, preg, off);
141+ LambBrc::Load<T>(pWd, vC, preg, off);
142+ Reg::Mul(vB, vB, vC, preg);
143+ LambBrc::Load<T>(pDoUse, vC, preg, off);
144+ Reg::Mul(vB, vB, vC, preg);
145+ Reg::Add(vA, vA, vB, preg);
146+ LambBrc::Store<T>(update, vA, preg, off);
147+ }
148+ }
149+ }
150+};
151+} // namespace LambApplyOptimizerAssignOp
152+#endif // LAMB_APPLY_OPTIMIZER_ASSIGN_VF_H
Doptim/lamb_apply_optimizer_assign/op_kernel/lamb_apply_optimizer_assign.cpp+0-44
@@ -1,44 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_apply_optimizer_assign.cpp
13- * \brief lamb_apply_optimizer_assign.cpp
14- */
15- 
16-#include "kernel_operator.h"
17-#include "arch35/lamb_apply_optimizer_assign_dag.h"
18-#include "arch35/lamb_apply_optimizer_assign_tiling_key.h"
19-#include "atvoss/broadcast/broadcast_sch.h"
20- 
21-using namespace AscendC;
22-using namespace Ops::Base;
23- 
24-template <uint64_t schMode>
25-__global__ __aicore__ void lamb_apply_optimizer_assign(GM_ADDR grad, GM_ADDR inputv, GM_ADDR inputm, GM_ADDR input3,
26- GM_ADDR mul0_x, GM_ADDR mul1_x, GM_ADDR mul2_x, GM_ADDR mul3_x,
27- GM_ADDR add2_y, GM_ADDR steps, GM_ADDR do_use_weight,
28- GM_ADDR weight_decay_rate, GM_ADDR output0, GM_ADDR output1,
29- GM_ADDR output2, GM_ADDR workspace, GM_ADDR tiling)
30-{
31- KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
32- // half/fp32 both compute in float (div/sqrt/log/exp need float precision).
33- if constexpr (std::is_same<DTYPE_GRAD, half>::value) {
34- using OpDag = LambApplyOptimizerAssignOp::LambApplyOptimizerAssignCompute<half, float>::OpDag;
35- BroadcastSch<schMode, OpDag> sch(tiling);
36- sch.Process(grad, inputv, inputm, input3, mul0_x, mul1_x, mul2_x, mul3_x, add2_y, steps, do_use_weight,
37- weight_decay_rate, output0, output1, output2);
38- } else {
39- using OpDag = LambApplyOptimizerAssignOp::LambApplyOptimizerAssignCompute<float, float>::OpDag;
40- BroadcastSch<schMode, OpDag> sch(tiling);
41- sch.Process(grad, inputv, inputm, input3, mul0_x, mul1_x, mul2_x, mul3_x, add2_y, steps, do_use_weight,
42- weight_decay_rate, output0, output1, output2);
43- }
44-}
Moptim/lamb_apply_optimizer_assign/tests/assets/golden.py+60-15
@@ -28,13 +28,17 @@ __golden__ = {
28 28 
29def _scalars(*xs):29def _scalars(*xs):
30 """取标量并落到 float32 torch 标量张量上。不能返回 Python float——那是 fp64，标量30 """取标量并落到 float32 torch 标量张量上。不能返回 Python float——那是 fp64，标量
31- 运算会被抬到双精度，而算子在 fp32 上算（A2 的 TBE compute 里 dtype='float32'，31+ 运算会被抬到双精度，而 arch35 内核的计算类型 U = float（fp16 输入 unpack 成 fp32 再算）。
32+ 注：A2 并**不**升精度 —— canndev 的 tbe impl 是 tvm.placeholder(dtype=input_dtype)，全程无 cast_to，
33+ fp16 输入就在 fp16 上做 vmul/vdiv/vsqrt。此处跟随 arch35 的计算类型，不跟随 A2（
32 arch35 DAG 的计算类型 U = float）。numpy 只用于取值与 dtype 转换。"""34 arch35 DAG 的计算类型 U = float）。numpy 只用于取值与 dtype 转换。"""
33 # 标量落成 **0 维 float64**: torch 类型提升里 0 维不抬 dim>0 张量的档,35 # 标量落成 **0 维 float64**: torch 类型提升里 0 维不抬 dim>0 张量的档,
34 # 故数据 fp32 时结果仍 fp32、被 Promote 成 fp64 时标量自动跟到 fp64。36 # 故数据 fp32 时结果仍 fp32、被 Promote 成 fp64 时标量自动跟到 fp64。
35- return tuple(37+ # A2 语义: 这些"系数"输入是**可广播的 ND Tensor**(canndev ops/built-in/tbe/impl/lamb_*.py
36- torch.from_numpy(np.asarray(x, "float64").reshape(-1)[:1])[0] for x in xs38+ # 每步 mul/sub/div 都先 shape_util.broadcast_shapes 再 tbe.broadcast)。原先 reshape(-1)[:1]
37- )39+ # 只取首元素, 传多元素张量时静默按首元素计算 —— 与内核的广播实现不一致, 广播档必然假红。
40+ # 改为返回完整张量交给 torch 自然广播: 形状 (1,) 的行为与原标量完全一致, 故常规档不变。
41+ return tuple(_t(x) for x in xs)
38 42 
39 43 
40def _t(x):44def _t(x):
@@ -69,7 +73,7 @@ def lamb_apply_optimizer_assign_golden(
69):73):
70 """Golden for LambApplyOptimizerAssign. Params follow lamb_apply_optimizer_assign_def.cpp (without outputs). All inputs are numpy.ndarray.74 """Golden for LambApplyOptimizerAssign. Params follow lamb_apply_optimizer_assign_def.cpp (without outputs). All inputs are numpy.ndarray.
71 75 
72- Computed by composing torch tensor ops (torch.add/torch.addcmul/torch.sqrt) instead of a76+ Computed by composing torch tensor ops (torch.mul/torch.add/torch.sqrt/torch.expm1) instead of a
73 hand-written numpy formula: red line R3 requires the golden to be a competitor-operator77 hand-written numpy formula: red line R3 requires the golden to be a competitor-operator
74 composition, and a naive numpy expression tends to make exactly the same rounding mistakes78 composition, and a naive numpy expression tends to make exactly the same rounding mistakes
75 as the kernel under test, which would disguise a precision shortfall as a pass.79 as the kernel under test, which would disguise a precision shortfall as a pass.
@@ -83,13 +87,16 @@ def lamb_apply_optimizer_assign_golden(
83 # 与算子定义的「Muls 再 Add」两步舍入不是同一个运算序列。golden 要如实转写定义。87 # 与算子定义的「Muls 再 Add」两步舍入不是同一个运算序列。golden 要如实转写定义。
84 next_v = v * b2 + (g * g) * omb288 next_v = v * b2 + (g * g) * omb2
85 next_m = m * b1 + g * omb189 next_m = m * b1 + g * omb1
86- # 偏差校正按内核的算法写：arch35 DAG 用 Log/Mul/Exp 三条指令实现幂90+ # 偏差校正 corr = 1 - b^steps。**不能**照抄内核的 `1 + (-1)*exp(ln(b)*steps)` 指令序列：
87- # LnB1=Log(B1); ExpArg1=Mul(LnB1,Steps); B1Steps=Exp(ExpArg1)91+ # 当 ln(b)*steps -> 0 时 exp(.) -> 1，该式是灾难性抵消 —— fp32 在 1.0 附近的间距是 6e-8，
88- # NegB1Steps=Muls(B1Steps,-1); B1corr=Adds(NegB1Steps,1)92+ # 结果为 1e-6 量级时只剩 4 位有效位。照抄会让 golden 和被测一样失准，失去参照资格
89- # 而不是 b1 ** t。两者数学等价、浮点下不等价：Log/Exp 各是标称 1 ULP 的单指令，93+ # （精度标准 §4.5 要求"以更高精度的 CPU 实现为真值"；本文件 R3 注释也写明 golden 不得
90- # 且 b1 接近 1 时 ln(b1) 有相消损失，直接幂运算得不到同一个中间量。94+ # 复制被测的舍入错误）。
91- b1_corr = 1.0 + (-1.0) * torch.exp(torch.log(b1) * t)95+ # 用 expm1 在 fp64 上算再落回工作精度：fp32 能精确表示 1e-6 量级，丢精度的是那步减法，
92- b2_corr = 1.0 + (-1.0) * torch.exp(torch.log(b2) * t)96+ # 不是表示能力。内核侧同样绕开该减法（小 |x| 走 Taylor 支），两边才在同一个真值上。
97+ _hi = torch.float64
98+ b1_corr = (-torch.expm1(torch.log(b1.to(_hi)) * t.to(_hi))).to(b1.dtype)
99+ b2_corr = (-torch.expm1(torch.log(b2.to(_hi)) * t.to(_hi))).to(b2.dtype)
93 update = (next_m / b1_corr) / (torch.sqrt(next_v / b2_corr) + eps) + w * wd * du100 update = (next_m / b1_corr) / (torch.sqrt(next_v / b2_corr) + eps) + w * wd * du
94 return [101 return [
95 update.numpy().astype(dt),102 update.numpy().astype(dt),
@@ -125,8 +132,32 @@ def _tp_t(x):
125 )132 )
126 133 
127 134 
135+def _tp_widen(t):
136+ """按 NPU 的加宽行为把三方入参落到**内核计算类型 U = float32**。
137+ 
138+ 规范依据 ttk_golden_logic.md §四/§五「浮点 + 三方」一格: 三方腿"按 NPU 加宽行为同步 cast"。
139+ 内核对 fp16 输入 unpack 成 fp32 计算(dag.h 的 `Cast<U, T, 0>`, U = float), 对 fp32 输入
140+ 原生 fp32, 两种情况计算类型都是 fp32, 故此处一律落到 fp32。
141+ 
142+ **不能只处理 fp16**: cross_check 下 TTK 会把 fp32 Promote 成 fp64 下发, 原样透传会让三方腿
143+ 在 fp64 上算 —— 内核 fp32 溢出/下溢的地方它都不溢出, 两条腿不在同一精度上, 比值无意义。
144+ 整型不动: 内核对 int 是原生/int32 累加, 走 fp32 会抹掉 >2^24 的低位。
145+ """
146+ return t.to(torch.float32) if t.is_floating_point() else t
147+ 
148+ 
149+def _tp_narrow(outs, dt):
150+ """出口复刻内核的 `Cast<T, U, 1>`: 窄回算子**声明**的 dtype T(不是 Promote 后的 dtype)。
151+ 
152+ 少了这一步或窄错目标, 三方腿会与走 Promote 的 golden 逐位相等 —— 双标杆塌成单标杆,
153+ 三比值分母被 safe_div 的 small_value 夹底, mare/rmse 恒为 1.0, 阈值永不触发。
154+ """
155+ return [o.to(dt) if o.is_floating_point() else o for o in outs]
156+ 
157+ 
128def _tp_s(x):158def _tp_s(x):
129- return _tp_t(x).reshape(-1)[0]159+ # 同 _scalars: 三方腿也必须广播, 不能只取首元素
160+ return _tp_t(x)
130 161 
131 162 
132class _LambApplyOptimizerAssignCompose:163class _LambApplyOptimizerAssignCompose:
@@ -164,8 +195,11 @@ class _LambApplyOptimizerAssignCompose:
164 # 不能用 addcmul 的 FMA 单次舍入, 否则竞品凭空更准, ratio 误判内核。195 # 不能用 addcmul 的 FMA 单次舍入, 否则竞品凭空更准, ratio 误判内核。
165 next_v = v * b2 + (g * g) * omb2196 next_v = v * b2 + (g * g) * omb2
166 next_m = m * b1 + g * omb1197 next_m = m * b1 + g * omb1
167- b1_corr = 1.0 - torch.pow(b1, t)198+ # 竞品的精确写法: torch 自带 expm1, 1-b^t 用 -expm1(log(b)*t) 避开对 1.0 的减法。
168- b2_corr = 1.0 - torch.pow(b2, t)199+ # 若写成 1-pow(b,t), b^t->1 时抵消归零 -> 除零 -> 整片 NaN, 三方腿会比 NPU 和 golden
200+ # 都差, 成为稻草人标杆 —— 比值 err_npu/err_third 反而更易通过, 是放水不是严格。
201+ b1_corr = -torch.expm1(torch.log(b1) * t)
202+ b2_corr = -torch.expm1(torch.log(b2) * t)
169 update = (next_m / b1_corr) / (torch.sqrt(next_v / b2_corr) + eps) + w * wd * du203 update = (next_m / b1_corr) / (torch.sqrt(next_v / b2_corr) + eps) + w * wd * du
170 return [update, next_v, next_m]204 return [update, next_v, next_m]
171 205 
@@ -184,6 +218,16 @@ class _LambApplyOptimizerAssignCompose:
184 218 
185 219 
186def _kf_widen(a, seen):220def _kf_widen(a, seen):
221+ """按 NPU 的加宽行为把三方入参加宽, 并记下算子声明的 dtype T 供出口窄回。
222+ 
223+ 三方腿拿到的是**原始 dtype T**(TTK 的 Promote 只作用于 golden, 见 profiling.py 的
224+ golden_mode_override; 三方腿走 _xpu_inputs -> original_input_arrays), 故这里按 T 判断:
225+ - T = fp16/bf16: 内核 unpack 成 fp32 全程不落回(dag.h 的 Cast<U, T, 0>, U = float),
226+ 而 torch 只在单个算子内部用 opmath=float、算子之间每步落回 T。不加宽就等于拿
227+ "逐步截断的实现"当竞品, 与被测内核不是同一个算法 -> 加宽到 fp32, 出口窄回 T。
228+ - T = fp32: 内核计算类型 U 即 fp32, **不加宽**; torch 同样原生 fp32 -> 原样不动。
229+ - 整型: 内核原生/int32 累加, 走 fp32 会抹掉 >2^24 的低位 -> 不动。
230+ """
187 if isinstance(a, (list, tuple)):231 if isinstance(a, (list, tuple)):
188 return type(a)(_kf_widen(x, seen) for x in a)232 return type(a)(_kf_widen(x, seen) for x in a)
189 if isinstance(a, torch.Tensor):233 if isinstance(a, torch.Tensor):
@@ -206,6 +250,7 @@ def _kf_widen(a, seen):
206 250 
207 251 
208def _kf_narrow(o, dt):252def _kf_narrow(o, dt):
253+ """出口复刻内核的 `Cast<T, U, 1>`: 窄回算子声明的 dtype T。T = fp32 时无需窄回。"""
209 if isinstance(o, (list, tuple)):254 if isinstance(o, (list, tuple)):
210 return type(o)(_kf_narrow(x, dt) for x in o)255 return type(o)(_kf_narrow(x, dt) for x in o)
211 if isinstance(o, torch.Tensor) and o.is_floating_point():256 if isinstance(o, torch.Tensor) and o.is_floating_point():
Moptim/lamb_apply_optimizer_assign/tests/ut/op_host/test_lamb_apply_optimizer_assign_infershape.cpp+27-5
@@ -107,14 +107,15 @@ namespace {
107// 按 grad / inputv / inputm / input3 四个张量的形状跑一次 infershape，107// 按 grad / inputv / inputm / input3 四个张量的形状跑一次 infershape，
108// 标量输入统一用 {1}。用于核对 infershape 与 tiling 的支持范围是否一致。108// 标量输入统一用 {1}。用于核对 infershape 与 tiling 的支持范围是否一致。
109ge::graphStatus RunInferShape(const gert::Shape& grad, const gert::Shape& inputv, const gert::Shape& inputm,109ge::graphStatus RunInferShape(const gert::Shape& grad, const gert::Shape& inputv, const gert::Shape& inputm,
110- const gert::Shape& input3, gert::Shape* out0, gert::Shape* out1, gert::Shape* out2)110+ const gert::Shape& input3, gert::Shape* out0, gert::Shape* out1, gert::Shape* out2,
111+ const gert::Shape& scalar = gert::Shape({1}))
111{112{
112 auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyOptimizerAssign")->infer_shape;113 auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyOptimizerAssign")->infer_shape;
113 gert::Shape g = grad;114 gert::Shape g = grad;
114 gert::Shape v = inputv;115 gert::Shape v = inputv;
115 gert::Shape m = inputm;116 gert::Shape m = inputm;
116 gert::Shape p = input3;117 gert::Shape p = input3;
117- gert::Shape s = {1};118+ gert::Shape s = scalar;
118 auto holder = gert::InferShapeContextFaker()119 auto holder = gert::InferShapeContextFaker()
119 .NodeIoNum(12, 3)120 .NodeIoNum(12, 3)
120 .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})121 .IrInstanceNum({1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
@@ -159,11 +160,15 @@ TEST_F(LambApplyOptimizerAssignProtoTest, moment_shape_decides_output_shape)
159 ASSERT_EQ(Ops::Base::ToString(o2), Ops::Base::ToString(expect));160 ASSERT_EQ(Ops::Base::ToString(o2), Ops::Base::ToString(expect));
160}161}
161 162 
162-// grad 不参与广播：小于动量形状同样拒收（底层广播模板不支持对 In0 广播）。163+// grad 与其余输入一样参与广播：小于动量形状时按右对齐广播进动量形状，须放行。
163-TEST_F(LambApplyOptimizerAssignProtoTest, grad_smaller_than_moment_is_rejected)164+TEST_F(LambApplyOptimizerAssignProtoTest, grad_broadcast_into_moment_is_accepted)
164{165{
165 gert::Shape o0 = {}, o1 = {}, o2 = {};166 gert::Shape o0 = {}, o1 = {}, o2 = {};
166- ASSERT_EQ(RunInferShape({1, 1024}, {512, 1024}, {512, 1024}, {512, 1024}, &o0, &o1, &o2), ge::GRAPH_FAILED);167+ gert::Shape expect = {512, 1024};
168+ ASSERT_EQ(RunInferShape({1, 1024}, {512, 1024}, {512, 1024}, {512, 1024}, &o0, &o1, &o2), ge::GRAPH_SUCCESS);
169+ ASSERT_EQ(Ops::Base::ToString(o0), Ops::Base::ToString(expect));
170+ ASSERT_EQ(Ops::Base::ToString(o1), Ops::Base::ToString(expect));
171+ ASSERT_EQ(Ops::Base::ToString(o2), Ops::Base::ToString(expect));
167}172}
168 173 
169// input3 是唯一参与广播的输入：小于动量形状时按右对齐广播，须放行。174// input3 是唯一参与广播的输入：小于动量形状时按右对齐广播，须放行。
@@ -196,6 +201,23 @@ TEST_F(LambApplyOptimizerAssignProtoTest, inputv_inputm_mismatch_is_rejected)
196 ASSERT_EQ(RunInferShape({512, 1024}, {512, 1024}, {1, 1024}, {512, 1024}, &o0, &o1, &o2), ge::GRAPH_FAILED);201 ASSERT_EQ(RunInferShape({512, 1024}, {512, 1024}, {1, 1024}, {512, 1024}, &o0, &o1, &o2), ge::GRAPH_FAILED);
197}202}
198 203 
204+// 系数输入不再限定为标量: 与动量同形的系数张量须放行(对齐 A2 的可广播 ND Tensor 声明)。
205+TEST_F(LambApplyOptimizerAssignProtoTest, non_scalar_coefficient_is_accepted)
206+{
207+ gert::Shape o0 = {}, o1 = {}, o2 = {};
208+ gert::Shape expect = {512, 1024};
209+ ASSERT_EQ(RunInferShape({512, 1024}, {512, 1024}, {512, 1024}, {512, 1024}, &o0, &o1, &o2, {512, 1024}),
210+ ge::GRAPH_SUCCESS);
211+ ASSERT_EQ(Ops::Base::ToString(o0), Ops::Base::ToString(expect));
212+}
213+ 
214+// 系数输入大于动量形状: 广播结果无处容纳(动量是原地输出), 须拒收。
215+TEST_F(LambApplyOptimizerAssignProtoTest, coefficient_larger_than_moment_is_rejected)
216+{
217+ gert::Shape o0 = {}, o1 = {}, o2 = {};
218+ ASSERT_EQ(RunInferShape({1, 1024}, {1, 1024}, {1, 1024}, {1, 1024}, &o0, &o1, &o2, {512, 1024}), ge::GRAPH_FAILED);
219+}
220+ 
199TEST_F(LambApplyOptimizerAssignProtoTest, lambapplyoptimizerassign_infer_datatype)221TEST_F(LambApplyOptimizerAssignProtoTest, lambapplyoptimizerassign_infer_datatype)
200{222{
201 ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyOptimizerAssign"), nullptr);223 ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("LambApplyOptimizerAssign"), nullptr);
Moptim/lamb_apply_weight_assign/CMakeLists.txt+3-3
@@ -1,9 +1,9 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4# CANN Open Software License Agreement Version 2.0 (the "License").4# CANN Open Software License Agreement Version 2.0 (the "License").
5# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
6-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
@@ -13,4 +13,4 @@ set(SUPPORT_COMPUTE_UNIT "ascend950")
13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14set(SUPPORT_TILING_DIR "arch35")14set(SUPPORT_TILING_DIR "arch35")
15 15 
16-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_apply_weight_assign ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)16+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_apply_weight_assign ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DEPENDENCIES lamb_apply_common DISABLE_IN_OPP TRUE)
Moptim/lamb_apply_weight_assign/README.md+8-6
@@ -44,42 +44,42 @@
44 <tr>44 <tr>
45 <td>input0</td>45 <td>input0</td>
46 <td>输入</td>46 <td>输入</td>
47- <td>不支持空Tensor。公式中的input0（权重范数），标量。</td>47+ <td>支持空Tensor。公式中的input0（权重范数），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
48 <td>FLOAT16、FLOAT</td>48 <td>FLOAT16、FLOAT</td>
49 <td>ND</td>49 <td>ND</td>
50 </tr>50 </tr>
51 <tr>51 <tr>
52 <td>input1</td>52 <td>input1</td>
53 <td>输入</td>53 <td>输入</td>
54- <td>不支持空Tensor。公式中的input1（梯度范数），标量。</td>54+ <td>支持空Tensor。公式中的input1（梯度范数），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
55 <td>FLOAT16、FLOAT</td>55 <td>FLOAT16、FLOAT</td>
56 <td>ND</td>56 <td>ND</td>
57 </tr>57 </tr>
58 <tr>58 <tr>
59 <td>input2</td>59 <td>input2</td>
60 <td>输入</td>60 <td>输入</td>
61- <td>不支持空Tensor。公式中的input2（学习率），标量。</td>61+ <td>支持空Tensor。公式中的input2（学习率），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
62 <td>FLOAT16、FLOAT</td>62 <td>FLOAT16、FLOAT</td>
63 <td>ND</td>63 <td>ND</td>
64 </tr>64 </tr>
65 <tr>65 <tr>
66 <td>input3</td>66 <td>input3</td>
67 <td>输入</td>67 <td>输入</td>
68- <td>支持空Tensor。公式中的input3（update）。允许小于input_param并向上广播，但其shape必须能broadcast进input_param的shape。</td>68+ <td>支持空Tensor。公式中的input3（update），shape需与其他输入满足broadcast关系。</td>
69 <td>FLOAT16、FLOAT</td>69 <td>FLOAT16、FLOAT</td>
70 <td>ND</td>70 <td>ND</td>
71 </tr>71 </tr>
72 <tr>72 <tr>
73 <td>input_param</td>73 <td>input_param</td>
74 <td>输入</td>74 <td>输入</td>
75- <td>支持空Tensor。公式中的input_param（参数）。input_param为**原地(in-place)更新**输出，其shape必须等于input3与input_param广播后的完整输出shape（即input3须能broadcast进input_param）。</td>75+ <td>支持空Tensor。公式中的input_param（参数）。input_param为<b>原地(in-place)更新</b>输出，其shape必须等于全部输入的broadcast结果（即其余输入均须能broadcast进input_param）。</td>
76 <td>FLOAT16、FLOAT</td>76 <td>FLOAT16、FLOAT</td>
77 <td>ND</td>77 <td>ND</td>
78 </tr>78 </tr>
79 <tr>79 <tr>
80 <td>input_param</td>80 <td>input_param</td>
81 <td>输出</td>81 <td>输出</td>
82- <td>支持空Tensor。更新后的input_param（原地更新），shape取input3与input_param的broadcast结果。</td>82+ <td>支持空Tensor。更新后的input_param（原地更新），shape取全部输入的broadcast结果，等于input_param的shape。</td>
83 <td>FLOAT16、FLOAT</td>83 <td>FLOAT16、FLOAT</td>
84 <td>ND</td>84 <td>ND</td>
85 </tr>85 </tr>
@@ -87,6 +87,8 @@
87 87 
88## 约束说明88## 约束说明
89 89 
90+- 所有输入的shape需两两满足broadcast规则（右对齐，对应维相等或为1）。`input_param`是原地(in-place)更新的输出，故全部输入的broadcast结果必须恰好等于`input_param`的shape，否则原地写回会越界。
91+- 所有输入及输出的维度数不超过8。
90- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。92- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。
91 93 
92## 调用说明94## 调用说明
Moptim/lamb_apply_weight_assign/op_host/arch35/lamb_apply_weight_assign_tiling_arch35.cpp+4-24
@@ -14,7 +14,7 @@
14 */14 */
15 15 
16#include "lamb_apply_weight_assign_tiling_arch35.h"16#include "lamb_apply_weight_assign_tiling_arch35.h"
17-#include "../../../lamb_apply_common/lamb_apply_check_util.h"17+#include "../../../lamb_apply_common/op_host/arch35/lamb_apply_check_util.h"
18#include <graph/utils/type_utils.h>18#include <graph/utils/type_utils.h>
19#include <string>19#include <string>
20#include "infershape_broadcast_util.h"20#include "infershape_broadcast_util.h"
@@ -34,6 +34,7 @@ namespace optiling {
34constexpr static uint64_t LAMB_APPLY_WEIGHT_ASSIGN_TILING_PRIORITY = 0;34constexpr static uint64_t LAMB_APPLY_WEIGHT_ASSIGN_TILING_PRIORITY = 0;
35constexpr static int32_t INPUT_NUM = 5;35constexpr static int32_t INPUT_NUM = 5;
36constexpr static int32_t OUTPUT_NUM = 1;36constexpr static int32_t OUTPUT_NUM = 1;
37+constexpr static int32_t PARAM_IDX = 4; // input_param: ref(原地)输出
37static const char* const kInputNames[] = {"input0", "input1", "input2", "input3", "input_param"};38static const char* const kInputNames[] = {"input0", "input1", "input2", "input3", "input_param"};
38static const char* const kOutputNames[] = {"input_param"};39static const char* const kOutputNames[] = {"input_param"};
39 40 
@@ -51,39 +52,18 @@ static ge::graphStatus TilingPrepareForLambApplyWeightAssign(gert::TilingParseCo
51 52 
52ge::graphStatus LambApplyWeightAssignTiling::GetShapeAttrsInfo()53ge::graphStatus LambApplyWeightAssignTiling::GetShapeAttrsInfo()
53{54{
54- static const int32_t kScalarInputIdx[] = {0, 1, 2};
55 if (CheckLambApplyDtypeConsistency(context_, INPUT_NUM, kInputNames, OUTPUT_NUM, kOutputNames) !=55 if (CheckLambApplyDtypeConsistency(context_, INPUT_NUM, kInputNames, OUTPUT_NUM, kOutputNames) !=
56 ge::GRAPH_SUCCESS) {56 ge::GRAPH_SUCCESS) {
57 return ge::GRAPH_FAILED;57 return ge::GRAPH_FAILED;
58 }58 }
59- if (CheckLambApplyScalarNotEmpty(context_, kScalarInputIdx, sizeof(kScalarInputIdx) / sizeof(kScalarInputIdx[0]),
60- kInputNames) != ge::GRAPH_SUCCESS) {
61- return ge::GRAPH_FAILED;
62- }
63 return CheckInplaceShapeConstraint();59 return CheckInplaceShapeConstraint();
64}60}
65 61 
66// input_param 是 in-place 更新的参数输出(next_param 原地写回其输入 buffer,见 proto "(in-place)"),62// input_param 是 in-place 更新的参数输出(next_param 原地写回其输入 buffer,见 proto "(in-place)"),
67-// 内核按 broadcast(input3, input_param) 的完整网格计算并写回,故 input_param 形状必须 == 该网格。63+// 内核按全部输入广播的完整网格计算并写回,故 input_param 形状必须 == 该网格。
68-// 等价充要条件:input3 能广播进 input_param。
69ge::graphStatus LambApplyWeightAssignTiling::CheckInplaceShapeConstraint()64ge::graphStatus LambApplyWeightAssignTiling::CheckInplaceShapeConstraint()
70{65{
71- auto input3Shape = context_->GetInputShape(3);66+ return CheckLambApplyBroadcastIntoRef(context_, INPUT_NUM, PARAM_IDX, "input_param");
72- auto inputParamShape = context_->GetInputShape(4);
73- OP_CHECK_NULL_WITH_CONTEXT(context_, input3Shape);
74- OP_CHECK_NULL_WITH_CONTEXT(context_, inputParamShape);
75- const auto& i3s = input3Shape->GetStorageShape();
76- const auto& ips = inputParamShape->GetStorageShape();
77- // input3 能广播进 input_param <=> broadcast(input3, input_param) == input_param
78- gert::Shape bcShape;
79- if (!Ops::Base::BroadcastShape(&i3s, &ips, &bcShape) || !(bcShape == ips)) {
80- OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
81- context_->GetNodeName(), "input3", Ops::Base::ToString(i3s).c_str(),
82- "input3 must be broadcastable into the in-place param shape input_param (input_param is updated "
83- "in-place and must equal the broadcast output shape)");
84- return ge::GRAPH_FAILED;
85- }
86- return ge::GRAPH_SUCCESS;
87}67}
88 68 
89bool LambApplyWeightAssignTiling::IsCapable() { return true; }69bool LambApplyWeightAssignTiling::IsCapable() { return true; }
Moptim/lamb_apply_weight_assign/op_host/arch35/lamb_apply_weight_assign_tiling_arch35.h+1-1
@@ -44,7 +44,7 @@ protected:
44 44 
45private:45private:
46 // 校验 in-place 更新的 input_param 形状 == broadcast(input3, input_param)(input3 可向上广播)。46 // 校验 in-place 更新的 input_param 形状 == broadcast(input3, input_param)(input3 可向上广播)。
47- // dtype 一致性、标量非空的通用校验见 lamb_apply_common/lamb_apply_check_util.h。47+ // dtype 一致性、标量非空的通用校验见 lamb_apply_common/op_host/arch35/lamb_apply_check_util.h。
48 ge::graphStatus CheckInplaceShapeConstraint();48 ge::graphStatus CheckInplaceShapeConstraint();
49 uint64_t tilingKey = 0;49 uint64_t tilingKey = 0;
50};50};
Moptim/lamb_apply_weight_assign/op_host/lamb_apply_weight_assign_infershape.cpp+48-20文件内容审核中，请稍后刷新重试
Moptim/lamb_apply_weight_assign/op_kernel/arch35/lamb_apply_weight_assign_dag.h+26-5文件内容审核中，请稍后刷新重试
Moptim/lamb_apply_weight_assign/tests/assets/golden.py+65-19
@@ -32,24 +32,27 @@ def _scalars(*xs):
32 # 标量落成 **0 维 float64** 张量: torch 的类型提升里 0 维张量不会把 dim>0 的张量抬档，32 # 标量落成 **0 维 float64** 张量: torch 的类型提升里 0 维张量不会把 dim>0 的张量抬档，
33 # 所以数据是 fp32 时结果仍是 fp32(与改动前一致)，数据被 Promote 成 fp64 时标量自动33 # 所以数据是 fp32 时结果仍是 fp32(与改动前一致)，数据被 Promote 成 fp64 时标量自动
34 # 跟到 fp64，不会用一个先降到 fp32 的标量去污染高精度真值。34 # 跟到 fp64，不会用一个先降到 fp32 的标量去污染高精度真值。
35- return tuple(35+ # A2 语义: 这些"系数"输入是**可广播的 ND Tensor**(canndev ops/built-in/tbe/impl/lamb_*.py
36- torch.from_numpy(np.asarray(x, "float64").reshape(-1)[:1])[0] for x in xs36+ # 每步 mul/sub/div 都先 shape_util.broadcast_shapes 再 tbe.broadcast)。原先 reshape(-1)[:1]
37- )37+ # 只取首元素, 传多元素张量时静默按首元素计算 —— 与内核的广播实现不一致, 广播档必然假红。
38- 38+ # 改为返回完整张量交给 torch 自然广播: 形状 (1,) 的行为与原标量完全一致, 故常规档不变。
39- 39+ return tuple(_t(x) for x in xs)
40-_F32_TINY = torch.tensor(float(np.finfo(np.float32).tiny), dtype=torch.float32)
41 40 
42 41 
43def _div_ftz(a, b):42def _div_ftz(a, b):
44- """torch.div 的 fp32 除法 + FTZ，对齐算子两代共同的行为。43+ """fp32 除法，**不做 FTZ**。
45 44 
46- A2(910B) 的 Div 没有 config 参数（asc-devkit Div.md 里带 config 的原型对 Atlas A245+ 原先这里模拟了 FTZ（把落在 [1e-45, 1.1754944e-38) 的商冲刷成 0），依据是 arch35 的
47- 标注"不支持"），只有单指令一条路，Subnormal 必然 FTZ；arch35 的 Vec::Div 用默认46+ Vec::Div 走默认 DivConfig{DivAlgo::INTRINSIC}、该档 Subnormal 均被 FTZ。
48- DivConfig{DivAlgo::INTRINSIC}，文档写明该档"Subnormal 均被 FTZ"。CPU 默认不 FTZ，47+ 但实测该档把**非规格化操作数**也冲成 0：w_norm=1.93e-39 / g_norm=5.91e-39（商 0.327
49- 所以要显式补这一步，否则商落入 [1e-45, 1.1754944e-38) 时 golden 会给出算子实际48+ 是完全正常的值）时内核输出 3159/3159 全为 NaN（0/0）。这是缺陷不是语义，已把内核的
50- 不会产出的值。"""49+ Div 改成 PRECISION_0ULP_FTZ_FALSE（见 lamb_apply_weight_assign_dag.h 的 DivFtzFalse）。
51- q = torch.div(a, b)50+ 
52- return torch.where((q != 0) & (q.abs() < _F32_TINY), torch.zeros_like(q), q)51+ 精度判据以竞品为准：实测 torch 2.11 在 CPU 与 A100 上，fp32 非规格化数在存取/乘/加/减/除
52+ 各环节全程保留（CUDA 默认 -ftz=false，A100 对 fp32 subnormal 是全速硬件支持）。
53+ 故 golden 也不冲刷。函数保留是为了留住这段结论，除法本身就是 torch.div。
54+ """
55+ return torch.div(a, b)
53 56 
54 57 
55def _t(x):58def _t(x):
@@ -88,8 +91,14 @@ def lamb_apply_weight_assign_golden(
88 wn, gn, lr = _scalars(input0, input1, input2)91 wn, gn, lr = _scalars(input0, input1, input2)
89 upd, param = _t(input3), _t(input_param)92 upd, param = _t(input3), _t(input_param)
90 one = torch.tensor(1.0, dtype=torch.float32)93 one = torch.tensor(1.0, dtype=torch.float32)
91- inner = _div_ftz(wn, gn) if gn > 0 else one94+ # 逐元素选择: 输入改为可广播张量后, `if gn > 0` 这类标量控制流会对整个张量只判一次,
92- ratio = inner if wn > 0 else one95+ # 与内核的逐元素语义(A2 亦然)不符, 广播档必然假红。用 torch.where 逐元素选。
96+ # 注意 where 的两个分支都会被求值, _div_ftz 对 gn==0 必须自身安全(不产生 inf/nan 污染)。
97+ safe_gn = torch.where(gn > 0, gn, torch.ones_like(gn))
98+ inner = torch.where(
99+ gn > 0, _div_ftz(wn, safe_gn), one.expand_as(safe_gn) if safe_gn.dim() else one
100+ )
101+ ratio = torch.where(wn > 0, inner, torch.ones_like(inner))
93 # 运算序列对齐实现：A2 (tbe.vmul(update, lr) 再 vmul(ratio, ·)) 与 arch35 DAG102 # 运算序列对齐实现：A2 (tbe.vmul(update, lr) 再 vmul(ratio, ·)) 与 arch35 DAG
94 # (UpdLr = Update*Lr; RatioUpdLr = Ratio*UpdLr) 都是先算 update*lr。按 README 原先的103 # (UpdLr = Update*Lr; RatioUpdLr = Ratio*UpdLr) 都是先算 update*lr。按 README 原先的
95 # 字面顺序写成 (lr*ratio)*upd 在浮点下不等价：update 与 lr 同时较大时实现会先溢出成 inf，104 # 字面顺序写成 (lr*ratio)*upd 在浮点下不等价：update 与 lr 同时较大时实现会先溢出成 inf，
@@ -137,8 +146,32 @@ def _tp_t(x):
137 return torch.from_numpy(a)146 return torch.from_numpy(a)
138 147 
139 148 
149+def _tp_widen(t):
150+ """按 NPU 的加宽行为把三方入参落到**内核计算类型 U = float32**。
151+ 
152+ 规范依据 ttk_golden_logic.md §四/§五「浮点 + 三方」一格: 三方腿"按 NPU 加宽行为同步 cast"。
153+ 内核对 fp16 输入 unpack 成 fp32 计算(dag.h 的 `Cast<U, T, 0>`, U = float), 对 fp32 输入
154+ 原生 fp32, 两种情况计算类型都是 fp32, 故此处一律落到 fp32。
155+ 
156+ **不能只处理 fp16**: cross_check 下 TTK 会把 fp32 Promote 成 fp64 下发, 原样透传会让三方腿
157+ 在 fp64 上算 —— 内核 fp32 溢出/下溢的地方它都不溢出, 两条腿不在同一精度上, 比值无意义。
158+ 整型不动: 内核对 int 是原生/int32 累加, 走 fp32 会抹掉 >2^24 的低位。
159+ """
160+ return t.to(torch.float32) if t.is_floating_point() else t
161+ 
162+ 
163+def _tp_narrow(outs, dt):
164+ """出口复刻内核的 `Cast<T, U, 1>`: 窄回算子**声明**的 dtype T(不是 Promote 后的 dtype)。
165+ 
166+ 少了这一步或窄错目标, 三方腿会与走 Promote 的 golden 逐位相等 —— 双标杆塌成单标杆,
167+ 三比值分母被 safe_div 的 small_value 夹底, mare/rmse 恒为 1.0, 阈值永不触发。
168+ """
169+ return [o.to(dt) if o.is_floating_point() else o for o in outs]
170+ 
171+ 
140def _tp_s(x):172def _tp_s(x):
141- return _tp_t(x).reshape(-1)[0]173+ # 同 _scalars: 三方腿也必须广播, 不能只取首元素
174+ return _tp_t(x)
142 175 
143 176 
144class _LambApplyWeightAssignCompose:177class _LambApplyWeightAssignCompose:
@@ -146,8 +179,10 @@ class _LambApplyWeightAssignCompose:
146 wn, gn, lr = (_tp_s(v_) for v_ in (input0, input1, input2))179 wn, gn, lr = (_tp_s(v_) for v_ in (input0, input1, input2))
147 upd, param = _tp_t(input3), _tp_t(input_param)180 upd, param = _tp_t(input3), _tp_t(input_param)
148 one = torch.ones((), dtype=torch.float32, device=upd.device)181 one = torch.ones((), dtype=torch.float32, device=upd.device)
149- inner = torch.div(wn, gn) if bool(gn > 0) else one182+ # 逐元素选择(理由同 golden 主体): 三方腿也必须按元素选, 不能整张量只判一次
150- ratio = inner if bool(wn > 0) else one183+ safe_gn = torch.where(gn > 0, gn, torch.ones_like(gn))
184+ inner = torch.where(gn > 0, torch.div(wn, safe_gn), one)
185+ ratio = torch.where(wn > 0, inner, one)
151 return [param - ratio * (upd * lr)]186 return [param - ratio * (upd * lr)]
152 187 
153 188 
@@ -165,6 +200,16 @@ class _LambApplyWeightAssignCompose:
165 200 
166 201 
167def _kf_widen(a, seen):202def _kf_widen(a, seen):
203+ """按 NPU 的加宽行为把三方入参加宽, 并记下算子声明的 dtype T 供出口窄回。
204+ 
205+ 三方腿拿到的是**原始 dtype T**(TTK 的 Promote 只作用于 golden, 见 profiling.py 的
206+ golden_mode_override; 三方腿走 _xpu_inputs -> original_input_arrays), 故这里按 T 判断:
207+ - T = fp16/bf16: 内核 unpack 成 fp32 全程不落回(dag.h 的 Cast<U, T, 0>, U = float),
208+ 而 torch 只在单个算子内部用 opmath=float、算子之间每步落回 T。不加宽就等于拿
209+ "逐步截断的实现"当竞品, 与被测内核不是同一个算法 -> 加宽到 fp32, 出口窄回 T。
210+ - T = fp32: 内核计算类型 U 即 fp32, **不加宽**; torch 同样原生 fp32 -> 原样不动。
211+ - 整型: 内核原生/int32 累加, 走 fp32 会抹掉 >2^24 的低位 -> 不动。
212+ """
168 if isinstance(a, (list, tuple)):213 if isinstance(a, (list, tuple)):
169 return type(a)(_kf_widen(x, seen) for x in a)214 return type(a)(_kf_widen(x, seen) for x in a)
170 if isinstance(a, torch.Tensor):215 if isinstance(a, torch.Tensor):
@@ -187,6 +232,7 @@ def _kf_widen(a, seen):
187 232 
188 233 
189def _kf_narrow(o, dt):234def _kf_narrow(o, dt):
235+ """出口复刻内核的 `Cast<T, U, 1>`: 窄回算子声明的 dtype T。T = fp32 时无需窄回。"""
190 if isinstance(o, (list, tuple)):236 if isinstance(o, (list, tuple)):
191 return type(o)(_kf_narrow(x, dt) for x in o)237 return type(o)(_kf_narrow(x, dt) for x in o)
192 if isinstance(o, torch.Tensor) and o.is_floating_point():238 if isinstance(o, torch.Tensor) and o.is_floating_point():
Moptim/lamb_next_m_v/CMakeLists.txt+3-3
@@ -1,9 +1,9 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4# CANN Open Software License Agreement Version 2.0 (the "License").4# CANN Open Software License Agreement Version 2.0 (the "License").
5# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
6-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
@@ -13,4 +13,4 @@ set(SUPPORT_COMPUTE_UNIT "ascend950")
13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14set(SUPPORT_TILING_DIR "arch35")14set(SUPPORT_TILING_DIR "arch35")
15 15 
16-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_next_m_v ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)16+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_next_m_v ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DEPENDENCIES lamb_apply_common DISABLE_IN_OPP TRUE)
Moptim/lamb_next_m_v/README.md+15-13
@@ -48,7 +48,7 @@
48 <tr>48 <tr>
49 <td>input_mul3</td>49 <td>input_mul3</td>
50 <td>输入</td>50 <td>输入</td>
51- <td>支持空Tensor。公式中的input_mul3（g^2），主张量，shape需与input_mul0满足broadcast关系。</td>51+ <td>支持空Tensor。公式中的input_mul3（g^2），主张量，shape需与其他输入满足broadcast关系。</td>
52 <td>FLOAT16、FLOAT</td>52 <td>FLOAT16、FLOAT</td>
53 <td>ND</td>53 <td>ND</td>
54 </tr>54 </tr>
@@ -76,7 +76,7 @@
76 <tr>76 <tr>
77 <td>input_mul0</td>77 <td>input_mul0</td>
78 <td>输入</td>78 <td>输入</td>
79- <td>支持空Tensor。公式中的input_mul0（一阶矩m），主张量，shape需与input_mul3满足broadcast关系，其broadcast结果决定各输出的shape。</td>79+ <td>支持空Tensor。公式中的input_mul0（一阶矩m），主张量，shape需与其他输入满足broadcast关系，全部输入的broadcast结果决定各输出的shape。</td>
80 <td>FLOAT16、FLOAT</td>80 <td>FLOAT16、FLOAT</td>
81 <td>ND</td>81 <td>ND</td>
82 </tr>82 </tr>
@@ -97,70 +97,70 @@
97 <tr>97 <tr>
98 <td>mul0_x</td>98 <td>mul0_x</td>
99 <td>输入</td>99 <td>输入</td>
100- <td>不支持空Tensor。公式中的mul0_x（beta1），标量。</td>100+ <td>支持空Tensor。公式中的mul0_x（beta1），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
101 <td>FLOAT16、FLOAT</td>101 <td>FLOAT16、FLOAT</td>
102 <td>ND</td>102 <td>ND</td>
103 </tr>103 </tr>
104 <tr>104 <tr>
105 <td>mul1_sub</td>105 <td>mul1_sub</td>
106 <td>输入</td>106 <td>输入</td>
107- <td>不支持空Tensor。公式中的mul1_sub（1-beta1），标量。</td>107+ <td>支持空Tensor。公式中的mul1_sub（1-beta1），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
108 <td>FLOAT16、FLOAT</td>108 <td>FLOAT16、FLOAT</td>
109 <td>ND</td>109 <td>ND</td>
110 </tr>110 </tr>
111 <tr>111 <tr>
112 <td>mul2_x</td>112 <td>mul2_x</td>
113 <td>输入</td>113 <td>输入</td>
114- <td>不支持空Tensor。公式中的mul2_x（beta2），标量。</td>114+ <td>支持空Tensor。公式中的mul2_x（beta2），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
115 <td>FLOAT16、FLOAT</td>115 <td>FLOAT16、FLOAT</td>
116 <td>ND</td>116 <td>ND</td>
117 </tr>117 </tr>
118 <tr>118 <tr>
119 <td>mul3_sub1</td>119 <td>mul3_sub1</td>
120 <td>输入</td>120 <td>输入</td>
121- <td>不支持空Tensor。公式中的mul3_sub1（1-beta2），标量。</td>121+ <td>支持空Tensor。公式中的mul3_sub1（1-beta2），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
122 <td>FLOAT16、FLOAT</td>122 <td>FLOAT16、FLOAT</td>
123 <td>ND</td>123 <td>ND</td>
124 </tr>124 </tr>
125 <tr>125 <tr>
126 <td>mul4_x</td>126 <td>mul4_x</td>
127 <td>输入</td>127 <td>输入</td>
128- <td>不支持空Tensor。公式中的mul4_x（权重衰减系数），标量。</td>128+ <td>支持空Tensor。公式中的mul4_x（权重衰减系数），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
129 <td>FLOAT16、FLOAT</td>129 <td>FLOAT16、FLOAT</td>
130 <td>ND</td>130 <td>ND</td>
131 </tr>131 </tr>
132 <tr>132 <tr>
133 <td>add2_y</td>133 <td>add2_y</td>
134 <td>输入</td>134 <td>输入</td>
135- <td>不支持空Tensor。公式中的add2_y（epsilon），标量。</td>135+ <td>支持空Tensor。公式中的add2_y（epsilon），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
136 <td>FLOAT16、FLOAT</td>136 <td>FLOAT16、FLOAT</td>
137 <td>ND</td>137 <td>ND</td>
138 </tr>138 </tr>
139 <tr>139 <tr>
140 <td>y1</td>140 <td>y1</td>
141 <td>输出</td>141 <td>输出</td>
142- <td>支持空Tensor。公式中的y1（update），shape取input_mul3与input_mul0的broadcast结果。</td>142+ <td>支持空Tensor。公式中的y1（update），shape取全部输入的broadcast结果。</td>
143 <td>FLOAT16、FLOAT</td>143 <td>FLOAT16、FLOAT</td>
144 <td>ND</td>144 <td>ND</td>
145 </tr>145 </tr>
146 <tr>146 <tr>
147 <td>y2</td>147 <td>y2</td>
148 <td>输出</td>148 <td>输出</td>
149- <td>支持空Tensor。公式中的y2（next_m），shape取input_mul3与input_mul0的broadcast结果。</td>149+ <td>支持空Tensor。公式中的y2（next_m），shape取全部输入的broadcast结果。</td>
150 <td>FLOAT16、FLOAT</td>150 <td>FLOAT16、FLOAT</td>
151 <td>ND</td>151 <td>ND</td>
152 </tr>152 </tr>
153 <tr>153 <tr>
154 <td>y3</td>154 <td>y3</td>
155 <td>输出</td>155 <td>输出</td>
156- <td>支持空Tensor。公式中的y3（next_v），shape取input_mul3与input_mul0的broadcast结果。</td>156+ <td>支持空Tensor。公式中的y3（next_v），shape取全部输入的broadcast结果。</td>
157 <td>FLOAT16、FLOAT</td>157 <td>FLOAT16、FLOAT</td>
158 <td>ND</td>158 <td>ND</td>
159 </tr>159 </tr>
160 <tr>160 <tr>
161 <td>y4</td>161 <td>y4</td>
162 <td>输出</td>162 <td>输出</td>
163- <td>支持空Tensor。公式中的y4，shape取input_mul3与input_mul0的broadcast结果。</td>163+ <td>支持空Tensor。公式中的y4，shape取全部输入的broadcast结果。</td>
164 <td>FLOAT16、FLOAT</td>164 <td>FLOAT16、FLOAT</td>
165 <td>ND</td>165 <td>ND</td>
166 </tr>166 </tr>
@@ -168,8 +168,10 @@
168 168 
169## 约束说明169## 约束说明
170 170 
171+- 所有输入的shape需两两满足broadcast规则（右对齐，对应维相等或为1），输出shape为全部输入的broadcast结果。
172+- 所有输入及输出的维度数不超过8。
171- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。173- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。
172-- input_mul0/input_mul1/input_mul2/input_mul3/input_mul4 为主张量，其shape需保持一致（或可相互广播到同一shape）；各输出y1/y2/y3/y4的shape均取该广播结果（实现以input_mul3与input_mul0的broadcast结果为准）。174+- 各输出y1/y2/y3/y4的shape均取全部输入的broadcast结果。
173 175 
174## 调用说明176## 调用说明
175 177 
Moptim/lamb_next_m_v/op_host/arch35/lamb_next_m_v_tiling_arch35.cpp+48-51
@@ -15,20 +15,24 @@
15 15 
16#include "lamb_next_m_v_tiling_arch35.h"16#include "lamb_next_m_v_tiling_arch35.h"
17#include <graph/utils/type_utils.h>17#include <graph/utils/type_utils.h>
18-#include "../../op_kernel/arch35/lamb_next_m_v_dag.h"18+#include <securec.h>
19-#include "atvoss/broadcast/broadcast_tiling.h"19+#include <algorithm>
20+#include <vector>
20#include "log/log.h"21#include "log/log.h"
21#include "platform/platform_info.h"22#include "platform/platform_info.h"
22#include "register/op_impl_registry.h"23#include "register/op_impl_registry.h"
23#include "register/tilingdata_base.h"24#include "register/tilingdata_base.h"
24#include "op_host/tiling_templates_registry.h"25#include "op_host/tiling_templates_registry.h"
25 26 
26-using namespace AscendC;
27using namespace ge;27using namespace ge;
28 28 
29namespace optiling {29namespace optiling {
30 30 
31constexpr static uint64_t LAMB_NEXT_M_V_TILING_PRIORITY = 0;31constexpr static uint64_t LAMB_NEXT_M_V_TILING_PRIORITY = 0;
32+constexpr static uint64_t LNMV_TILING_KEY_FP32 = 100;
33+constexpr static uint64_t LNMV_TILING_KEY_FP16 = 200;
34+constexpr static uint64_t TILING_KEY_FP32 = 100;
35+constexpr static uint64_t TILING_KEY_FP16 = 200;
32constexpr static int32_t INPUT_NUM = 13;36constexpr static int32_t INPUT_NUM = 13;
33constexpr static int32_t OUTPUT_NUM = 4;37constexpr static int32_t OUTPUT_NUM = 4;
34 38 
@@ -80,18 +84,6 @@ ge::graphStatus LambNextMVTiling::GetShapeAttrsInfo()
80 return ge::GRAPH_FAILED;84 return ge::GRAPH_FAILED;
81 }85 }
82 }86 }
83- // 标量类输入为每元素计算所必需的系数, 空Tensor 视为缺失必选值(畸形输入), 不支持。
84- static const int32_t kScalarInputIdx[] = {7, 8, 9, 10, 11, 12};
85- for (int32_t scalarIdx : kScalarInputIdx) {
86- auto scalarShape = context_->GetInputShape(scalarIdx);
87- OP_CHECK_NULL_WITH_CONTEXT(context_, scalarShape);
88- if (scalarShape->GetStorageShape().GetShapeSize() == 0) {
89- OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), kInputNames[scalarIdx],
90- Ops::Base::ToString(scalarShape->GetStorageShape()).c_str(),
91- "scalar input does not support empty tensor");
92- return ge::GRAPH_FAILED;
93- }
94- }
95 return ge::GRAPH_SUCCESS;87 return ge::GRAPH_SUCCESS;
96}88}
97 89 
@@ -99,49 +91,44 @@ bool LambNextMVTiling::IsCapable() { return true; }
99 91 
100ge::graphStatus LambNextMVTiling::DoOpTiling()92ge::graphStatus LambNextMVTiling::DoOpTiling()
101{93{
102- // 空 tensor 应对(空进空出): 输出为空(0元素)时设 1 核(空转), 配合全0 tiling 数据(blockFormer=0)使 kernel 空转退出,94+ auto rawTilingData = context_->GetRawTilingData();
103- // 直接成功。95+ OP_CHECK_NULL_WITH_CONTEXT(context_, rawTilingData);
104- auto emptyTensorOutShape0 = context_->GetOutputShape(0);
105- if (emptyTensorOutShape0 != nullptr && emptyTensorOutShape0->GetStorageShape().GetShapeSize() == 0) {
106- auto emptyRawTiling = context_->GetRawTilingData();
107- if (emptyRawTiling != nullptr && emptyRawTiling->GetData() != nullptr) {
108- size_t emptyCap = emptyRawTiling->GetCapacity();
109- uint8_t* emptyPtr = static_cast<uint8_t*>(emptyRawTiling->GetData());
110- for (size_t emptyIdx = 0; emptyIdx < emptyCap; ++emptyIdx) {
111- emptyPtr[emptyIdx] = 0;
112- }
113- emptyRawTiling->SetDataSize(emptyCap);
114- }
115- size_t* emptyWs = context_->GetWorkspaceSizes(1);
116- if (emptyWs != nullptr) {
117- emptyWs[0] = 0;
118- }
119- context_->SetBlockDim(1);
120- tilingKey = GET_TPL_TILING_KEY(1); // schMode=1(已编译), 配合全0 tiling(blockFormer=0)空转
121- return ge::GRAPH_SUCCESS;
122- }
123 auto input0Desc = context_->GetInputDesc(0);96 auto input0Desc = context_->GetInputDesc(0);
124 OP_CHECK_NULL_WITH_CONTEXT(context_, input0Desc);97 OP_CHECK_NULL_WITH_CONTEXT(context_, input0Desc);
98+ 
125 ge::DataType input0DType = input0Desc->GetDataType();99 ge::DataType input0DType = input0Desc->GetDataType();
126- if (input0DType == ge::DT_FLOAT16) {100+ uint32_t dtSize = 0;
127- BroadcastBaseTiling<LambNextMVOp::LambNextMVCompute<half, float>::OpDag> brcBaseTiling(101+ if (input0DType == ge::DT_FLOAT) {
128- context_, static_cast<uint32_t>(BROADCAST_KERNEL_TYPE::KERNEL_TYPE_NDDMA));102+ tilingKey = TILING_KEY_FP32;
129- OP_CHECK_IF(brcBaseTiling.DoTiling() == ge::GRAPH_FAILED,103+ dtSize = sizeof(float);
130- OP_LOGE(context_->GetNodeName(), "Do tiling failed. Please check the detailed log."),104+ } else if (input0DType == ge::DT_FLOAT16) {
131- return ge::GRAPH_FAILED);105+ tilingKey = TILING_KEY_FP16;
132- tilingKey = GET_TPL_TILING_KEY(brcBaseTiling.GetSchMode());106+ dtSize = sizeof(uint16_t);
133- } else if (input0DType == ge::DT_FLOAT) {
134- BroadcastBaseTiling<LambNextMVOp::LambNextMVCompute<float, float>::OpDag> brcBaseTiling(
135- context_, static_cast<uint32_t>(BROADCAST_KERNEL_TYPE::KERNEL_TYPE_NDDMA));
136- OP_CHECK_IF(brcBaseTiling.DoTiling() == ge::GRAPH_FAILED,
137- OP_LOGE(context_->GetNodeName(), "Do tiling failed. Please check the detailed log."),
138- return ge::GRAPH_FAILED);
139- tilingKey = GET_TPL_TILING_KEY(brcBaseTiling.GetSchMode());
140 } else {107 } else {
141 OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "input_mul3", Ops::Base::ToString(input0DType).c_str(),108 OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "input_mul3", Ops::Base::ToString(input0DType).c_str(),
142 "fp16 or fp32");109 "fp16 or fp32");
143 return ge::GRAPH_FAILED;110 return ge::GRAPH_FAILED;
144 }111 }
112+ 
113+ using PlanTiling = LambBrcTilingData<13, 4>;
114+ td_ = PlanTiling{};
115+ // 空进空出: 输出 0 元素时 tiling 全 0, kernel 按 usedCoreNum=0 直接退出。
116+ auto outShape0 = context_->GetOutputShape(0);
117+ if (outShape0 == nullptr || outShape0->GetStorageShape().GetShapeSize() != 0) {
118+ // 先取返回值再判: 模板实参里的逗号会被预处理器当成 OP_CHECK_IF 的参数分隔符。
119+ ge::graphStatus planRet = BuildLambBrcPlan<13, 4>(context_, coreNum_, ubSize_, dtSize, td_);
120+ OP_CHECK_IF(planRet != ge::GRAPH_SUCCESS, OP_LOGE(context_->GetNodeName(), "build broadcast plan failed"),
121+ return ge::GRAPH_FAILED);
122+ }
123+ 
124+ auto ret = memcpy_s(rawTilingData->GetData(), rawTilingData->GetCapacity(), &td_, sizeof(td_));
125+ OP_CHECK_IF(ret != EOK, OP_LOGE(context_->GetNodeName(), "copy tiling data failed, ret %d", ret),
126+ return ge::GRAPH_FAILED);
127+ rawTilingData->SetDataSize(sizeof(td_));
128+ context_->SetBlockDim(std::max<uint32_t>(td_.usedCoreNum, 1));
129+ size_t* ws = context_->GetWorkspaceSizes(1);
130+ OP_CHECK_NULL_WITH_CONTEXT(context_, ws);
131+ ws[0] = 0U;
145 return ge::GRAPH_SUCCESS;132 return ge::GRAPH_SUCCESS;
146}133}
147 134 
@@ -153,7 +140,17 @@ ge::graphStatus LambNextMVTiling::GetWorkspaceSize() { return ge::GRAPH_SUCCESS;
153 140 
154ge::graphStatus LambNextMVTiling::PostTiling() { return ge::GRAPH_SUCCESS; }141ge::graphStatus LambNextMVTiling::PostTiling() { return ge::GRAPH_SUCCESS; }
155 142 
156-ge::graphStatus LambNextMVTiling::GetPlatformInfo() { return ge::GRAPH_SUCCESS; }143+ge::graphStatus LambNextMVTiling::GetPlatformInfo()
144+{
145+ auto compileInfo = static_cast<const LambNextMVCompileInfo*>(context_->GetCompileInfo());
146+ OP_CHECK_NULL_WITH_CONTEXT(context_, compileInfo);
147+ coreNum_ = compileInfo->coreNum;
148+ ubSize_ = compileInfo->ubSize;
149+ OP_CHECK_IF(coreNum_ == 0 || ubSize_ == 0,
150+ OP_LOGE(context_->GetNodeName(), "invalid platform info: coreNum %lu ubSize %lu", coreNum_, ubSize_),
151+ return ge::GRAPH_FAILED);
152+ return ge::GRAPH_SUCCESS;
153+}
157 154 
158static ge::graphStatus TilingForLambNextMV(gert::TilingContext* context)155static ge::graphStatus TilingForLambNextMV(gert::TilingContext* context)
159{156{
Moptim/lamb_next_m_v/op_host/arch35/lamb_next_m_v_tiling_arch35.h+6-2文件内容审核中，请稍后刷新重试
Moptim/lamb_next_m_v/op_host/lamb_next_m_v_infershape.cpp+34-12
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include <vector>
16#include "register/op_impl_registry.h"17#include "register/op_impl_registry.h"
17#include "log/log.h"18#include "log/log.h"
18#include "infershape_broadcast_util.h"19#include "infershape_broadcast_util.h"
@@ -20,28 +21,49 @@
20using namespace Ops::Base;21using namespace Ops::Base;
21using namespace ge;22using namespace ge;
22namespace ops {23namespace ops {
23-// full tensors: input_mul3(0,g^2), input_mul0(4,m). All four outputs share their broadcast shape.24+// A2 语义: 本族算子的所有输入都是可广播的 ND Tensor(见 canndev
24-constexpr size_t IN_MUL3 = 0;25+// ops/built-in/tbe/impl/lamb_*.py, 每一步 mul/sub/div 都先 shape_util.broadcast_shapes
25-constexpr size_t IN_MUL0 = 4;26+// 再 tbe.broadcast), 输出形状为全部输入广播的结果。A2 的 op_proto 只声明了其中两个输入,
27+// 属声明宽松, 不作为支持面依据。
28+constexpr size_t IN_NUM = 13;
26constexpr size_t OUT_NUM = 4;29constexpr size_t OUT_NUM = 4;
27 30 
28static ge::graphStatus InferShape4LambNextMV(gert::InferShapeContext* context)31static ge::graphStatus InferShape4LambNextMV(gert::InferShapeContext* context)
29{32{
30- auto g2 = context->GetInputShape(IN_MUL3);33+ std::vector<const gert::Shape*> inShapes;
31- OP_CHECK_NULL_WITH_CONTEXT(context, g2);34+ inShapes.reserve(IN_NUM);
32- auto m = context->GetInputShape(IN_MUL0);35+ for (size_t i = 0; i < IN_NUM; i++) {
33- OP_CHECK_NULL_WITH_CONTEXT(context, m);36+ auto in = context->GetInputShape(i);
37+ OP_CHECK_NULL_WITH_CONTEXT(context, in);
38+ inShapes.push_back(in);
39+ }
40+ gert::Shape bcShape;
41+ // 逐对折叠广播: 只用两参数重载。vector 重载虽在 op_common/op_host/infershape_broadcast_util.h
42+ // 以 Ops::Base 声明, 但 libops_base.so 只导出两参数版, vector 版的实现在 libop_common.so 的
43+ // 小写 ops 命名空间下 —— 声明与实现命名空间不一致, 用它会编译期通过、加载期
44+ // undefined symbol 而装不上包。
45+ bcShape = *inShapes[0];
46+ for (size_t i = 1; i < inShapes.size(); i++) {
47+ gert::Shape tmp;
48+ OP_CHECK_IF(!BroadcastShape(&bcShape, inShapes[i], &tmp),
49+ OP_LOGE(context->GetNodeName(), "input shapes cannot broadcast together"), return ge::GRAPH_FAILED);
50+ bcShape = tmp;
51+ }
52+ // 标量归一: 全标量输入 broadcast 得 0 维空 shape (), 与 A2 的 shape_util.scalar2tensor_one
53+ // 对齐, 归一为 (1,)。否则动态 shape 编译期 DFX 生成会对空 shape 做 reduce 连乘(无初值)而报
54+ // TypeError 编译失败。
55+ if (bcShape.GetDimNum() == 0) {
56+ bcShape.SetDimNum(1);
57+ bcShape.SetDim(0, 1);
58+ }
34 for (size_t i = 0; i < OUT_NUM; i++) {59 for (size_t i = 0; i < OUT_NUM; i++) {
35 auto out = context->GetOutputShape(i);60 auto out = context->GetOutputShape(i);
36 OP_CHECK_NULL_WITH_CONTEXT(context, out);61 OP_CHECK_NULL_WITH_CONTEXT(context, out);
37- OP_CHECK_IF(!BroadcastShape(g2, m, out),62+ *out = bcShape;
38- OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "input_mul3 and input_mul0",
39- (ToString(*g2) + " and " + ToString(*m)).c_str(),
40- "shapes cannot broadcast"),
41- return ge::GRAPH_FAILED);
42 }63 }
43 return GRAPH_SUCCESS;64 return GRAPH_SUCCESS;
44}65}
66+ 
45static ge::graphStatus InferDataType4LambNextMV(gert::InferDataTypeContext* context)67static ge::graphStatus InferDataType4LambNextMV(gert::InferDataTypeContext* context)
46{68{
47 if (context == nullptr) {69 if (context == nullptr) {
Aoptim/lamb_next_m_v/op_host/lamb_next_m_v_tiling_def.h+51-0
@@ -0,0 +1,51 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file lamb_next_m_v_tiling_def.h
13+ * \brief LambNextMV tiling data definition (arch35 手写内核)
14+ *
15+ * 只用于给 framework 定 raw tiling buffer 的容量; 实际下发由 host 侧
16+ * memcpy 整个 LambBrcTilingData<13, 4> POD 完成, 故字段与 POD 严格一一对应。
17+ */
18+ 
19+#ifndef LAMB_NEXT_M_V_TILING_DEF_H
20+#define LAMB_NEXT_M_V_TILING_DEF_H
21+ 
22+#include "register/tilingdata_base.h"
23+#include "register/op_impl_registry.h"
24+ 
25+namespace optiling {
26+constexpr int32_t LAMB_NEXT_M_V_IN_NUM = 13;
27+constexpr int32_t LAMB_NEXT_M_V_MAX_DIM = 8;
28+constexpr int32_t LAMB_NEXT_M_V_STRIDE_NUM = LAMB_NEXT_M_V_IN_NUM * LAMB_NEXT_M_V_MAX_DIM;
29+ 
30+BEGIN_TILING_DATA_DEF(LambNextMVTilingData)
31+TILING_DATA_FIELD_DEF(uint64_t, totalNum);
32+TILING_DATA_FIELD_DEF(uint64_t, totalRows);
33+TILING_DATA_FIELD_DEF_ARR(uint64_t, LAMB_NEXT_M_V_STRIDE_NUM, effStride);
34+TILING_DATA_FIELD_DEF_ARR(uint64_t, LAMB_NEXT_M_V_IN_NUM, srcBlockLen);
35+TILING_DATA_FIELD_DEF(uint64_t, perCoreElems);
36+TILING_DATA_FIELD_DEF(uint64_t, rowsPerCore);
37+TILING_DATA_FIELD_DEF(uint32_t, tileLen);
38+TILING_DATA_FIELD_DEF(uint32_t, blockLen);
39+TILING_DATA_FIELD_DEF(uint32_t, rowsPerTile);
40+TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum);
41+TILING_DATA_FIELD_DEF(uint32_t, splitAxis);
42+TILING_DATA_FIELD_DEF(uint32_t, collapsedRank);
43+TILING_DATA_FIELD_DEF(uint32_t, tilingKey);
44+TILING_DATA_FIELD_DEF_ARR(uint32_t, LAMB_NEXT_M_V_IN_NUM, inKind);
45+TILING_DATA_FIELD_DEF_ARR(uint32_t, LAMB_NEXT_M_V_MAX_DIM, outShape);
46+TILING_DATA_FIELD_DEF_ARR(uint32_t, LAMB_NEXT_M_V_STRIDE_NUM, inShape);
47+END_TILING_DATA_DEF;
48+ 
49+REGISTER_TILING_DATA_CLASS(LambNextMV, LambNextMVTilingData)
50+} // namespace optiling
51+#endif // LAMB_NEXT_M_V_TILING_DEF_H
Aoptim/lamb_next_m_v/op_kernel/CMakeLists.txt+16-0文件内容审核中，请稍后刷新重试
Doptim/lamb_next_m_v/op_kernel/arch35/lamb_next_m_v_dag.h+0-104
@@ -1,104 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_next_m_v_dag.h
13- * \brief lamb_next_m_v_dag head file
14- */
15- 
16-#ifndef LAMB_NEXT_M_V_DAG_H
17-#define LAMB_NEXT_M_V_DAG_H
18-#include "atvoss/util/dag.h"
19-#include "atvoss/util/vec.h"
20-#include "atvoss/util/placeholder.h"
21- 
22-// The atvoss Placeholder ships In0..In11 (12 inputs). LambNextMV has 13 inputs; extend with In12
23-// following the framework's exact pattern (additive: the struct + its IsInHolder trait, same namespace).
24-namespace Ops {
25-namespace Base {
26-namespace Placeholder {
27-template <class T, class Attr = InAttr<>>
28-struct In12 : public Holder<T, Attr, 12> {};
29-template <class U, class Attr>
30-struct IsInHolder<In12<U, Attr>> {
31- constexpr static bool Value = true;
32-};
33-} // namespace Placeholder
34-} // namespace Base
35-} // namespace Ops
36- 
37-namespace LambNextMVOp {
38-using namespace AscendC;
39-using namespace Ops::Base;
40-// In : 0 input_mul3(g^2) 1 input_mul2(v) 2 input_realdiv1(1-b2^t) 3 input_mul1(g) 4 input_mul0(m)
41-// 5 input_realdiv0(1-b1^t) 6 input_mul4(param) 7 mul0_x(b1) 8 mul1_sub(1-b1) 9 mul2_x(b2)
42-// 10 mul3_sub1(1-b2) 11 mul4_x(wd) 12 add2_y(eps) [0,1,3,4,6 full; rest scalar]
43-// Out: 0 y1(update) 1 y2(next_m) 2 y3(next_v) 3 y4(m_unbiased/(sqrt(v_unbiased)+eps))
44-// Compute in U (=float): half casts to float (div/sqrt need float precision). rsqrt -> Div(_,Sqrt).
45-template <typename T, typename U>
46-struct LambNextMVCompute {
47- // full tensors
48- using InG2 = Bind<Vec::CopyInBrc<T>, Placeholder::In0<T>>;
49- using InV = Bind<Vec::CopyInBrc<T>, Placeholder::In1<T>>;
50- using InG = Bind<Vec::CopyInBrc<T>, Placeholder::In3<T>>;
51- using InM = Bind<Vec::CopyInBrc<T>, Placeholder::In4<T>>;
52- using InParam = Bind<Vec::CopyInBrc<T>, Placeholder::In6<T>>;
53- // scalar coefficients
54- using InRd1 = Bind<Vec::Duplicate<T>, Placeholder::In2<T, Placeholder::ScalarAttr<true>>>;
55- using InRd0 = Bind<Vec::Duplicate<T>, Placeholder::In5<T, Placeholder::ScalarAttr<true>>>;
56- using InB1 = Bind<Vec::Duplicate<T>, Placeholder::In7<T, Placeholder::ScalarAttr<true>>>;
57- using InOmB1 = Bind<Vec::Duplicate<T>, Placeholder::In8<T, Placeholder::ScalarAttr<true>>>;
58- using InB2 = Bind<Vec::Duplicate<T>, Placeholder::In9<T, Placeholder::ScalarAttr<true>>>;
59- using InOmB2 = Bind<Vec::Duplicate<T>, Placeholder::In10<T, Placeholder::ScalarAttr<true>>>;
60- using InWd = Bind<Vec::Duplicate<T>, Placeholder::In11<T, Placeholder::ScalarAttr<true>>>;
61- using InEps = Bind<Vec::Duplicate<T>, Placeholder::In12<T, Placeholder::ScalarAttr<true>>>;
62- 
63- using G2 = Bind<Vec::Cast<U, T, 0>, InG2>;
64- using V = Bind<Vec::Cast<U, T, 0>, InV>;
65- using G = Bind<Vec::Cast<U, T, 0>, InG>;
66- using M = Bind<Vec::Cast<U, T, 0>, InM>;
67- using Param = Bind<Vec::Cast<U, T, 0>, InParam>;
68- using Rd1 = Bind<Vec::Cast<U, T, 0>, InRd1>;
69- using Rd0 = Bind<Vec::Cast<U, T, 0>, InRd0>;
70- using B1 = Bind<Vec::Cast<U, T, 0>, InB1>;
71- using OmB1 = Bind<Vec::Cast<U, T, 0>, InOmB1>;
72- using B2 = Bind<Vec::Cast<U, T, 0>, InB2>;
73- using OmB2 = Bind<Vec::Cast<U, T, 0>, InOmB2>;
74- using Wd = Bind<Vec::Cast<U, T, 0>, InWd>;
75- using Eps = Bind<Vec::Cast<U, T, 0>, InEps>;
76- 
77- // next_v = v*b2 + g^2*(1-b2)
78- using NextV = Bind<Vec::Add<U>, Bind<Vec::Mul<U>, V, B2>, Bind<Vec::Mul<U>, G2, OmB2>>;
79- using VUnbias = Bind<Vec::Div<U>, NextV, Rd1>;
80- using SqrtVeps = Bind<Vec::Sqrt<U>, Bind<Vec::Add<U>, VUnbias, Eps>>; // sqrt(v_unbiased+eps)
81- using SqrtVAddEps = Bind<Vec::Add<U>, Bind<Vec::Sqrt<U>, VUnbias>, Eps>; // sqrt(v_unbiased)+eps
82- // next_m = m*b1 + g*(1-b1)
83- using NextM = Bind<Vec::Add<U>, Bind<Vec::Mul<U>, M, B1>, Bind<Vec::Mul<U>, G, OmB1>>;
84- using MUnbias = Bind<Vec::Div<U>, NextM, Rd0>;
85- // y1 = param*wd + m_unbiased/sqrt(v_unbiased+eps)
86- using Y1 = Bind<Vec::Add<U>, Bind<Vec::Mul<U>, Param, Wd>, Bind<Vec::Div<U>, MUnbias, SqrtVeps>>;
87- using Y4 = Bind<Vec::Div<U>, MUnbias, SqrtVAddEps>;
88- 
89- using Y1Cast = Bind<Vec::Cast<T, U, 1>, Y1>;
90- using Y2Cast = Bind<Vec::Cast<T, U, 1>, NextM>;
91- using Y3Cast = Bind<Vec::Cast<T, U, 1>, NextV>;
92- using Y4Cast = Bind<Vec::Cast<T, U, 1>, Y4>;
93- 
94- using OpOut0 = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, Y1Cast>;
95- using OpOut1 = Bind<Vec::CopyOut<T>, Placeholder::Out1<T>, Y2Cast>;
96- using OpOut2 = Bind<Vec::CopyOut<T>, Placeholder::Out2<T>, Y3Cast>;
97- using OpOut3 = Bind<Vec::CopyOut<T>, Placeholder::Out3<T>, Y4Cast>;
98- 
99- using Outputs = Elems<OpOut0, OpOut1, OpOut2, OpOut3>;
100- using MemCfg = MemOptCfg<MemLevel::LEVEL_1>;
101- using OpDag = DAGSch<Outputs, void, MemCfg>;
102-};
103-} // namespace LambNextMVOp
104-#endif // LAMB_NEXT_M_V_DAG_H
Doptim/lamb_next_m_v/op_kernel/arch35/lamb_next_m_v_tiling_key.h+0-27
@@ -1,27 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_next_m_v_tiling_key.h
13- * \brief lamb_next_m_v_tiling_key head file
14- */
15- 
16-#ifndef ADAM_APPLY_ONE_STRUCT_H
17-#define ADAM_APPLY_ONE_STRUCT_H
18- 
19-#include "atvoss/broadcast/broadcast_base_struct.h"
20- 
21-using namespace Ops::Base;
22-// 算子自定义的tiling key字段
23-ASCENDC_TPL_ARGS_DECL(LambNextMV, BRC_NDDMA_SCH_MODE_KEY_DECL(schMode));
24- 
25-ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(BRC_NDDMA_SCH_MODE_KEY_SEL(schMode)));
26- 
27-#endif // ADAM_APPLY_ONE_STRUCT_H
Roptim/lamb_next_m_v_with_decay/op_kernel/lamb_next_m_v_with_decay.cpp→optim/lamb_next_m_v/op_kernel/arch35/lamb_next_mv.cpp+25-19文件内容审核中，请稍后刷新重试
Doptim/lamb_next_m_v/op_kernel/lamb_next_m_v.cpp+0-43
@@ -1,43 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_next_m_v.cpp
13- * \brief lamb_next_m_v.cpp
14- */
15- 
16-#include "kernel_operator.h"
17-#include "arch35/lamb_next_m_v_dag.h"
18-#include "arch35/lamb_next_m_v_tiling_key.h"
19-#include "atvoss/broadcast/broadcast_sch.h"
20- 
21-using namespace AscendC;
22-using namespace Ops::Base;
23- 
24-template <uint64_t schMode>
25-__global__ __aicore__ void lamb_next_mv(GM_ADDR input_mul3, GM_ADDR input_mul2, GM_ADDR input_realdiv1,
26- GM_ADDR input_mul1, GM_ADDR input_mul0, GM_ADDR input_realdiv0,
27- GM_ADDR input_mul4, GM_ADDR mul0_x, GM_ADDR mul1_sub, GM_ADDR mul2_x,
28- GM_ADDR mul3_sub1, GM_ADDR mul4_x, GM_ADDR add2_y, GM_ADDR y1, GM_ADDR y2,
29- GM_ADDR y3, GM_ADDR y4, GM_ADDR workspace, GM_ADDR tiling)
30-{
31- KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
32- if constexpr (std::is_same<DTYPE_INPUT_MUL3, half>::value) {
33- using OpDag = LambNextMVOp::LambNextMVCompute<half, float>::OpDag;
34- BroadcastSch<schMode, OpDag> sch(tiling);
35- sch.Process(input_mul3, input_mul2, input_realdiv1, input_mul1, input_mul0, input_realdiv0, input_mul4, mul0_x,
36- mul1_sub, mul2_x, mul3_sub1, mul4_x, add2_y, y1, y2, y3, y4);
37- } else {
38- using OpDag = LambNextMVOp::LambNextMVCompute<float, float>::OpDag;
39- BroadcastSch<schMode, OpDag> sch(tiling);
40- sch.Process(input_mul3, input_mul2, input_realdiv1, input_mul1, input_mul0, input_realdiv0, input_mul4, mul0_x,
41- mul1_sub, mul2_x, mul3_sub1, mul4_x, add2_y, y1, y2, y3, y4);
42- }
43-}
Moptim/lamb_next_m_v/tests/assets/golden.py+17-20
@@ -19,14 +19,18 @@ __golden__ = {"kernel": {"lamb_next_mv": "lamb_next_mv_golden"}}
19 19 
20def _scalars(*xs):20def _scalars(*xs):
21 """取标量并落到 float32 torch 标量张量上。不能返回 Python float——那是 fp64，标量21 """取标量并落到 float32 torch 标量张量上。不能返回 Python float——那是 fp64，标量
22- 运算会被抬到双精度，而算子在 fp32 上算（A2 的 TBE compute 里 dtype='float32'，22+ 运算会被抬到双精度，而 arch35 内核的计算类型 U = float（fp16 输入 unpack 成 fp32 再算）。
23+ 注：A2 并**不**升精度 —— canndev 的 tbe impl 是 tvm.placeholder(dtype=input_dtype)，全程无 cast_to，
24+ fp16 输入就在 fp16 上做 vmul/vdiv/vsqrt。此处跟随 arch35 的计算类型，不跟随 A2（
23 arch35 DAG 的计算类型 U = float）。numpy 只用于取值与 dtype 转换。"""25 arch35 DAG 的计算类型 U = float）。numpy 只用于取值与 dtype 转换。"""
24 # 标量落成 **0 维 float64** 张量: torch 的类型提升里 0 维张量不会把 dim>0 的张量抬档，26 # 标量落成 **0 维 float64** 张量: torch 的类型提升里 0 维张量不会把 dim>0 的张量抬档，
25 # 所以数据是 fp32 时结果仍是 fp32(与改动前一致)，数据被 Promote 成 fp64 时标量自动27 # 所以数据是 fp32 时结果仍是 fp32(与改动前一致)，数据被 Promote 成 fp64 时标量自动
26 # 跟到 fp64，不会用一个先降到 fp32 的标量去污染高精度真值。28 # 跟到 fp64，不会用一个先降到 fp32 的标量去污染高精度真值。
27- return tuple(29+ # A2 语义: 这些"系数"输入是**可广播的 ND Tensor**(canndev ops/built-in/tbe/impl/lamb_*.py
28- torch.from_numpy(np.asarray(x, "float64").reshape(-1)[:1])[0] for x in xs30+ # 每步 mul/sub/div 都先 shape_util.broadcast_shapes 再 tbe.broadcast)。原先 reshape(-1)[:1]
29- )31+ # 只取首元素, 传多元素张量时静默按首元素计算 —— 与内核的广播实现不一致, 广播档必然假红。
32+ # 改为返回完整张量交给 torch 自然广播: 形状 (1,) 的行为与原标量完全一致, 故常规档不变。
33+ return tuple(_t(x) for x in xs)
30 34 
31 35 
32def _t(x):36def _t(x):
@@ -140,31 +144,24 @@ def _tp_t(x):
140 144 
141 145 
142def _tp_widen(t):146def _tp_widen(t):
143- """把三方入参加宽到**内核的计算类型**, 复刻 op_kernel/arch35/lamb_next_m_v_dag.h 的147+ """按 NPU 的加宽行为加宽三方入参。
144- `Cast<U, T, 0>`(见该文件注释: "Compute in U (=float): half casts to float")。
145 148 
146- 这不是"给竞品放水抬精度", 而是**同算法转写**——本算子是融合 DAG, 十几步中间量全程留在149+ 三方腿拿到的是**原始 dtype T**(TTK 的 Promote 只作用于 golden)。内核对 fp16 输入 unpack
147- fp32、一次都不落回 T; 而 torch 只在**单个算子内部**用 opmath=float, 算子之间每一步都把150+ 成 fp32 全程不落回(dag.h 的 `Cast<U, T, 0>`, U = float), 而 torch 只在单个算子内部用
148- 结果落回 fp16(A100 实测: fp16 的 300*300 直接得 inf, 而 (a*a)/a 真值 300 明明存得下)。151+ opmath=float、算子之间每步落回 T —— 不加宽就等于拿"逐步截断的实现"当竞品。
149- 不加宽就等于拿"逐步截断的实现"当竞品, 与被测内核不是同一个算法。152+ T = fp32 时内核计算类型即 fp32, 不加宽; 整型不动(走 fp32 会抹掉 >2^24 的低位)。
150- 
151- 整型不动: 内核对 int 也是原生/int32 累加, 走 fp32 会把 >2^24 抹掉低位。
152 """153 """
153 return t.float() if t.dtype in (torch.float16, torch.bfloat16) else t154 return t.float() if t.dtype in (torch.float16, torch.bfloat16) else t
154 155 
155 156 
156def _tp_narrow(outs, dt):157def _tp_narrow(outs, dt):
157- """出口复刻内核的 `Cast<T, U, 1>`: 算完窄回算子输出 dtype。158+ """出口复刻内核的 `Cast<T, U, 1>`: 窄回算子声明的 dtype T。"""
158- 
159- 这个 cast **必须与 _tp_widen 成对出现**: 少了它, 三方停在 fp32, 会与走 TTK Promote
160- (fp16->fp32) 的 golden 逐位相等 —— 双标杆塌成单标杆, 三比值分母被夹到 §4.5.1 的 err,
161- 有量纲的 RMSE 比值随输出量级线性放大而假红。
162- """
163 return [o.to(dt) if o.is_floating_point() else o for o in outs]159 return [o.to(dt) if o.is_floating_point() else o for o in outs]
164 160 
165 161 
166def _tp_s(x):162def _tp_s(x):
167- return _tp_t(x).reshape(-1)[0]163+ # 同 _scalars: 三方腿也必须广播, 不能只取首元素
164+ return _tp_t(x)
168 165 
169 166 
170class _LambNextMVCompose:167class _LambNextMVCompose:
@@ -185,7 +182,7 @@ class _LambNextMVCompose:
185 add2_y,182 add2_y,
186 **kwargs,183 **kwargs,
187 ):184 ):
188- _dt = _tp_t(input_mul3).dtype # 算子输出 dtype185+ _dt = _tp_t(input_mul3).dtype # 算子声明的 dtype T(三方腿入参不经 Promote)
189 g2, v, g, m, param = (186 g2, v, g, m, param = (
190 _tp_widen(_tp_t(t))187 _tp_widen(_tp_t(t))
191 for t in (input_mul3, input_mul2, input_mul1, input_mul0, input_mul4)188 for t in (input_mul3, input_mul2, input_mul1, input_mul0, input_mul4)
Moptim/lamb_next_m_v_with_decay/CMakeLists.txt+3-3
@@ -1,9 +1,9 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4# CANN Open Software License Agreement Version 2.0 (the "License").4# CANN Open Software License Agreement Version 2.0 (the "License").
5# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
6-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
@@ -13,4 +13,4 @@ set(SUPPORT_COMPUTE_UNIT "ascend950")
13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译13# 设置每种芯片类型对应的tiling文件目录，即采用op_host目录下哪个文件夹下的tiling文件编译
14set(SUPPORT_TILING_DIR "arch35")14set(SUPPORT_TILING_DIR "arch35")
15 15 
16-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_next_m_v_with_decay ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)16+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE lamb_next_m_v_with_decay ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DEPENDENCIES lamb_apply_common DISABLE_IN_OPP TRUE)
Moptim/lamb_next_m_v_with_decay/README.md+15-13
@@ -48,7 +48,7 @@
48 <tr>48 <tr>
49 <td>input_mul3</td>49 <td>input_mul3</td>
50 <td>输入</td>50 <td>输入</td>
51- <td>支持空Tensor。公式中的input_mul3（g^2），主张量，shape需与input_mul0满足broadcast关系。</td>51+ <td>支持空Tensor。公式中的input_mul3（g^2），主张量，shape需与其他输入满足broadcast关系。</td>
52 <td>FLOAT16、FLOAT</td>52 <td>FLOAT16、FLOAT</td>
53 <td>ND</td>53 <td>ND</td>
54 </tr>54 </tr>
@@ -76,7 +76,7 @@
76 <tr>76 <tr>
77 <td>input_mul0</td>77 <td>input_mul0</td>
78 <td>输入</td>78 <td>输入</td>
79- <td>支持空Tensor。公式中的input_mul0（一阶矩m），主张量，shape需与input_mul3满足broadcast关系，其broadcast结果决定各输出的shape。</td>79+ <td>支持空Tensor。公式中的input_mul0（一阶矩m），主张量，shape需与其他输入满足broadcast关系，全部输入的broadcast结果决定各输出的shape。</td>
80 <td>FLOAT16、FLOAT</td>80 <td>FLOAT16、FLOAT</td>
81 <td>ND</td>81 <td>ND</td>
82 </tr>82 </tr>
@@ -97,70 +97,70 @@
97 <tr>97 <tr>
98 <td>mul0_x</td>98 <td>mul0_x</td>
99 <td>输入</td>99 <td>输入</td>
100- <td>不支持空Tensor。公式中的mul0_x（beta1），标量。</td>100+ <td>支持空Tensor。公式中的mul0_x（beta1），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
101 <td>FLOAT16、FLOAT</td>101 <td>FLOAT16、FLOAT</td>
102 <td>ND</td>102 <td>ND</td>
103 </tr>103 </tr>
104 <tr>104 <tr>
105 <td>mul1_sub</td>105 <td>mul1_sub</td>
106 <td>输入</td>106 <td>输入</td>
107- <td>不支持空Tensor。公式中的mul1_sub（1-beta1），标量。</td>107+ <td>支持空Tensor。公式中的mul1_sub（1-beta1），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
108 <td>FLOAT16、FLOAT</td>108 <td>FLOAT16、FLOAT</td>
109 <td>ND</td>109 <td>ND</td>
110 </tr>110 </tr>
111 <tr>111 <tr>
112 <td>mul2_x</td>112 <td>mul2_x</td>
113 <td>输入</td>113 <td>输入</td>
114- <td>不支持空Tensor。公式中的mul2_x（beta2），标量。</td>114+ <td>支持空Tensor。公式中的mul2_x（beta2），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
115 <td>FLOAT16、FLOAT</td>115 <td>FLOAT16、FLOAT</td>
116 <td>ND</td>116 <td>ND</td>
117 </tr>117 </tr>
118 <tr>118 <tr>
119 <td>mul3_sub1</td>119 <td>mul3_sub1</td>
120 <td>输入</td>120 <td>输入</td>
121- <td>不支持空Tensor。公式中的mul3_sub1（1-beta2），标量。</td>121+ <td>支持空Tensor。公式中的mul3_sub1（1-beta2），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
122 <td>FLOAT16、FLOAT</td>122 <td>FLOAT16、FLOAT</td>
123 <td>ND</td>123 <td>ND</td>
124 </tr>124 </tr>
125 <tr>125 <tr>
126 <td>mul4_x</td>126 <td>mul4_x</td>
127 <td>输入</td>127 <td>输入</td>
128- <td>不支持空Tensor。公式中的mul4_x（权重衰减系数），标量。</td>128+ <td>支持空Tensor。公式中的mul4_x（权重衰减系数），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
129 <td>FLOAT16、FLOAT</td>129 <td>FLOAT16、FLOAT</td>
130 <td>ND</td>130 <td>ND</td>
131 </tr>131 </tr>
132 <tr>132 <tr>
133 <td>add2_y</td>133 <td>add2_y</td>
134 <td>输入</td>134 <td>输入</td>
135- <td>不支持空Tensor。公式中的add2_y（epsilon），标量。</td>135+ <td>支持空Tensor。公式中的add2_y（epsilon），shape支持1-8维，需与其他输入满足broadcast规则（右对齐，对应维相等或为1）。</td>
136 <td>FLOAT16、FLOAT</td>136 <td>FLOAT16、FLOAT</td>
137 <td>ND</td>137 <td>ND</td>
138 </tr>138 </tr>
139 <tr>139 <tr>
140 <td>y1</td>140 <td>y1</td>
141 <td>输出</td>141 <td>输出</td>
142- <td>支持空Tensor。公式中的y1（update），shape取input_mul3与input_mul0的broadcast结果。</td>142+ <td>支持空Tensor。公式中的y1（update），shape取全部输入的broadcast结果。</td>
143 <td>FLOAT16、FLOAT</td>143 <td>FLOAT16、FLOAT</td>
144 <td>ND</td>144 <td>ND</td>
145 </tr>145 </tr>
146 <tr>146 <tr>
147 <td>y2</td>147 <td>y2</td>
148 <td>输出</td>148 <td>输出</td>
149- <td>支持空Tensor。公式中的y2（next_m），shape取input_mul3与input_mul0的broadcast结果。</td>149+ <td>支持空Tensor。公式中的y2（next_m），shape取全部输入的broadcast结果。</td>
150 <td>FLOAT16、FLOAT</td>150 <td>FLOAT16、FLOAT</td>
151 <td>ND</td>151 <td>ND</td>
152 </tr>152 </tr>
153 <tr>153 <tr>
154 <td>y3</td>154 <td>y3</td>
155 <td>输出</td>155 <td>输出</td>
156- <td>支持空Tensor。公式中的y3（next_v），shape取input_mul3与input_mul0的broadcast结果。</td>156+ <td>支持空Tensor。公式中的y3（next_v），shape取全部输入的broadcast结果。</td>
157 <td>FLOAT16、FLOAT</td>157 <td>FLOAT16、FLOAT</td>
158 <td>ND</td>158 <td>ND</td>
159 </tr>159 </tr>
160 <tr>160 <tr>
161 <td>y4</td>161 <td>y4</td>
162 <td>输出</td>162 <td>输出</td>
163- <td>支持空Tensor。公式中的y4，shape取input_mul3与input_mul0的broadcast结果。</td>163+ <td>支持空Tensor。公式中的y4，shape取全部输入的broadcast结果。</td>
164 <td>FLOAT16、FLOAT</td>164 <td>FLOAT16、FLOAT</td>
165 <td>ND</td>165 <td>ND</td>
166 </tr>166 </tr>
@@ -168,8 +168,10 @@
168 168 
169## 约束说明169## 约束说明
170 170 
171+- 所有输入的shape需两两满足broadcast规则（右对齐，对应维相等或为1），输出shape为全部输入的broadcast结果。
172+- 所有输入及输出的维度数不超过8。
171- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。173- 所有输入的数据类型必须一致，同为FLOAT16或同为FLOAT。
172-- input_mul0/input_mul1/input_mul2/input_mul3/input_mul4 为主张量，其shape需保持一致（或可相互广播到同一shape）；各输出y1/y2/y3/y4的shape均取该广播结果（实现以input_mul3与input_mul0的broadcast结果为准）。174+- 各输出y1/y2/y3/y4的shape均取全部输入的broadcast结果。
173 175 
174## 调用说明176## 调用说明
175 177 
Moptim/lamb_next_m_v_with_decay/op_host/arch35/lamb_next_m_v_with_decay_tiling_arch35.cpp+45-51
@@ -15,20 +15,21 @@
15 15 
16#include "lamb_next_m_v_with_decay_tiling_arch35.h"16#include "lamb_next_m_v_with_decay_tiling_arch35.h"
17#include <graph/utils/type_utils.h>17#include <graph/utils/type_utils.h>
18-#include "../../op_kernel/arch35/lamb_next_m_v_with_decay_dag.h"18+#include <securec.h>
19-#include "atvoss/broadcast/broadcast_tiling.h"19+#include <algorithm>
20#include "log/log.h"20#include "log/log.h"
21#include "platform/platform_info.h"21#include "platform/platform_info.h"
22#include "register/op_impl_registry.h"22#include "register/op_impl_registry.h"
23#include "register/tilingdata_base.h"23#include "register/tilingdata_base.h"
24#include "op_host/tiling_templates_registry.h"24#include "op_host/tiling_templates_registry.h"
25 25 
26-using namespace AscendC;
27using namespace ge;26using namespace ge;
28 27 
29namespace optiling {28namespace optiling {
30 29 
31constexpr static uint64_t LAMB_NEXT_M_V_TILING_PRIORITY = 0;30constexpr static uint64_t LAMB_NEXT_M_V_TILING_PRIORITY = 0;
31+constexpr static uint64_t TILING_KEY_FP32 = 100;
32+constexpr static uint64_t TILING_KEY_FP16 = 200;
32constexpr static int32_t INPUT_NUM = 13;33constexpr static int32_t INPUT_NUM = 13;
33constexpr static int32_t OUTPUT_NUM = 4;34constexpr static int32_t OUTPUT_NUM = 4;
34 35 
@@ -80,18 +81,6 @@ ge::graphStatus LambNextMVWithDecayTiling::GetShapeAttrsInfo()
80 return ge::GRAPH_FAILED;81 return ge::GRAPH_FAILED;
81 }82 }
82 }83 }
83- // 标量类输入为每元素计算所必需的系数, 空Tensor 视为缺失必选值(畸形输入), 不支持。
84- static const int32_t kScalarInputIdx[] = {7, 8, 9, 10, 11, 12};
85- for (int32_t scalarIdx : kScalarInputIdx) {
86- auto scalarShape = context_->GetInputShape(scalarIdx);
87- OP_CHECK_NULL_WITH_CONTEXT(context_, scalarShape);
88- if (scalarShape->GetStorageShape().GetShapeSize() == 0) {
89- OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), kInputNames[scalarIdx],
90- Ops::Base::ToString(scalarShape->GetStorageShape()).c_str(),
91- "scalar input does not support empty tensor");
92- return ge::GRAPH_FAILED;
93- }
94- }
95 return ge::GRAPH_SUCCESS;84 return ge::GRAPH_SUCCESS;
96}85}
97 86 
@@ -99,49 +88,44 @@ bool LambNextMVWithDecayTiling::IsCapable() { return true; }
99 88 
100ge::graphStatus LambNextMVWithDecayTiling::DoOpTiling()89ge::graphStatus LambNextMVWithDecayTiling::DoOpTiling()
101{90{
102- // 空 tensor 应对(空进空出): 输出为空(0元素)时设 1 核(空转), 配合全0 tiling 数据(blockFormer=0)使 kernel 空转退出,91+ auto rawTilingData = context_->GetRawTilingData();
103- // 直接成功。92+ OP_CHECK_NULL_WITH_CONTEXT(context_, rawTilingData);
104- auto emptyTensorOutShape0 = context_->GetOutputShape(0);
105- if (emptyTensorOutShape0 != nullptr && emptyTensorOutShape0->GetStorageShape().GetShapeSize() == 0) {
106- auto emptyRawTiling = context_->GetRawTilingData();
107- if (emptyRawTiling != nullptr && emptyRawTiling->GetData() != nullptr) {
108- size_t emptyCap = emptyRawTiling->GetCapacity();
109- uint8_t* emptyPtr = static_cast<uint8_t*>(emptyRawTiling->GetData());
110- for (size_t emptyIdx = 0; emptyIdx < emptyCap; ++emptyIdx) {
111- emptyPtr[emptyIdx] = 0;
112- }
113- emptyRawTiling->SetDataSize(emptyCap);
114- }
115- size_t* emptyWs = context_->GetWorkspaceSizes(1);
116- if (emptyWs != nullptr) {
117- emptyWs[0] = 0;
118- }
119- context_->SetBlockDim(1);
120- tilingKey = GET_TPL_TILING_KEY(1); // schMode=1(已编译), 配合全0 tiling(blockFormer=0)空转
121- return ge::GRAPH_SUCCESS;
122- }
123 auto input0Desc = context_->GetInputDesc(0);93 auto input0Desc = context_->GetInputDesc(0);
124 OP_CHECK_NULL_WITH_CONTEXT(context_, input0Desc);94 OP_CHECK_NULL_WITH_CONTEXT(context_, input0Desc);
95+ 
125 ge::DataType input0DType = input0Desc->GetDataType();96 ge::DataType input0DType = input0Desc->GetDataType();
126- if (input0DType == ge::DT_FLOAT16) {97+ uint32_t dtSize = 0;
127- BroadcastBaseTiling<LambNextMVWithDecayOp::LambNextMVWithDecayCompute<half, float>::OpDag> brcBaseTiling(98+ if (input0DType == ge::DT_FLOAT) {
128- context_, static_cast<uint32_t>(BROADCAST_KERNEL_TYPE::KERNEL_TYPE_NDDMA));99+ tilingKey = TILING_KEY_FP32;
129- OP_CHECK_IF(brcBaseTiling.DoTiling() == ge::GRAPH_FAILED,100+ dtSize = sizeof(float);
130- OP_LOGE(context_->GetNodeName(), "Do tiling failed. Please check the detailed log."),101+ } else if (input0DType == ge::DT_FLOAT16) {
131- return ge::GRAPH_FAILED);102+ tilingKey = TILING_KEY_FP16;
132- tilingKey = GET_TPL_TILING_KEY(brcBaseTiling.GetSchMode());103+ dtSize = sizeof(uint16_t);
133- } else if (input0DType == ge::DT_FLOAT) {
134- BroadcastBaseTiling<LambNextMVWithDecayOp::LambNextMVWithDecayCompute<float, float>::OpDag> brcBaseTiling(
135- context_, static_cast<uint32_t>(BROADCAST_KERNEL_TYPE::KERNEL_TYPE_NDDMA));
136- OP_CHECK_IF(brcBaseTiling.DoTiling() == ge::GRAPH_FAILED,
137- OP_LOGE(context_->GetNodeName(), "Do tiling failed. Please check the detailed log."),
138- return ge::GRAPH_FAILED);
139- tilingKey = GET_TPL_TILING_KEY(brcBaseTiling.GetSchMode());
140 } else {104 } else {
141 OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "input_mul3", Ops::Base::ToString(input0DType).c_str(),105 OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "input_mul3", Ops::Base::ToString(input0DType).c_str(),
142 "fp16 or fp32");106 "fp16 or fp32");
143 return ge::GRAPH_FAILED;107 return ge::GRAPH_FAILED;
144 }108 }
109+ 
110+ using PlanTiling = LambBrcTilingData<13, 4>;
111+ td_ = PlanTiling{};
112+ // 空进空出: 输出 0 元素时 tiling 全 0, kernel 按 usedCoreNum=0 直接退出。
113+ auto outShape0 = context_->GetOutputShape(0);
114+ if (outShape0 == nullptr || outShape0->GetStorageShape().GetShapeSize() != 0) {
115+ // 先取返回值再判: 模板实参里的逗号会被预处理器当成 OP_CHECK_IF 的参数分隔符。
116+ ge::graphStatus planRet = BuildLambBrcPlan<13, 4>(context_, coreNum_, ubSize_, dtSize, td_);
117+ OP_CHECK_IF(planRet != ge::GRAPH_SUCCESS, OP_LOGE(context_->GetNodeName(), "build broadcast plan failed"),
118+ return ge::GRAPH_FAILED);
119+ }
120+ 
121+ auto ret = memcpy_s(rawTilingData->GetData(), rawTilingData->GetCapacity(), &td_, sizeof(td_));
122+ OP_CHECK_IF(ret != EOK, OP_LOGE(context_->GetNodeName(), "copy tiling data failed, ret %d", ret),
123+ return ge::GRAPH_FAILED);
124+ rawTilingData->SetDataSize(sizeof(td_));
125+ context_->SetBlockDim(std::max<uint32_t>(td_.usedCoreNum, 1));
126+ size_t* ws = context_->GetWorkspaceSizes(1);
127+ OP_CHECK_NULL_WITH_CONTEXT(context_, ws);
128+ ws[0] = 0U;
145 return ge::GRAPH_SUCCESS;129 return ge::GRAPH_SUCCESS;
146}130}
147 131 
@@ -153,7 +137,17 @@ ge::graphStatus LambNextMVWithDecayTiling::GetWorkspaceSize() { return ge::GRAPH
153 137 
154ge::graphStatus LambNextMVWithDecayTiling::PostTiling() { return ge::GRAPH_SUCCESS; }138ge::graphStatus LambNextMVWithDecayTiling::PostTiling() { return ge::GRAPH_SUCCESS; }
155 139 
156-ge::graphStatus LambNextMVWithDecayTiling::GetPlatformInfo() { return ge::GRAPH_SUCCESS; }140+ge::graphStatus LambNextMVWithDecayTiling::GetPlatformInfo()
141+{
142+ auto compileInfo = static_cast<const LambNextMVWithDecayCompileInfo*>(context_->GetCompileInfo());
143+ OP_CHECK_NULL_WITH_CONTEXT(context_, compileInfo);
144+ coreNum_ = compileInfo->coreNum;
145+ ubSize_ = compileInfo->ubSize;
146+ OP_CHECK_IF(coreNum_ == 0 || ubSize_ == 0,
147+ OP_LOGE(context_->GetNodeName(), "invalid platform info: coreNum %lu ubSize %lu", coreNum_, ubSize_),
148+ return ge::GRAPH_FAILED);
149+ return ge::GRAPH_SUCCESS;
150+}
157 151 
158static ge::graphStatus TilingForLambNextMVWithDecay(gert::TilingContext* context)152static ge::graphStatus TilingForLambNextMVWithDecay(gert::TilingContext* context)
159{153{
Moptim/lamb_next_m_v_with_decay/op_host/lamb_next_m_v_with_decay_infershape.cpp+34-12
@@ -13,6 +13,7 @@
13 * \brief13 * \brief
14 */14 */
15 15 
16+#include <vector>
16#include "register/op_impl_registry.h"17#include "register/op_impl_registry.h"
17#include "log/log.h"18#include "log/log.h"
18#include "infershape_broadcast_util.h"19#include "infershape_broadcast_util.h"
@@ -20,28 +21,49 @@
20using namespace Ops::Base;21using namespace Ops::Base;
21using namespace ge;22using namespace ge;
22namespace ops {23namespace ops {
23-// full tensors: input_mul3(0,g^2), input_mul0(4,m). All four outputs share their broadcast shape.24+// A2 语义: 本族算子的所有输入都是可广播的 ND Tensor(见 canndev
24-constexpr size_t IN_MUL3 = 0;25+// ops/built-in/tbe/impl/lamb_*.py, 每一步 mul/sub/div 都先 shape_util.broadcast_shapes
25-constexpr size_t IN_MUL0 = 4;26+// 再 tbe.broadcast), 输出形状为全部输入广播的结果。A2 的 op_proto 只声明了其中两个输入,
27+// 属声明宽松, 不作为支持面依据。
28+constexpr size_t IN_NUM = 13;
26constexpr size_t OUT_NUM = 4;29constexpr size_t OUT_NUM = 4;
27 30 
28static ge::graphStatus InferShape4LambNextMVWithDecay(gert::InferShapeContext* context)31static ge::graphStatus InferShape4LambNextMVWithDecay(gert::InferShapeContext* context)
29{32{
30- auto g2 = context->GetInputShape(IN_MUL3);33+ std::vector<const gert::Shape*> inShapes;
31- OP_CHECK_NULL_WITH_CONTEXT(context, g2);34+ inShapes.reserve(IN_NUM);
32- auto m = context->GetInputShape(IN_MUL0);35+ for (size_t i = 0; i < IN_NUM; i++) {
33- OP_CHECK_NULL_WITH_CONTEXT(context, m);36+ auto in = context->GetInputShape(i);
37+ OP_CHECK_NULL_WITH_CONTEXT(context, in);
38+ inShapes.push_back(in);
39+ }
40+ gert::Shape bcShape;
41+ // 逐对折叠广播: 只用两参数重载。vector 重载虽在 op_common/op_host/infershape_broadcast_util.h
42+ // 以 Ops::Base 声明, 但 libops_base.so 只导出两参数版, vector 版的实现在 libop_common.so 的
43+ // 小写 ops 命名空间下 —— 声明与实现命名空间不一致, 用它会编译期通过、加载期
44+ // undefined symbol 而装不上包。
45+ bcShape = *inShapes[0];
46+ for (size_t i = 1; i < inShapes.size(); i++) {
47+ gert::Shape tmp;
48+ OP_CHECK_IF(!BroadcastShape(&bcShape, inShapes[i], &tmp),
49+ OP_LOGE(context->GetNodeName(), "input shapes cannot broadcast together"), return ge::GRAPH_FAILED);
50+ bcShape = tmp;
51+ }
52+ // 标量归一: 全标量输入 broadcast 得 0 维空 shape (), 与 A2 的 shape_util.scalar2tensor_one
53+ // 对齐, 归一为 (1,)。否则动态 shape 编译期 DFX 生成会对空 shape 做 reduce 连乘(无初值)而报
54+ // TypeError 编译失败。
55+ if (bcShape.GetDimNum() == 0) {
56+ bcShape.SetDimNum(1);
57+ bcShape.SetDim(0, 1);
58+ }
34 for (size_t i = 0; i < OUT_NUM; i++) {59 for (size_t i = 0; i < OUT_NUM; i++) {
35 auto out = context->GetOutputShape(i);60 auto out = context->GetOutputShape(i);
36 OP_CHECK_NULL_WITH_CONTEXT(context, out);61 OP_CHECK_NULL_WITH_CONTEXT(context, out);
37- OP_CHECK_IF(!BroadcastShape(g2, m, out),62+ *out = bcShape;
38- OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(context->GetNodeName(), "input_mul3 and input_mul0",
39- (ToString(*g2) + " and " + ToString(*m)).c_str(),
40- "shapes cannot broadcast"),
41- return ge::GRAPH_FAILED);
42 }63 }
43 return GRAPH_SUCCESS;64 return GRAPH_SUCCESS;
44}65}
66+ 
45static ge::graphStatus InferDataType4LambNextMVWithDecay(gert::InferDataTypeContext* context)67static ge::graphStatus InferDataType4LambNextMVWithDecay(gert::InferDataTypeContext* context)
46{68{
47 if (context == nullptr) {69 if (context == nullptr) {
Aoptim/lamb_next_m_v_with_decay/op_host/lamb_next_m_v_with_decay_tiling_def.h+52-0文件内容审核中，请稍后刷新重试
Aoptim/lamb_next_m_v_with_decay/op_kernel/CMakeLists.txt+16-0文件内容审核中，请稍后刷新重试
Doptim/lamb_next_m_v_with_decay/op_kernel/arch35/lamb_next_m_v_with_decay_dag.h+0-105
@@ -1,105 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_next_m_v_with_decay_dag.h
13- * \brief lamb_next_m_v_with_decay_dag head file
14- */
15- 
16-#ifndef LAMB_NEXT_M_V_DAG_H
17-#define LAMB_NEXT_M_V_DAG_H
18-#include "atvoss/util/dag.h"
19-#include "atvoss/util/vec.h"
20-#include "atvoss/util/placeholder.h"
21- 
22-// The atvoss Placeholder ships In0..In11 (12 inputs). LambNextMVWithDecay has 13 inputs; extend with In12
23-// following the framework's exact pattern (additive: the struct + its IsInHolder trait, same namespace).
24-namespace Ops {
25-namespace Base {
26-namespace Placeholder {
27-template <class T, class Attr = InAttr<>>
28-struct In12 : public Holder<T, Attr, 12> {};
29-template <class U, class Attr>
30-struct IsInHolder<In12<U, Attr>> {
31- constexpr static bool Value = true;
32-};
33-} // namespace Placeholder
34-} // namespace Base
35-} // namespace Ops
36- 
37-namespace LambNextMVWithDecayOp {
38-using namespace AscendC;
39-using namespace Ops::Base;
40-// In : 0 input_mul3(g^2) 1 input_mul2(v) 2 input_realdiv1(1-b2^t) 3 input_mul1(g) 4 input_mul0(m)
41-// 5 input_realdiv0(1-b1^t) 6 input_mul4(param) 7 mul0_x(b1) 8 mul1_sub(1-b1) 9 mul2_x(b2)
42-// 10 mul3_sub1(1-b2) 11 mul4_x(wd) 12 add2_y(eps) [0,1,3,4,6 full; rest scalar]
43-// Out: 0 y1(update) 1 y2(next_m) 2 y3(next_v) 3 y4(m_unbiased/(sqrt(v_unbiased)+eps))
44-// Compute in U (=float): half casts to float (div/sqrt need float precision). rsqrt -> Div(_,Sqrt).
45-template <typename T, typename U>
46-struct LambNextMVWithDecayCompute {
47- // full tensors
48- using InG2 = Bind<Vec::CopyInBrc<T>, Placeholder::In0<T>>;
49- using InV = Bind<Vec::CopyInBrc<T>, Placeholder::In1<T>>;
50- using InG = Bind<Vec::CopyInBrc<T>, Placeholder::In3<T>>;
51- using InM = Bind<Vec::CopyInBrc<T>, Placeholder::In4<T>>;
52- using InParam = Bind<Vec::CopyInBrc<T>, Placeholder::In6<T>>;
53- // scalar coefficients
54- using InRd1 = Bind<Vec::Duplicate<T>, Placeholder::In2<T, Placeholder::ScalarAttr<true>>>;
55- using InRd0 = Bind<Vec::Duplicate<T>, Placeholder::In5<T, Placeholder::ScalarAttr<true>>>;
56- using InB1 = Bind<Vec::Duplicate<T>, Placeholder::In7<T, Placeholder::ScalarAttr<true>>>;
57- using InOmB1 = Bind<Vec::Duplicate<T>, Placeholder::In8<T, Placeholder::ScalarAttr<true>>>;
58- using InB2 = Bind<Vec::Duplicate<T>, Placeholder::In9<T, Placeholder::ScalarAttr<true>>>;
59- using InOmB2 = Bind<Vec::Duplicate<T>, Placeholder::In10<T, Placeholder::ScalarAttr<true>>>;
60- using InWd = Bind<Vec::Duplicate<T>, Placeholder::In11<T, Placeholder::ScalarAttr<true>>>;
61- using InEps = Bind<Vec::Duplicate<T>, Placeholder::In12<T, Placeholder::ScalarAttr<true>>>;
62- 
63- using G2 = Bind<Vec::Cast<U, T, 0>, InG2>;
64- using V = Bind<Vec::Cast<U, T, 0>, InV>;
65- using G = Bind<Vec::Cast<U, T, 0>, InG>;
66- using M = Bind<Vec::Cast<U, T, 0>, InM>;
67- using Param = Bind<Vec::Cast<U, T, 0>, InParam>;
68- using Rd1 = Bind<Vec::Cast<U, T, 0>, InRd1>;
69- using Rd0 = Bind<Vec::Cast<U, T, 0>, InRd0>;
70- using B1 = Bind<Vec::Cast<U, T, 0>, InB1>;
71- using OmB1 = Bind<Vec::Cast<U, T, 0>, InOmB1>;
72- using B2 = Bind<Vec::Cast<U, T, 0>, InB2>;
73- using OmB2 = Bind<Vec::Cast<U, T, 0>, InOmB2>;
74- using Wd = Bind<Vec::Cast<U, T, 0>, InWd>;
75- using Eps = Bind<Vec::Cast<U, T, 0>, InEps>;
76- 
77- // next_v = v*b2 + g^2*(1-b2)
78- using NextV = Bind<Vec::Add<U>, Bind<Vec::Mul<U>, V, B2>, Bind<Vec::Mul<U>, G2, OmB2>>;
79- using VUnbias = Bind<Vec::Div<U>, NextV, Rd1>;
80- using SqrtVeps = Bind<Vec::Sqrt<U>, Bind<Vec::Add<U>, VUnbias, Eps>>; // sqrt(v_unbiased+eps)
81- using SqrtVAddEps = Bind<Vec::Add<U>, Bind<Vec::Sqrt<U>, VUnbias>, Eps>; // sqrt(v_unbiased)+eps
82- // next_m = m*b1 + g*(1-b1)
83- using NextM = Bind<Vec::Add<U>, Bind<Vec::Mul<U>, M, B1>, Bind<Vec::Mul<U>, G, OmB1>>;
84- using MUnbias = Bind<Vec::Div<U>, NextM, Rd0>;
85- // y1 = param*wd + m_unbiased/sqrt(v_unbiased+eps)
86- using Y1 = Bind<Vec::Add<U>, Bind<Vec::Mul<U>, Param, Wd>, Bind<Vec::Div<U>, MUnbias, SqrtVeps>>;
87- // with_decay: y4 = param*wd + m_unbiased/(sqrt(v_unbiased)+eps) (NextMV has no param*wd term in y4)
88- using Y4 = Bind<Vec::Add<U>, Bind<Vec::Mul<U>, Param, Wd>, Bind<Vec::Div<U>, MUnbias, SqrtVAddEps>>;
89- 
90- using Y1Cast = Bind<Vec::Cast<T, U, 1>, Y1>;
91- using Y2Cast = Bind<Vec::Cast<T, U, 1>, NextM>;
92- using Y3Cast = Bind<Vec::Cast<T, U, 1>, NextV>;
93- using Y4Cast = Bind<Vec::Cast<T, U, 1>, Y4>;
94- 
95- using OpOut0 = Bind<Vec::CopyOut<T>, Placeholder::Out0<T>, Y1Cast>;
96- using OpOut1 = Bind<Vec::CopyOut<T>, Placeholder::Out1<T>, Y2Cast>;
97- using OpOut2 = Bind<Vec::CopyOut<T>, Placeholder::Out2<T>, Y3Cast>;
98- using OpOut3 = Bind<Vec::CopyOut<T>, Placeholder::Out3<T>, Y4Cast>;
99- 
100- using Outputs = Elems<OpOut0, OpOut1, OpOut2, OpOut3>;
101- using MemCfg = MemOptCfg<MemLevel::LEVEL_1>;
102- using OpDag = DAGSch<Outputs, void, MemCfg>;
103-};
104-} // namespace LambNextMVWithDecayOp
105-#endif // LAMB_NEXT_M_V_DAG_H
Doptim/lamb_next_m_v_with_decay/op_kernel/arch35/lamb_next_m_v_with_decay_tiling_key.h+0-27
@@ -1,27 +0,0 @@
1-/**
2- * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Please refer to the License for details. You may not use this file except in compliance with the License.
6- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7- * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8- * See LICENSE in the root of the software repository for the full text of the License.
9- */
10- 
11-/*!
12- * \file lamb_next_m_v_with_decay_tiling_key.h
13- * \brief lamb_next_m_v_with_decay_tiling_key head file
14- */
15- 
16-#ifndef ADAM_APPLY_ONE_STRUCT_H
17-#define ADAM_APPLY_ONE_STRUCT_H
18- 
19-#include "atvoss/broadcast/broadcast_base_struct.h"
20- 
21-using namespace Ops::Base;
22-// 算子自定义的tiling key字段
23-ASCENDC_TPL_ARGS_DECL(LambNextMVWithDecay, BRC_NDDMA_SCH_MODE_KEY_DECL(schMode));
24- 
25-ASCENDC_TPL_SEL(ASCENDC_TPL_ARGS_SEL(BRC_NDDMA_SCH_MODE_KEY_SEL(schMode)));
26- 
27-#endif // ADAM_APPLY_ONE_STRUCT_H
Aoptim/lamb_next_m_v_with_decay/op_kernel/arch35/lamb_next_mv_with_decay.cpp+48-0文件内容审核中，请稍后刷新重试
Moptim/lamb_next_right/op_kernel/arch35/lamb_next_right_dag.h+27-6文件内容审核中，请稍后刷新重试
Moptim/lamb_update_with_lr/op_kernel/arch35/lamb_update_with_lr_dag.h+30-9文件内容审核中，请稍后刷新重试
Moptim/lamb_update_with_lr/op_kernel/arch35/lamb_update_with_lr_tiling_key.h+6-1文件内容审核中，请稍后刷新重试
Moptim/lamb_update_with_lr_v2/op_kernel/arch35/lamb_update_with_lr_v2_dag.h+28-7文件内容审核中，请稍后刷新重试