Star945
1861
代码
代码
Issues223
Pull Requests333
Actions
项目讨论
Wiki
分析
Star945
1861
已合并
feat(sgd): 新增SGD优化器算子arch35实现（Ascend 950PR/950DT，GE图模式） #8155
raoliang_sac创建于 8月1日
feat(sgd): 新增SGD优化器算子arch35实现（Ascend 950PR/950DT，GE图模式） #8155
已合并
raoliang_sac创建于 8月1日
从raoliang_sac/ops-nnSGD-810合入到cann/ops-nnmaster
共 22 个文件变更+2719-0
Mdocs/zh/op_list.md+10-0
@@ -3878,6 +3878,16 @@
3878 <td>AI Core</td>3878 <td>AI Core</td>
3879 <td>实现FusedSgd融合优化器功能。</td>3879 <td>实现FusedSgd融合优化器功能。</td>
3880 </tr>3880 </tr>
3881+ <tr>
3882+ <td>optim</td>
3883+ <td><a href="../../optim/sgd/README.md">sgd</a></td>
3884+ <td>✓</td>
3885+ <td>✓</td>
3886+ <td>✗</td>
3887+ <td>✓</td>
3888+ <td>AI Core</td>
3889+ <td>带动量的随机梯度下降优化器更新算子，按weight_decay、dampening、Nesterov动量组合的SGD公式in-place更新权重parameters，同时原地回写动量累加器accum与首步标志stat。</td>
3890+ </tr>
3881 <tr>3891 <tr>
3882 <td>pooling</td>3892 <td>pooling</td>
3883 <td><a href="../../pooling/adaptive_avg_pool3d/README.md">adaptive_avg_pool3d</a></td>3893 <td><a href="../../pooling/adaptive_avg_pool3d/README.md">adaptive_avg_pool3d</a></td>
Aoptim/sgd/CMakeLists.txt+31-0
@@ -0,0 +1,31 @@
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+# ⚠️ 本算子【不提供 aclnn 单算子接口】，只走 GE 图模式下发。
12+# 依据：CANN 9.1.0 未定义 aclnnSgd（include/aclnnop/ 下只有 aclnn_fused_sgd.h，
13+# libopapi.so 无 aclnnSgd* 符号），canndev 全仓亦无 aclnn_sgd —— SGD 在上游本就是
14+# 纯图模式算子。故本仓不再引入 op_api/ 目录。
15+#
16+# add_modules_sources 仍放在【算子根 CMakeLists】而非 op_host/CMakeLists.txt：
17+# func.cmake:391 仅在 <op>/op_host/CMakeLists.txt 不存在时 add_subdirectory(<op>)，
18+# 二者【只能存在其一】，否则 add_modules_sources 被调用两次。
19+# DIR 传算子根路径是安全的：func.cmake:513 会在 SOURCE_DIR 不以 /op_host 结尾时自动补上，
20+# OPDEF_SRCS 仍正确 glob 到 <op>/op_host/sgd_def*.cpp。
21+ 
22+# 设置算子定义时支持的芯片类型
23+set(SUPPORT_COMPUTE_UNIT "ascend950")
24+# 设置每种芯片类型对应的 tiling 文件目录，即采用 op_host 目录下哪个文件夹下的 tiling 文件编译
25+set(SUPPORT_TILING_DIR "arch35")
26+ 
27+# ACLNNTYPE aclnn_exclude：把 sgd_def.cpp 归入 aclnnExc 桶 —— 该桶【不触发 aclnn 接口自动生成】
28+# （opbuild.cmake:57 分支只收集手写的 op_api/aclnn_*.h，本算子已无该目录，故不产出任何 aclnn 头）。
29+# 算子定义本身照常注册，GE 图通路不受影响。
30+# ⛔ 不得改成 aclnn / aclnn_inner —— 那会让构建系统重新生成 aclnnSgd 接口。
31+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE sgd ACLNNTYPE aclnn_exclude COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
Aoptim/sgd/README.md+185-0
@@ -0,0 +1,185 @@
1+# SGD
2+ 
3+## 产品支持情况
4+ 
5+|产品 | 是否支持 |
6+|:-------------------------|:----------:|
7+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
8+| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
9+| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |
10+| <term>Atlas 200I/500 A2 推理产品</term> | √ |
11+| <term>Atlas 推理系列产品</term> | √ |
12+| <term>Atlas 训练系列产品</term> | √ |
13+ 
14+> 上表写的是SGD在各产品形态上的**可得性**，不是本次交付的架构范围。本仓的Ascend C实现只适配 <term>Ascend 950PR/Ascend 950DT</term>（`sgd_def.cpp`中仅`AddConfig("ascend950")`）；其余产品形态上的SGD由CANN内置的TBE实现提供，语义一致，但不由本算子承载。
15+ 
16+## 功能说明
17+ 
18+- **算子功能**：带动量的随机梯度下降（SGD）优化器更新算子，训练迭代中就地更新一组权重。
19+ 
20+- **计算公式**：
21+ 
22+ 记$d$为`dampening`、$wd$为`weightDecay`、$lr$为`learningRate[0]`、$m$为`momentum[0]`，逐元素计算：
23+ 
24+ **步骤一** 权重衰减（仅$wd \neq 0$时执行，否则$grad = gradient$）：
25+ 
26+ $$
27+ grad = gradient + parameters \times wd
28+ $$
29+ 
30+ **步骤二** 动量累积（**无条件执行**）：
31+ 
32+ $$
33+ accum_t = accum \times m + grad
34+ $$
35+ 
36+ **步骤三** 阻尼修正（仅$d \neq 0$时执行）。$stat$是**逐元素**的首步标记，取值1表示该元素处于首步、不施加阻尼：
37+ 
38+ $$
39+ accum_t = accum_t - grad \times (1 - stat) \times d
40+ $$
41+ 
42+ **步骤四** 权重更新（**无条件写出**）：
43+ 
44+ $$
45+ parameters_{out} =
46+ \begin{cases}
47+ parameters - (grad \times lr + accum_t \times m \times lr), & nesterov = true \\
48+ parameters - accum_t \times lr, & nesterov = false
49+ \end{cases}
50+ $$
51+ 
52+ **步骤五** 动量与标记回写，受$m \neq 0$掩码控制：
53+ 
54+ $$
55+ accum_{out}, stat_{out} =
56+ \begin{cases}
57+ accum_t,\ 0, & m \neq 0 \\
58+ \text{保持输入原值（不回写）}, & m = 0
59+ \end{cases}
60+ $$
61+ 
62+- **计算精度**：中间计算在float32域进行，结果按就近偶数舍入（round-half-to-even）回目标数据类型。`learningRate`与`momentum`与`parameters`同数据类型，故float16/bfloat16下这两个标量本身已被量化。
63+ 
64+## 参数说明
65+ 
66+<table style="table-layout: auto; width: 100%">
67+<thead>
68+ <tr>
69+ <th style="white-space: nowrap">参数名</th>
70+ <th style="white-space: nowrap">输入/输出/属性</th>
71+ <th style="white-space: nowrap">描述</th>
72+ <th style="white-space: nowrap">数据类型</th>
73+ <th style="white-space: nowrap">数据格式</th>
74+ </tr>
75+</thead>
76+<tbody>
77+ <tr>
78+ <td>parameters</td>
79+ <td>输入 / 输出（原地）</td>
80+ <td>待更新的权重。<b>无条件被改写</b>。维度数（rank）须在<code>[1, 8]</code>内；不支持空tensor。</td>
81+ <td>FLOAT、FLOAT16、BFLOAT16</td>
82+ <td>ND</td>
83+ </tr>
84+ <tr>
85+ <td>gradient</td>
86+ <td>输入</td>
87+ <td>梯度。shape与数据类型须与parameters一致。</td>
88+ <td>FLOAT、FLOAT16、BFLOAT16</td>
89+ <td>ND</td>
90+ </tr>
91+ <tr>
92+ <td>learning_rate</td>
93+ <td>输入</td>
94+ <td>学习率。<b>shape须为<code>[1]</code>或0维标量</b>，数据类型须与parameters一致。</td>
95+ <td>FLOAT、FLOAT16、BFLOAT16</td>
96+ <td>ND</td>
97+ </tr>
98+ <tr>
99+ <td>accum</td>
100+ <td>输入 / 输出（原地）</td>
101+ <td>动量累积量。<b>仅momentum ≠ 0时被改写</b>；momentum = 0时逐位保持原值。shape与数据类型须与parameters一致。</td>
102+ <td>FLOAT、FLOAT16、BFLOAT16</td>
103+ <td>ND</td>
104+ </tr>
105+ <tr>
106+ <td>momentum</td>
107+ <td>输入</td>
108+ <td>动量因子。<b>shape须为<code>[1]</code>或0维标量</b>，数据类型须与parameters一致。取值为0（含<code>-0.0</code>）时触发“不回写”语义。</td>
109+ <td>FLOAT、FLOAT16、BFLOAT16</td>
110+ <td>ND</td>
111+ </tr>
112+ <tr>
113+ <td>stat</td>
114+ <td>输入 / 输出（原地）</td>
115+ <td>逐元素首步标记，取值1表示该元素处于首步、不施加阻尼。<b>仅momentum ≠ 0时被改写为0</b>；momentum = 0时逐位保持原值。shape与数据类型须与parameters一致。</td>
116+ <td>FLOAT、FLOAT16、BFLOAT16</td>
117+ <td>ND</td>
118+ </tr>
119+ <tr>
120+ <td>dampening</td>
121+ <td>属性</td>
122+ <td>动量阻尼系数，默认值0.0。<b>nesterov为true时必须为0</b>。</td>
123+ <td>FLOAT</td>
124+ <td>-</td>
125+ </tr>
126+ <tr>
127+ <td>weight_decay</td>
128+ <td>属性</td>
129+ <td>权重衰减系数，默认值0.0。<b>必须大于或等于0</b>。</td>
130+ <td>FLOAT</td>
131+ <td>-</td>
132+ </tr>
133+ <tr>
134+ <td>nesterov</td>
135+ <td>属性</td>
136+ <td>是否启用Nesterov动量，默认值false。</td>
137+ <td>BOOL</td>
138+ <td>-</td>
139+ </tr>
140+ <tr>
141+ <td>parameters</td>
142+ <td>输出</td>
143+ <td>更新后的权重，与输入parameters为同一块内存。shape、数据类型、数据格式均与输入parameters一致。</td>
144+ <td>FLOAT、FLOAT16、BFLOAT16</td>
145+ <td>ND</td>
146+ </tr>
147+</tbody>
148+</table>
149+ 
150+## 约束说明
151+ 
152+- **三路原地回写，但图上仅声明1个输出**：算子实际就地更新`parameters`、`accum`、`stat`三个张量，而图原型只声明`parameters`一个输出，`accum`与`stat`通过覆写其输入内存返回。调用方必须把这三者都视为可写。此形态与CANN内置实现一致。
153+ 
154+- **momentum = 0时的回写语义**：`momentum`为0（含`-0.0`）时，`accum`与`stat` **完全不被写入**，逐位保持输入原值（包括NaN的具体位模式、`±inf`、`-0.0`）；`parameters`不受该掩码影响，任何`momentum`取值下都照常计算并写出。`momentum`为极小非零值（如`1e-8`、`1e-30`）时按非零处理，正常回写。
155+ 
156+- **⚠️ 从PyTorch迁移的差异告警**：本算子`momentum = 0`时的“不回写”方向与`torch.optim.SGD`一致（PyTorch在`momentum == 0`时整块跳过动量更新）。但 **PyTorch的`dampening`施加在该判断之内，本算子（与CANN内置实现一致）施加在判断之外**。因此当`momentum = 0`且`dampening > 0`且`stat = 0`时，`parameters`的更新量与PyTorch相差$(1 - dampening)$倍。仅当`dampening = 0`或`stat = 1`时两者一致。从PyTorch迁移的调用方须感知此差异。
157+ 
158+- **rank与空tensor**：`parameters`的维度数须在`[1, 8]`内，**0维标量被拒绝**。**不支持空tensor** —— 任意一轴或多轴为0均判为非法并返回错误码，不存在“空进空出”语义（`accum`/`stat`的原地回写在元素数为0时无定义）。
159+ 
160+- **属性取值**：`nesterov = true`时`dampening`必须为0；`weight_decay`必须大于或等于0。违反者返回参数非法错误码。
161+ 
162+- **inf/NaN**：按IEEE 754语义传播，不做钳制或特判。特别地，`accum`含`±inf`且`momentum = 0`时，$accum \times momentum$产生的NaN会按IEEE语义传播进`parameters`。
163+ 
164+- **确定性**：输出逐位可复现。算子为纯逐元素计算、无跨元素累加，多核切分不改变任一元素的计算顺序。
165+ 
166+- **张量连续性**：所有输入须为连续张量。本算子不提供aclnn接口，无接口层做转连续/回填，非连续视图由调用方（GE图编译期）负责处理。
167+ 
168+## 调用说明
169+ 
170+> **不提供aclnn单算子接口。** SGD在CANN上游本就是纯图模式算子：CANN 9.1.0的
171+> `include/aclnnop/`下无`aclnn_sgd.h`（只有语义不同的`aclnn_fused_sgd.h`），
172+> `libopapi.so`未导出`aclnnSgd*`符号，canndev全仓亦无`aclnn_sgd`定义。
173+> 本算子与之对齐，只支持GE图模式下发。
174+ 
175+| 调用方式 | 样例代码 | 说明 |
176+| ---------------- | --------------------------- | --------------------------------------------------- |
177+| 图模式 | [test_geir_sgd.cpp](examples/test_geir_sgd.cpp) | 通过GE图方式调用SGD算子。 |
178+ 
179+## 参考资源
180+ 
181+- [《Ascend C算子开发》](https://hiascend.com/document/redirect/CannCommunityOpdevAscendC)：算子开发的概念原理与编程模型。
182+- [算子列表](../../docs/zh/op_list.md)：本项目全部算子的分类、调用方式与功能说明。
183+- [算子调用快速入门](../../docs/zh/invocation/quick_op_invocation.md)：算子样例的编译与运行步骤。
184+- [apply_momentum](../apply_momentum/README.md)：同族的动量优化器算子，与本算子结构最相近，可对照阅读。
185+- [fused_sgd](../fused_sgd/README.md)：语义**不同**的另一个算子（多TensorList融合、`dampening`施加在`momentum`分支之内、无`stat`）。名称相近，请勿混用。
Aoptim/sgd/examples/test_geir_sgd.cpp+342-0
@@ -0,0 +1,342 @@
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 test_geir_sgd.cpp
13+ * \brief SGD 的 GE 图（GEIR）通路验证样例。
14+ *
15+ * 用法：
16+ * ./test_geir_sgd # momentum = 0.9，走 accum/stat 回写分支
17+ * ./test_geir_sgd 0 # momentum = 0 ，走掩码分支（accum/stat 逐位保持输入原值）
18+ *
19+ * ⚠️ GEInitialize 每进程只能调用一次，故两个分支必须【分两次独立运行】，不要在同一进程里跑两遍。
20+ * ⚠️ GE SelectEngine 不认 ASCEND_CUSTOM_OPP_PATH —— 本样例必须在 .run 包【正式 install】后运行，
21+ * 仅靠隔离 vendor 会报 "Cannot find engine"。
22+ *
23+ * 结构照 optim/apply_momentum/examples/test_geir_apply_momentum.cpp，改动点：
24+ * ① 算子换为 op::SGD，输入 6 路（顺序按 REG_OP(SGD)）、属性 3 个；
25+ * ② ADD_INPUT 宏增加 inputValue 形参 —— 原版硬编码填 2，而本算子需要给 momentum 单独设 0；
26+ * ③ dtype 默认改 DT_FLOAT（便于手工核对期望值）。
27+ */
28+ 
29+#include <iostream>
30+#include <fstream>
31+#include <string.h>
32+#include <stdint.h>
33+#include <vector>
34+#include <string>
35+#include <map>
36+#include "assert.h"
37+#include "graph.h"
38+#include "types.h"
39+#include "tensor.h"
40+#include "ge_error_codes.h"
41+#include "ge_api_types.h"
42+#include "ge_api.h"
43+#include "array_ops.h"
44+#include "ge_ir_build.h"
45+#include "../op_graph/sgd_proto.h"
46+ 
47+#define FAILED -1
48+#define SUCCESS 0
49+ 
50+using namespace ge;
51+using std::map;
52+using std::string;
53+using std::vector;
54+#define ADD_INPUT(intputIndex, intputName, intputDtype, inputShape, inputValue) \
55+ vector<int64_t> placeholder##intputIndex##_shape = inputShape; \
56+ auto placeholder##intputIndex = op::Data("placeholder" + intputIndex).set_attr_index(0); \
57+ TensorDesc placeholder##intputIndex##_desc = TensorDesc(ge::Shape(placeholder##intputIndex##_shape), FORMAT_ND, \
58+ intputDtype); \
59+ placeholder##intputIndex##_desc.SetPlacement(ge::kPlacementHost); \
60+ placeholder##intputIndex##_desc.SetFormat(FORMAT_ND); \
61+ Tensor tensor_placeholder##intputIndex; \
62+ ret = GenOnesDataFloat32(placeholder##intputIndex##_shape, tensor_placeholder##intputIndex, \
63+ placeholder##intputIndex##_desc, inputValue); \
64+ if (ret != SUCCESS) { \
65+ printf("%s - ERROR - [XIR]: Generate input data failed\n", GetTime().c_str()); \
66+ return FAILED; \
67+ } \
68+ placeholder##intputIndex.update_input_desc_x(placeholder##intputIndex##_desc); \
69+ placeholder##intputIndex.update_output_desc_y(placeholder##intputIndex##_desc); \
70+ input.push_back(tensor_placeholder##intputIndex); \
71+ graph.AddOp(placeholder##intputIndex); \
72+ sgd1.set_input_##intputName(placeholder##intputIndex); \
73+ inputs.push_back(placeholder##intputIndex)
74+ 
75+#define ADD_INPUT_ATTR(attrName, attrValue) sgd1.set_attr_##attrName(attrValue)
76+ 
77+#define ADD_CONST_INPUT(intputIndex, intputName, intputDtype, inputShape) \
78+ vector<int64_t> placeholder##intputIndex##_shape = inputShape; \
79+ auto placeholder##intputIndex = op::Const("placeholder" + intputIndex); \
80+ TensorDesc placeholder##intputIndex##_desc = TensorDesc(ge::Shape(placeholder##intputIndex##_shape), FORMAT_ND, \
81+ intputDtype); \
82+ placeholder##intputIndex##_desc.SetPlacement(ge::kPlacementHost); \
83+ placeholder##intputIndex##_desc.SetFormat(FORMAT_ND); \
84+ Tensor tensor_placeholder##intputIndex; \
85+ ret = GenOnesData(placeholder##intputIndex##_shape, tensor_placeholder##intputIndex, \
86+ placeholder##intputIndex##_desc, intputDtype, 2); \
87+ if (ret != SUCCESS) { \
88+ printf("%s - ERROR - [XIR]: Generate input data failed\n", GetTime().c_str()); \
89+ return FAILED; \
90+ } \
91+ placeholder##intputIndex.SetAttr("value", tensor_placeholder##intputIndex); \
92+ placeholder##intputIndex.update_output_desc_y(placeholder##intputIndex##_desc); \
93+ graph.AddOp(placeholder##intputIndex); \
94+ sgd1.set_input_##intputName(placeholder##intputIndex); \
95+ sgd1.update_input_desc_##intputName(placeholder##intputIndex##_desc); \
96+ inputs.push_back(placeholder##intputIndex);
97+ 
98+#define ADD_OUTPUT(outputIndex, outputName, outputDtype, outputShape) \
99+ TensorDesc outputName##outputIndex##_desc_ = TensorDesc(ge::Shape(outputShape), FORMAT_ND, outputDtype); \
100+ sgd1.update_output_desc_##outputName(outputName##outputIndex##_desc)
101+ 
102+#define LOG_PRINT(message, ...) \
103+ do { \
104+ printf(message, ##__VA_ARGS__); \
105+ } while (0)
106+ 
107+string GetTime()
108+{
109+ time_t timep;
110+ time(&timep);
111+ char tmp[64];
112+ strftime(tmp, sizeof(tmp), "%Y-%m-%d %H:%M:%S,000", localtime(&timep));
113+ return tmp;
114+}
115+ 
116+uint32_t GetDataTypeSize(DataType dt)
117+{
118+ uint32_t dilation = 1;
119+ uint32_t oneByte = 1;
120+ uint32_t twoByte = 2;
121+ uint32_t fourByte = 4;
122+ uint32_t eightByte = 8;
123+ 
124+ if (dt == ge::DT_FLOAT) {
125+ dilation = fourByte;
126+ } else if (dt == ge::DT_FLOAT16) {
127+ dilation = twoByte;
128+ } else if (dt == ge::DT_BF16) {
129+ dilation = twoByte;
130+ } else if (dt == ge::DT_INT16) {
131+ dilation = twoByte;
132+ } else if (dt == ge::DT_UINT16) {
133+ dilation = twoByte;
134+ } else if (dt == ge::DT_INT32) {
135+ dilation = fourByte;
136+ } else if (dt == ge::DT_UINT32) {
137+ dilation = fourByte;
138+ } else if (dt == ge::DT_INT64) {
139+ dilation = eightByte;
140+ } else if (dt == ge::DT_UINT64) {
141+ dilation = eightByte;
142+ } else if (dt == ge::DT_INT8) {
143+ dilation = oneByte;
144+ }
145+ return dilation;
146+}
147+ 
148+int32_t GenOnesDataFloat32(vector<int64_t> shapes, Tensor& input_tensor, TensorDesc& input_tensor_desc, float value)
149+{
150+ input_tensor_desc.SetRealDimCnt(shapes.size());
151+ size_t size = 1;
152+ for (uint32_t i = 0; i < shapes.size(); i++) {
153+ size *= shapes[i];
154+ }
155+ uint32_t byteSizeFloat32 = 4;
156+ uint32_t data_len = size * byteSizeFloat32;
157+ float* pData = new (std::nothrow) float[size];
158+ 
159+ for (size_t i = 0; i < size; ++i) {
160+ *(pData + i) = value;
161+ }
162+ input_tensor = Tensor(input_tensor_desc, (uint8_t*)pData, data_len);
163+ return SUCCESS;
164+}
165+ 
166+int32_t GenOnesData(vector<int64_t> shapes, Tensor& input_tensor, TensorDesc& input_tensor_desc, DataType data_type,
167+ int value)
168+{
169+ input_tensor_desc.SetRealDimCnt(shapes.size());
170+ size_t size = 1;
171+ for (uint32_t i = 0; i < shapes.size(); i++) {
172+ size *= shapes[i];
173+ }
174+ uint32_t data_len = size * GetDataTypeSize(data_type);
175+ int32_t* pData = new (std::nothrow) int32_t[data_len];
176+ for (size_t i = 0; i < size; ++i) {
177+ *(pData + i) = value;
178+ }
179+ input_tensor = Tensor(input_tensor_desc, reinterpret_cast<uint8_t*>(pData), data_len);
180+ return SUCCESS;
181+}
182+ 
183+int32_t WriteDataToFile(string bin_file, uint64_t data_size, uint8_t* inputData)
184+{
185+ FILE* fp = fopen(bin_file.c_str(), "w");
186+ fwrite(inputData, sizeof(uint8_t), data_size, fp);
187+ fclose(fp);
188+ return SUCCESS;
189+}
190+ 
191+int CreateOppInGraph(DataType inDtype, float momentumValue, std::vector<ge::Tensor>& input,
192+ std::vector<Operator>& inputs, std::vector<Operator>& outputs, Graph& graph)
193+{
194+ Status ret = SUCCESS;
195+ auto sgd1 = op::SGD("sgd1");
196+ // 输入顺序严格按图原型 REG_OP(SGD)：parameters / gradient / learning_rate / accum / momentum / stat
197+ std::vector<int64_t> xShape = {2, 30};
198+ ADD_INPUT(1, parameters, inDtype, xShape, 1.0f);
199+ ADD_INPUT(2, gradient, inDtype, xShape, 0.5f);
200+ ADD_INPUT(3, learning_rate, inDtype, std::vector<int64_t>{1}, 0.1f);
201+ ADD_INPUT(4, accum, inDtype, xShape, 3.5f);
202+ // ★ momentum 由命令行给定：非 0 → accum/stat 回写分支；0 → 掩码分支（accum/stat 逐位保持）。
203+ // GEIR 每个场景必须独立进程（GEInitialize 只能调用一次），故两个分支分两次运行。
204+ ADD_INPUT(5, momentum, inDtype, std::vector<int64_t>{1}, momentumValue);
205+ ADD_INPUT(6, stat, inDtype, xShape, 1.0f);
206+ 
207+ // 属性：nesterov 为 true 时 dampening 必须为 0；weight_decay 必须 >= 0
208+ ADD_INPUT_ATTR(dampening, 0.0f);
209+ ADD_INPUT_ATTR(weight_decay, 0.0f);
210+ ADD_INPUT_ATTR(nesterov, false);
211+ 
212+ outputs.push_back(sgd1);
213+ return SUCCESS;
214+}
215+ 
216+int main(int argc, char* argv[])
217+{
218+ const char* graph_name = "tc_ge_irrun_test";
219+ Graph graph(graph_name);
220+ std::vector<ge::Tensor> input;
221+ 
222+ printf("%s - INFO - [XIR]: Start to initialize ge using ge global options\n", GetTime().c_str());
223+ std::map<AscendString, AscendString> global_options = {{"ge.exec.deviceId", "0"}, {"ge.graphRunMode", "1"}};
224+ Status ret = ge::GEInitialize(global_options);
225+ if (ret != SUCCESS) {
226+ printf("%s - INFO - [XIR]: Initialize ge using ge global options failed\n", GetTime().c_str());
227+ return FAILED;
228+ }
229+ printf("%s - INFO - [XIR]: Initialize ge using ge global options success\n", GetTime().c_str());
230+ 
231+ std::vector<Operator> inputs{};
232+ std::vector<Operator> outputs{};
233+ 
234+ // 用法：test_geir_sgd [momentum]
235+ // momentum 缺省 0.9（回写分支）；传 0 走掩码分支（accum/stat 不回写）。
236+ float momentumValue = 0.9f;
237+ if (argc > 1) {
238+ momentumValue = static_cast<float>(atof(argv[1]));
239+ }
240+ printf("%s - INFO - [XIR]: momentum = %f (%s)\n", GetTime().c_str(), momentumValue,
241+ (momentumValue != 0.0f) ? "writeback branch" : "mask branch (accum/stat kept bitwise)");
242+ 
243+ DataType inDtype = DT_FLOAT;
244+ std::cout << "inDtype = " << inDtype << std::endl;
245+ 
246+ ret = CreateOppInGraph(inDtype, momentumValue, input, inputs, outputs, graph);
247+ if (ret != SUCCESS) {
248+ printf("%s - ERROR - [XIR]: Create ir session using build options failed\n", GetTime().c_str());
249+ return FAILED;
250+ }
251+ 
252+ if (!inputs.empty() && !outputs.empty()) {
253+ graph.SetInputs(inputs).SetOutputs(outputs);
254+ }
255+ 
256+ std::map<AscendString, AscendString> build_options = {
257+ 
258+ };
259+ printf("%s - INFO - [XIR]: Start to create ir session using build options\n", GetTime().c_str());
260+ ge::Session* session = new Session(build_options);
261+ 
262+ if (session == nullptr) {
263+ printf("%s - ERROR - [XIR]: Create ir session using build options failed\n", GetTime().c_str());
264+ return FAILED;
265+ }
266+ printf("%s - INFO - [XIR]: Create ir session using build options success\n", GetTime().c_str());
267+ printf("%s - INFO - [XIR]: Start to add compute graph to ir session\n", GetTime().c_str());
268+ 
269+ std::map<AscendString, AscendString> graph_options = {
270+ 
271+ };
272+ uint32_t graph_id = 0;
273+ ret = session->AddGraph(graph_id, graph, graph_options);
274+ 
275+ printf("%s - INFO - [XIR]: Session add ir compute graph to ir session success\n", GetTime().c_str());
276+ printf("%s - INFO - [XIR]: dump graph to txt\n", GetTime().c_str());
277+ std::string file_path = "./dump";
278+ aclgrphDumpGraph(graph, file_path.c_str(), file_path.length());
279+ printf("%s - INFO - [XIR]: Start to run ir compute graph\n", GetTime().c_str());
280+ std::vector<ge::Tensor> output;
281+ ret = session->RunGraph(graph_id, input, output);
282+ if (ret != SUCCESS) {
283+ printf("%s - INFO - [XIR]: Run graph failed\n", GetTime().c_str());
284+ delete session;
285+ GEFinalize();
286+ return FAILED;
287+ }
288+ printf("%s - INFO - [XIR]: Session run ir compute graph success\n", GetTime().c_str());
289+ 
290+ int input_num = input.size();
291+ for (int i = 0; i < input_num; i++) {
292+ std::cout << "input " << i << " dtype : " << input[i].GetTensorDesc().GetDataType() << std::endl;
293+ string input_file = "./tc_ge_irrun_test_0008_npu_input_" + std::to_string(i) + ".bin";
294+ uint8_t* input_data_i = input[i].GetData();
295+ int64_t input_shape = input[i].GetTensorDesc().GetShape().GetShapeSize();
296+ std::cout << "this is " << i << "th input, input shape size =" << input_shape << std::endl;
297+ uint32_t data_size = input_shape * GetDataTypeSize(input[i].GetTensorDesc().GetDataType());
298+ WriteDataToFile((const char*)input_file.c_str(), data_size, input_data_i);
299+ }
300+ 
301+ int output_num = output.size();
302+ for (int i = 0; i < output_num; i++) {
303+ std::cout << "output " << i << " dtype : " << output[i].GetTensorDesc().GetDataType() << std::endl;
304+ string output_file = "./tc_ge_irrun_test_0008_npu_output_" + std::to_string(i) + ".bin";
305+ uint8_t* output_data_i = output[i].GetData();
306+ int64_t output_shape = output[i].GetTensorDesc().GetShape().GetShapeSize();
307+ std::cout << "this is " << i << "th output, output shape size =" << output_shape << std::endl;
308+ uint32_t data_size = output_shape * GetDataTypeSize(output[i].GetTensorDesc().GetDataType());
309+ WriteDataToFile((const char*)output_file.c_str(), data_size, output_data_i);
310+ float* resultData = (float*)output_data_i;
311+ for (int64_t j = 0; j < output_shape; j++) {
312+ LOG_PRINT("result[%ld] is: %f\n", j, resultData[j]);
313+ }
314+ }
315+ 
316+ ge::AscendString error_msg = ge::GEGetErrorMsgV2();
317+ std::string error_str(error_msg.GetString());
318+ std::cout << "Error message: " << error_str << std::endl;
319+ ge::AscendString warning_msg = ge::GEGetWarningMsgV2();
320+ std::string warning_str(warning_msg.GetString());
321+ std::cout << "Warning message: " << warning_str << std::endl;
322+ // 在 GEFinalize() 之前 delete session，与上游样例
323+ // /workspace/ops-nn/activation/celu/examples/test_geir_celu.cpp:241-243 对齐。
324+ // 本文件原先只在 RunGraph 失败分支删 session、成功路径漏删（session 泄漏）。
325+ //
326+ // ⚠ 澄清：补这两行【不能】消除进程收尾时的 `corrupted size vs. prev_size` 核心转储。
327+ // 实测该崩溃在 GEFinalize() 内部发生，且用**未经改动的上游样例 test_geir_celu
328+ // 打内置算子 Celu** 可等价复现（同样的报错、同样的位置、同样 exit 134），
329+ // 与本算子无关，属环境/GE 侧收尾问题。
330+ // 崩溃点在全部结果打印完毕之后，不影响计算结果的正确性；但也【不得】因此把它
331+ // 当成"跑通了"——排障时先跑一遍 celu 对照，再判断是不是自己的问题。
332+ delete session;
333+ session = nullptr;
334+ printf("%s - INFO - [XIR]: Start to finalize ir graph session\n", GetTime().c_str());
335+ ret = ge::GEFinalize();
336+ if (ret != SUCCESS) {
337+ printf("%s - INFO - [XIR]: Finalize ir graph session failed\n", GetTime().c_str());
338+ return FAILED;
339+ }
340+ printf("%s - INFO - [XIR]: Finalize ir graph session success\n", GetTime().c_str());
341+ return SUCCESS;
342+}
Aoptim/sgd/op_graph/CMakeLists.txt+12-0
@@ -0,0 +1,12 @@
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+message(STATUS "=== Debug: start ops.optim.sgd.graph_plugin.CMakeLists.txt ")
12+add_graph_plugin_sources()
Aoptim/sgd/op_graph/sgd_proto.h+102-0
@@ -0,0 +1,102 @@
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 sgd_proto.h
13+ * \brief
14+ */
15+#ifndef OPS_NN_OPTIM_SGD_PROTO_H
16+#define OPS_NN_OPTIM_SGD_PROTO_H
17+ 
18+#include "graph/operator_reg.h"
19+#include "graph/types.h"
20+ 
21+namespace ge {
22+/**
23+ *@brief Updates "parameters" according to the SGD algorithm with momentum. \n
24+ * computing process:
25+ *@code{.c}
26+ * // d = dampening, wd = weight_decay, lr = learning_rate[0], m = momentum[0]
27+ * if (wd != 0) {
28+ * grad = gradient + parameters * wd
29+ * } else {
30+ * grad = gradient
31+ * }
32+ * accum_t = accum * m + grad // unconditional
33+ * if (d != 0) {
34+ * accum_t -= grad * (1 - stat) * d
35+ * }
36+ * if (nesterov) {
37+ * parameters -= grad * lr + accum_t * m * lr
38+ * } else {
39+ * parameters -= accum_t * lr
40+ * }
41+ * if (m != 0) { // writeback mask
42+ * accum = accum_t
43+ * stat = 0
44+ * } // otherwise accum and stat keep their input values
45+ *@endcode
46+ *
47+ *@par Inputs:
48+ *@li parameters: A mutable tensor of ND. Must be of dtype float16, float32 or bfloat16.
49+ * Specifying parameters to be updated. Should be from a Variable().
50+ *@li gradient: A tensor of ND. Must be of the same shape and dtype as "parameters".
51+ * Specifying the gradient.
52+ *@li learning_rate: A scalar. Must be of the same dtype as "parameters".
53+ * Specifying the learning rate.
54+ *@li accum: A mutable tensor of ND. Must be of the same shape and dtype as "parameters".
55+ * Specifying the momentum accumulation. Should be from a Variable().
56+ *@li momentum: A scalar. Must be of the same dtype as "parameters". Specifying the momentum.
57+ *@li stat: A mutable tensor of ND. Must be of the same shape and dtype as "parameters".
58+ * Per-element first-step flag: a value of 1 means "first step" and suppresses the
59+ * dampening correction for that element. Should be from a Variable().
60+ *
61+ *@par Attributes:
62+ *@li dampening: An optional float. Defaults to "0.0". Must be 0 when "nesterov" is true.
63+ *@li weight_decay: An optional float. Defaults to "0.0". Must be greater than or equal to 0.
64+ *@li nesterov: An optional bool. Defaults to "false". If "true", uses Nesterov momentum.
65+ *
66+ *@par Outputs:
67+ * parameters: A mutable tensor. Has the same shape, dtype and format as input "parameters".
68+ *
69+ *@attention Constraints:
70+ *@li Only one output is declared on the graph, but the operator updates THREE tensors
71+ * in place: "parameters", "accum" and "stat". "accum" and "stat" are returned by
72+ * overwriting their input memory and are therefore not visible as graph outputs.
73+ * Callers must treat them as mutable. This mirrors the 910B/910C behaviour
74+ * (the TBE implementation declares reuse=('accum', 'parameters', 'stat')).
75+ *@li When "momentum" is 0 (including -0.0), "accum" and "stat" are NOT written at all
76+ * and keep their input values bit-for-bit; "parameters" is still updated as usual.
77+ *@li The rank of "parameters" must be in the range [1, 8]; rank-0 (scalar) is rejected.
78+ * Empty tensors (any axis being 0) are rejected as well.
79+ *
80+ *@par Third-party framework compatibility
81+ * The writeback mask matches PyTorch's torch.optim.SGD, which skips the whole momentum
82+ * block when momentum == 0. Note however that PyTorch applies "dampening" INSIDE that
83+ * block while this operator (like the 910B/910C implementation) applies it OUTSIDE.
84+ * Consequently, when momentum == 0 && dampening > 0 && stat == 0, "parameters" differs
85+ * from PyTorch by a factor of (1 - dampening).
86+ */
87+ 
88+REG_OP(SGD)
89+ .INPUT(parameters, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
90+ .INPUT(gradient, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
91+ .INPUT(learning_rate, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
92+ .INPUT(accum, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
93+ .INPUT(momentum, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
94+ .INPUT(stat, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
95+ .OUTPUT(parameters, TensorType({DT_FLOAT16, DT_FLOAT, DT_BF16}))
96+ .ATTR(dampening, Float, 0.0)
97+ .ATTR(weight_decay, Float, 0.0)
98+ .ATTR(nesterov, Bool, false)
99+ .OP_END_FACTORY_REG(SGD)
100+} // namespace ge
101+ 
102+#endif // OPS_NN_OPTIM_SGD_PROTO_H
Aoptim/sgd/op_host/arch35/sgd_tiling.cpp+316-0
@@ -0,0 +1,316 @@
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 sgd_tiling.cpp
13+ * \brief
14+ */
15+#include "sgd_tiling.h"
16+#include <graph/utils/type_utils.h>
17+#include "error_util.h"
18+#include "register/op_def_registry.h"
19+#include "register/tilingdata_base.h"
20+#include "tiling/platform/platform_ascendc.h"
21+#include "op_host/tiling_templates_registry.h"
22+ 
23+using namespace ge;
24+using namespace SgdOp;
25+ 
26+namespace optiling {
27+namespace {
28+constexpr size_t SYS_WORKSPACE = 16777216; // 16M
29+constexpr size_t MAX_DIM_NUM = 8; // 对齐 A2：rank 1 ~ 8（canndev kMaxDimNum = 8）
30+constexpr size_t MIN_DIM_NUM = 1;
31+ 
32+constexpr int32_t IDX_PARAMETERS = 0;
33+constexpr int32_t IDX_GRADIENT = 1;
34+constexpr int32_t IDX_LEARNING_RATE = 2;
35+constexpr int32_t IDX_ACCUM = 3;
36+constexpr int32_t IDX_MOMENTUM = 4;
37+constexpr int32_t IDX_STAT = 5;
38+ 
39+constexpr size_t ATTR_IDX_DAMPENING = 0;
40+constexpr size_t ATTR_IDX_WEIGHT_DECAY = 1;
41+constexpr size_t ATTR_IDX_NESTEROV = 2;
42+ 
43+// 与 parameters 严格同形同 dtype 的大张量
44+const std::map<int32_t, std::string> TENSOR_INDEX_LIST = {
45+ {IDX_GRADIENT, "gradient"}, {IDX_ACCUM, "accum"}, {IDX_STAT, "stat"}};
46+// shape 必须为 [1]（或 0D 标量）且与 parameters 同 dtype 的标量张量
47+const std::map<int32_t, std::string> SCALAR_INDEX_LIST = {{IDX_LEARNING_RATE, "learning_rate"},
48+ {IDX_MOMENTUM, "momentum"}};
49+} // namespace
50+ 
51+ge::graphStatus SgdRegbaseTiling::GetAttr()
52+{
53+ auto attrs = tilingContext_->GetAttrs();
54+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrs);
55+ 
56+ const float* dampeningAttr = attrs->GetAttrPointer<float>(ATTR_IDX_DAMPENING);
57+ dampening_ = (dampeningAttr != nullptr) ? *dampeningAttr : 0.0f;
58+ const float* weightDecayAttr = attrs->GetAttrPointer<float>(ATTR_IDX_WEIGHT_DECAY);
59+ weightDecay_ = (weightDecayAttr != nullptr) ? *weightDecayAttr : 0.0f;
60+ const bool* nesterovAttr = attrs->GetAttrPointer<bool>(ATTR_IDX_NESTEROV);
61+ nesterov_ = (nesterovAttr != nullptr) ? *nesterovAttr : false;
62+ 
63+ // 对齐 A2 的两条属性语义校验（canndev nn_training_ops.cc:1932-1955）。
64+ // 走结构化上报宏（R6），非裸 OP_LOGE。
65+ OP_CHECK_IF(nesterov_ && dampening_ != 0.0f,
66+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "dampening",
67+ std::to_string(dampening_).c_str(),
68+ "attr dampening must be 0 when attr nesterov is true"),
69+ return ge::GRAPH_FAILED);
70+ OP_CHECK_IF(weightDecay_ < 0.0f,
71+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "weight_decay",
72+ std::to_string(weightDecay_).c_str(),
73+ "attr weight_decay must be more than or equal to 0"),
74+ return ge::GRAPH_FAILED);
75+ 
76+ // 属性"是否为 0"固化为编译期分支（spec numerical_stability.skip_zero_branches）：
77+ // 条件不成立时【真正跳过】该子图，而不是乘 0 —— 0 * inf = NaN 会污染结果。
78+ useNesterovKey_ = nesterov_ ? 1U : 0U;
79+ hasWeightDecayKey_ = (weightDecay_ != 0.0f) ? 1U : 0U;
80+ hasDampeningKey_ = (dampening_ != 0.0f) ? 1U : 0U;
81+ return ge::GRAPH_SUCCESS;
82+}
83+ 
84+ge::graphStatus SgdRegbaseTiling::CheckRank(const gert::Shape& input0Shape)
85+{
86+ size_t dimNum = input0Shape.GetDimNum();
87+ OP_CHECK_IF(dimNum < MIN_DIM_NUM || dimNum > MAX_DIM_NUM,
88+ OP_LOGE_FOR_INVALID_SHAPEDIM(tilingContext_->GetNodeName(), "parameters",
89+ std::to_string(dimNum).c_str(), "1 ~ 8 dims"),
90+ return ge::GRAPH_FAILED);
91+ return ge::GRAPH_SUCCESS;
92+}
93+ 
94+ge::graphStatus SgdRegbaseTiling::CheckNotEmpty(const gert::Shape& input0Shape)
95+{
96+ // 空 Tensor（任意一轴或多轴为 0）按 spec error_codes 归 null_input，拒绝为非法。
97+ // 本算子无"空进空出"语义 —— accum / stat 的原地回写在 numel == 0 下无定义。
98+ OP_CHECK_IF(input0Shape.GetShapeSize() == 0,
99+ OP_LOGE_FOR_INVALID_SHAPESIZE_WITH_REASON(tilingContext_->GetNodeName(), "parameters",
100+ std::to_string(input0Shape.GetShapeSize()).c_str(),
101+ "empty tensor (any axis being 0) is not supported by SGD"),
102+ return ge::GRAPH_FAILED);
103+ return ge::GRAPH_SUCCESS;
104+}
105+ 
106+ge::graphStatus SgdRegbaseTiling::CheckScalarShape(int32_t inputIdx)
107+{
108+ auto inputShape = tilingContext_->GetInputShape(inputIdx);
109+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, inputShape);
110+ auto storageShape = inputShape->GetStorageShape();
111+ const std::string& paramName = SCALAR_INDEX_LIST.at(inputIdx);
112+ OP_CHECK_IF((!storageShape.IsScalar() && storageShape.GetShapeSize() != 1),
113+ OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(
114+ tilingContext_->GetNodeName(), paramName.c_str(), Ops::Base::ToString(storageShape).c_str(),
115+ (std::string("input ") + paramName + " must be scalar(0D) or have shape size 1").c_str()),
116+ return ge::GRAPH_FAILED);
117+ return ge::GRAPH_SUCCESS;
118+}
119+ 
120+ge::graphStatus SgdRegbaseTiling::CheckSameShape(int32_t inputIdx, const gert::Shape& input0Shape)
121+{
122+ auto inputShape = tilingContext_->GetInputShape(inputIdx);
123+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, inputShape);
124+ // 严格相等，不做广播 —— "可广播但不相等"（如 [2,3] vs [1,3]）同样判非法，
125+ // 对齐 spec 的 shape_mismatch 口径。
126+ if (inputShape->GetStorageShape() != input0Shape) {
127+ return ge::GRAPH_FAILED;
128+ }
129+ return ge::GRAPH_SUCCESS;
130+}
131+ 
132+ge::graphStatus SgdRegbaseTiling::CheckSameDtype(int32_t inputIdx, const ge::DataType& input0Dtype)
133+{
134+ auto inputDesc = tilingContext_->GetInputDesc(inputIdx);
135+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, inputDesc);
136+ if (inputDesc->GetDataType() != input0Dtype) {
137+ return ge::GRAPH_FAILED;
138+ }
139+ return ge::GRAPH_SUCCESS;
140+}
141+ 
142+ge::graphStatus SgdRegbaseTiling::CheckShapeAndType()
143+{
144+ auto inputShape = tilingContext_->GetInputShape(IDX_PARAMETERS);
145+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, inputShape);
146+ auto inputStorageShape = inputShape->GetStorageShape();
147+ 
148+ auto inputParamDesc = tilingContext_->GetInputDesc(IDX_PARAMETERS);
149+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, inputParamDesc);
150+ auto inputDtype = inputParamDesc->GetDataType();
151+ 
152+ // A2 只校验了 parameters 的 rank，其余 5 个输入的同形 / 同 dtype 完全未校验
153+ // （canndev CheckSgdDimension 仅看 parameters）。A5 侧补齐，见 01 §6.5。
154+ // 内层 CheckRank 已用 OP_LOGE_FOR_INVALID_SHAPEDIM 结构化上报，此处只作调用链留痕，
155+ // 不重复打 ERROR（R6 要求的是"校验必须走结构化宏"，不是"每层都再报一次"）。
156+ OP_CHECK_IF(CheckRank(inputStorageShape) != ge::GRAPH_SUCCESS, OP_LOGD(tilingContext_, "rank check failed"),
157+ return ge::GRAPH_FAILED);
158+ OP_CHECK_IF(CheckNotEmpty(inputStorageShape) != ge::GRAPH_SUCCESS,
159+ OP_LOGD(tilingContext_, "empty tensor check failed"), return ge::GRAPH_FAILED);
160+ 
161+ for (const auto& pair : SCALAR_INDEX_LIST) {
162+ OP_CHECK_IF(CheckScalarShape(pair.first) != ge::GRAPH_SUCCESS,
163+ OP_LOGD(tilingContext_, "scalar shape check failed for %s", pair.second.c_str()),
164+ return ge::GRAPH_FAILED);
165+ OP_CHECK_IF(
166+ CheckSameDtype(pair.first, inputDtype) != ge::GRAPH_SUCCESS,
167+ OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(
168+ tilingContext_->GetNodeName(), (std::string("parameters and ") + pair.second).c_str(),
169+ (ge::TypeUtils::DataTypeToSerialString(inputDtype) + " and " +
170+ ge::TypeUtils::DataTypeToSerialString(tilingContext_->GetInputDesc(pair.first)->GetDataType()))
171+ .c_str(),
172+ (std::string("the dtypes of input ") + pair.second + " and input parameters must be the same").c_str()),
173+ return ge::GRAPH_FAILED);
174+ }
175+ 
176+ for (const auto& pair : TENSOR_INDEX_LIST) {
177+ OP_CHECK_IF(
178+ CheckSameShape(pair.first, inputStorageShape) != ge::GRAPH_SUCCESS,
179+ OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(
180+ tilingContext_->GetNodeName(), (std::string("parameters and ") + pair.second).c_str(),
181+ (Ops::Base::ToString(inputStorageShape) + " and " +
182+ Ops::Base::ToString(tilingContext_->GetInputShape(pair.first)->GetStorageShape()))
183+ .c_str(),
184+ (std::string("the shapes of input ") + pair.second + " and input parameters must be the same").c_str()),
185+ return ge::GRAPH_FAILED);
186+ OP_CHECK_IF(
187+ CheckSameDtype(pair.first, inputDtype) != ge::GRAPH_SUCCESS,
188+ OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON(
189+ tilingContext_->GetNodeName(), (std::string("parameters and ") + pair.second).c_str(),
190+ (ge::TypeUtils::DataTypeToSerialString(inputDtype) + " and " +
191+ ge::TypeUtils::DataTypeToSerialString(tilingContext_->GetInputDesc(pair.first)->GetDataType()))
192+ .c_str(),
193+ (std::string("the dtypes of input ") + pair.second + " and input parameters must be the same").c_str()),
194+ return ge::GRAPH_FAILED);
195+ }
196+ return ge::GRAPH_SUCCESS;
197+}
198+ 
199+template <bool useNesterov, bool hasWeightDecay, bool hasDampening>
200+ge::graphStatus SgdRegbaseTiling::DoElewiseTilingByDtype(ElewiseBaseTiling& eleBaseTiling, ge::DataType dtype)
201+{
202+ // 一律用回写 DAG（doWriteback = true）反解 ubFormer：掩码 DAG 的 BufferNum 严格更小，
203+ // 按更保守的一套反解不会溢出；且 Host 收不到 Device 张量数据，本来也看不见 momentum 的值。
204+ if (dtype == ge::DT_FLOAT) {
205+ return eleBaseTiling
206+ .DoTiling<typename SgdDag<float, useNesterov, hasWeightDecay, hasDampening, true, float>::OpDag>(
207+ tiling_->elewiseTiling);
208+ }
209+ if (dtype == ge::DT_FLOAT16) {
210+ return eleBaseTiling
211+ .DoTiling<typename SgdDag<half, useNesterov, hasWeightDecay, hasDampening, true, float>::OpDag>(
212+ tiling_->elewiseTiling);
213+ }
214+ if (dtype == ge::DT_BF16) {
215+ return eleBaseTiling
216+ .DoTiling<typename SgdDag<bfloat16_t, useNesterov, hasWeightDecay, hasDampening, true, float>::OpDag>(
217+ tiling_->elewiseTiling);
218+ }
219+ OP_LOGE_FOR_INVALID_DTYPE(tilingContext_->GetNodeName(), "parameters",
220+ ge::TypeUtils::DataTypeToSerialString(dtype).c_str(), "float32, float16 or bfloat16");
221+ return ge::GRAPH_FAILED;
222+}
223+ 
224+ge::graphStatus SgdRegbaseTiling::DoElewiseTiling()
225+{
226+ ElewiseBaseTiling eleBaseTiling(tilingContext_);
227+ auto paramDesc = tilingContext_->GetInputDesc(IDX_PARAMETERS);
228+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, paramDesc);
229+ ge::DataType dtype = paramDesc->GetDataType();
230+ 
231+ OP_LOGI(tilingContext_->GetNodeName(),
232+ "Do elewise base tiling with nesterov=%d, hasWeightDecay=%d, hasDampening=%d", static_cast<int>(nesterov_),
233+ static_cast<int>(hasWeightDecayKey_), static_cast<int>(hasDampeningKey_));
234+ 
235+ // 6 个合法业务模板（K0~K5）。nesterov == true && dampening != 0 已在 GetAttr 拦下，
236+ // 故此处 nesterov 分支下 hasDampening 恒为 0，不生成该组合。
237+ if (useNesterovKey_ == 0) {
238+ if (hasWeightDecayKey_ == 0 && hasDampeningKey_ == 0) { // K0
239+ return DoElewiseTilingByDtype<false, false, false>(eleBaseTiling, dtype);
240+ }
241+ if (hasWeightDecayKey_ == 0 && hasDampeningKey_ == 1) { // K1
242+ return DoElewiseTilingByDtype<false, false, true>(eleBaseTiling, dtype);
243+ }
244+ if (hasWeightDecayKey_ == 1 && hasDampeningKey_ == 0) { // K2
245+ return DoElewiseTilingByDtype<false, true, false>(eleBaseTiling, dtype);
246+ }
247+ return DoElewiseTilingByDtype<false, true, true>(eleBaseTiling, dtype); // K3
248+ }
249+ if (hasWeightDecayKey_ == 1) { // K4
250+ return DoElewiseTilingByDtype<true, true, false>(eleBaseTiling, dtype);
251+ }
252+ return DoElewiseTilingByDtype<true, false, false>(eleBaseTiling, dtype); // K5
253+}
254+ 
255+ge::graphStatus SgdRegbaseTiling::SetTilingData()
256+{
257+ OP_LOGD(tilingContext_->GetNodeName(), "Enter SetTilingData");
258+ size_t* currentWorkspace = tilingContext_->GetWorkspaceSizes(1);
259+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, currentWorkspace);
260+ currentWorkspace[0] = SYS_WORKSPACE; // 算子不申请临时 GM，仅系统保留段
261+ 
262+ // 属性值下发给 Kernel，由 sch.SetVar 注入 DAG 的 Placeholder::Var 节点
263+ tiling_->dampening = dampening_;
264+ tiling_->weightDecay = weightDecay_;
265+ 
266+ tilingKey_ = GET_TPL_TILING_KEY(tiling_->elewiseTiling.scheMode, useNesterovKey_, hasWeightDecayKey_,
267+ hasDampeningKey_);
268+ OP_LOGI(tilingContext_->GetNodeName(),
269+ "scheMode=%ld, useNesterov=%ld, hasWeightDecay=%ld, hasDampening=%ld, tilingKey=%lu",
270+ tiling_->elewiseTiling.scheMode, useNesterovKey_, hasWeightDecayKey_, hasDampeningKey_, tilingKey_);
271+ tilingContext_->SetTilingKey(tilingKey_);
272+ 
273+ uint32_t blockDim = static_cast<uint32_t>(tiling_->elewiseTiling.blockNum);
274+ OP_CHECK_IF(blockDim <= 0, OP_LOGE(tilingContext_, "Get blockDim failed"), return ge::GRAPH_FAILED);
275+ tilingContext_->SetBlockDim(blockDim);
276+ return ge::GRAPH_SUCCESS;
277+}
278+ 
279+ge::graphStatus SgdRegbaseTiling::RunTiling()
280+{
281+ if (tilingContext_ == nullptr) {
282+ OP_LOGE("Sgd", "Get nullptr while obtaining tilingContext_.");
283+ return ge::GRAPH_FAILED;
284+ }
285+ // GetAttr / CheckShapeAndType 内部对非法输入已用 OP_LOGE_FOR_INVALID_* 结构化上报，
286+ // 此处只作调用链留痕、不重复打 ERROR —— 与上面 CheckRank / CheckNotEmpty 的处理一致。
287+ // R6 要求的是"校验必须走结构化宏"，不是"每层都再报一次"。
288+ OP_CHECK_IF(GetAttr() != ge::GRAPH_SUCCESS, OP_LOGD(tilingContext_, "Get attr failed."), return ge::GRAPH_FAILED);
289+ OP_CHECK_IF(CheckShapeAndType() != ge::GRAPH_SUCCESS, OP_LOGD(tilingContext_, "Shape and dtype check failed."),
290+ return ge::GRAPH_FAILED);
291+ tiling_ = tilingContext_->GetTilingData<SgdRegbaseTilingData>();
292+ OP_CHECK_IF((tiling_ == nullptr), OP_LOGE(tilingContext_, "Get SgdRegbaseTilingData from GE context failed"),
293+ return ge::GRAPH_FAILED);
294+ OP_CHECK_IF(DoElewiseTiling() != ge::GRAPH_SUCCESS, OP_LOGE(tilingContext_, "elewiseBaseTiling failed"),
295+ return ge::GRAPH_FAILED);
296+ return SetTilingData();
297+}
298+ 
299+ge::graphStatus Tiling4Sgd(gert::TilingContext* context)
300+{
301+ OP_LOGD(context, "Tiling4Sgd running begin");
302+ SgdRegbaseTiling tiling(context);
303+ return tiling.RunTiling();
304+}
305+ 
306+ge::graphStatus TilingPrepareForSgd(gert::TilingParseContext* context)
307+{
308+ OP_LOGD(context, "TilingPrepareForSgd running begin");
309+ return ge::GRAPH_SUCCESS;
310+}
311+ 
312+struct SgdCompileInfo {};
313+ 
314+IMPL_OP_OPTILING(SGD).Tiling(Tiling4Sgd).TilingParse<SgdCompileInfo>(TilingPrepareForSgd);
315+ 
316+} // namespace optiling
Aoptim/sgd/op_host/arch35/sgd_tiling.h+66-0
@@ -0,0 +1,66 @@
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 sgd_tiling.h
13+ * \brief
14+ */
15+#ifndef RUNTIME_V2_OP_IMPL_SGD_REGBASE_TILING_H_
16+#define RUNTIME_V2_OP_IMPL_SGD_REGBASE_TILING_H_
17+ 
18+#include "atvoss/elewise/elewise_tiling.h"
19+#include "register/tilingdata_base.h"
20+#include "../../op_kernel/arch35/sgd_dag.h"
21+#include "../../op_kernel/arch35/sgd_tiling_key.h"
22+#include "../../op_kernel/arch35/sgd_tiling_data.h"
23+ 
24+namespace optiling {
25+using namespace Ops::Base;
26+ 
27+class SgdRegbaseTiling {
28+public:
29+ explicit SgdRegbaseTiling(gert::TilingContext* context) : tilingContext_(context) {};
30+ 
31+ ge::graphStatus RunTiling();
32+ SgdRegbaseTilingData* tiling_ = nullptr;
33+ 
34+protected:
35+ ge::graphStatus GetAttr();
36+ ge::graphStatus CheckShapeAndType();
37+ ge::graphStatus CheckScalarShape(int32_t inputIdx);
38+ ge::graphStatus CheckSameShape(int32_t inputIdx, const gert::Shape& input0Shape);
39+ ge::graphStatus CheckSameDtype(int32_t inputIdx, const ge::DataType& input0Dtype);
40+ ge::graphStatus CheckRank(const gert::Shape& input0Shape);
41+ ge::graphStatus CheckNotEmpty(const gert::Shape& input0Shape);
42+ ge::graphStatus DoElewiseTiling();
43+ ge::graphStatus SetTilingData();
44+ 
45+ // DoElewiseTiling 的编译期分支展开：按 useNesterov / hasWeightDecay /
46+ // hasDampening 三个属性分支 × dtype 实例化对应的 OpDag 反解 ubFormer。
47+ // 【一律用回写 DAG（doWriteback = true）】—— 掩码 DAG 的 BufferNum 严格更小
48+ // （OutList::Size 3→1，GetLvl12Mte3Count()*2 由 6 降到 2，另有 OpZeroTsr 与两个
49+ // 降精度 Cast 被裁），按更保守的一套反解不会溢出。且 Host 看不见 momentum 的值，
50+ // 本来也无从按掩码分支反解。
51+ template <bool useNesterov, bool hasWeightDecay, bool hasDampening>
52+ ge::graphStatus DoElewiseTilingByDtype(ElewiseBaseTiling& eleBaseTiling, ge::DataType dtype);
53+ 
54+private:
55+ gert::TilingContext* tilingContext_ = nullptr;
56+ uint64_t tilingKey_ = 0;
57+ uint64_t useNesterovKey_ = 0;
58+ uint64_t hasWeightDecayKey_ = 0;
59+ uint64_t hasDampeningKey_ = 0;
60+ bool nesterov_ = false;
61+ float dampening_ = 0.0f;
62+ float weightDecay_ = 0.0f;
63+};
64+} // namespace optiling
65+ 
66+#endif // RUNTIME_V2_OP_IMPL_SGD_REGBASE_TILING_H_
Aoptim/sgd/op_host/config/ascend950/sgd_binary.json+308-0
@@ -0,0 +1,308 @@
1+{
2+ "op_type": "SGD",
3+ "op_list": [
4+ {
5+ "bin_filename": "SGD_b4f5a2k2g9d8h1o2e3e4b8f2b9m5g4n2",
6+ "inputs": [
7+ {
8+ "name": "parameters",
9+ "index": 0,
10+ "dtype": "float16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ],
16+ "format_match_mode": "FormatAgnostic"
17+ },
18+ {
19+ "name": "gradient",
20+ "index": 1,
21+ "dtype": "float16",
22+ "format": "ND",
23+ "paramType": "required",
24+ "shape": [
25+ -2
26+ ],
27+ "format_match_mode": "FormatAgnostic"
28+ },
29+ {
30+ "name": "learning_rate",
31+ "index": 2,
32+ "dtype": "float16",
33+ "format": "ND",
34+ "paramType": "required",
35+ "shape": [
36+ -2
37+ ],
38+ "format_match_mode": "FormatAgnostic"
39+ },
40+ {
41+ "name": "accum",
42+ "index": 3,
43+ "dtype": "float16",
44+ "format": "ND",
45+ "paramType": "required",
46+ "shape": [
47+ -2
48+ ],
49+ "format_match_mode": "FormatAgnostic"
50+ },
51+ {
52+ "name": "momentum",
53+ "index": 4,
54+ "dtype": "float16",
55+ "format": "ND",
56+ "paramType": "required",
57+ "shape": [
58+ -2
59+ ],
60+ "format_match_mode": "FormatAgnostic"
61+ },
62+ {
63+ "name": "stat",
64+ "index": 5,
65+ "dtype": "float16",
66+ "format": "ND",
67+ "paramType": "required",
68+ "shape": [
69+ -2
70+ ],
71+ "format_match_mode": "FormatAgnostic"
72+ }
73+ ],
74+ "outputs": [
75+ {
76+ "name": "parameters",
77+ "index": 0,
78+ "dtype": "float16",
79+ "format": "ND",
80+ "paramType": "required",
81+ "shape": [
82+ -2
83+ ],
84+ "format_match_mode": "FormatAgnostic"
85+ }
86+ ],
87+ "attrs": [
88+ {
89+ "name": "dampening",
90+ "dtype": "float",
91+ "value": 0.0
92+ },
93+ {
94+ "name": "weight_decay",
95+ "dtype": "float",
96+ "value": 0.0
97+ },
98+ {
99+ "name": "nesterov",
100+ "dtype": "bool",
101+ "value": false
102+ }
103+ ]
104+ },
105+ {
106+ "bin_filename": "SGD_e4h8j8n8h7m1n1k0n4f1b8p6e7i2b0b4",
107+ "inputs": [
108+ {
109+ "name": "parameters",
110+ "index": 0,
111+ "dtype": "float32",
112+ "format": "ND",
113+ "paramType": "required",
114+ "shape": [
115+ -2
116+ ],
117+ "format_match_mode": "FormatAgnostic"
118+ },
119+ {
120+ "name": "gradient",
121+ "index": 1,
122+ "dtype": "float32",
123+ "format": "ND",
124+ "paramType": "required",
125+ "shape": [
126+ -2
127+ ],
128+ "format_match_mode": "FormatAgnostic"
129+ },
130+ {
131+ "name": "learning_rate",
132+ "index": 2,
133+ "dtype": "float32",
134+ "format": "ND",
135+ "paramType": "required",
136+ "shape": [
137+ -2
138+ ],
139+ "format_match_mode": "FormatAgnostic"
140+ },
141+ {
142+ "name": "accum",
143+ "index": 3,
144+ "dtype": "float32",
145+ "format": "ND",
146+ "paramType": "required",
147+ "shape": [
148+ -2
149+ ],
150+ "format_match_mode": "FormatAgnostic"
151+ },
152+ {
153+ "name": "momentum",
154+ "index": 4,
155+ "dtype": "float32",
156+ "format": "ND",
157+ "paramType": "required",
158+ "shape": [
159+ -2
160+ ],
161+ "format_match_mode": "FormatAgnostic"
162+ },
163+ {
164+ "name": "stat",
165+ "index": 5,
166+ "dtype": "float32",
167+ "format": "ND",
168+ "paramType": "required",
169+ "shape": [
170+ -2
171+ ],
172+ "format_match_mode": "FormatAgnostic"
173+ }
174+ ],
175+ "outputs": [
176+ {
177+ "name": "parameters",
178+ "index": 0,
179+ "dtype": "float32",
180+ "format": "ND",
181+ "paramType": "required",
182+ "shape": [
183+ -2
184+ ],
185+ "format_match_mode": "FormatAgnostic"
186+ }
187+ ],
188+ "attrs": [
189+ {
190+ "name": "dampening",
191+ "dtype": "float",
192+ "value": 0.0
193+ },
194+ {
195+ "name": "weight_decay",
196+ "dtype": "float",
197+ "value": 0.0
198+ },
199+ {
200+ "name": "nesterov",
201+ "dtype": "bool",
202+ "value": false
203+ }
204+ ]
205+ },
206+ {
207+ "bin_filename": "SGD_l2l1h2j4p8p2p7p8f0f5l0a8b7n5i9g1",
208+ "inputs": [
209+ {
210+ "name": "parameters",
211+ "index": 0,
212+ "dtype": "bfloat16",
213+ "format": "ND",
214+ "paramType": "required",
215+ "shape": [
216+ -2
217+ ],
218+ "format_match_mode": "FormatAgnostic"
219+ },
220+ {
221+ "name": "gradient",
222+ "index": 1,
223+ "dtype": "bfloat16",
224+ "format": "ND",
225+ "paramType": "required",
226+ "shape": [
227+ -2
228+ ],
229+ "format_match_mode": "FormatAgnostic"
230+ },
231+ {
232+ "name": "learning_rate",
233+ "index": 2,
234+ "dtype": "bfloat16",
235+ "format": "ND",
236+ "paramType": "required",
237+ "shape": [
238+ -2
239+ ],
240+ "format_match_mode": "FormatAgnostic"
241+ },
242+ {
243+ "name": "accum",
244+ "index": 3,
245+ "dtype": "bfloat16",
246+ "format": "ND",
247+ "paramType": "required",
248+ "shape": [
249+ -2
250+ ],
251+ "format_match_mode": "FormatAgnostic"
252+ },
253+ {
254+ "name": "momentum",
255+ "index": 4,
256+ "dtype": "bfloat16",
257+ "format": "ND",
258+ "paramType": "required",
259+ "shape": [
260+ -2
261+ ],
262+ "format_match_mode": "FormatAgnostic"
263+ },
264+ {
265+ "name": "stat",
266+ "index": 5,
267+ "dtype": "bfloat16",
268+ "format": "ND",
269+ "paramType": "required",
270+ "shape": [
271+ -2
272+ ],
273+ "format_match_mode": "FormatAgnostic"
274+ }
275+ ],
276+ "outputs": [
277+ {
278+ "name": "parameters",
279+ "index": 0,
280+ "dtype": "bfloat16",
281+ "format": "ND",
282+ "paramType": "required",
283+ "shape": [
284+ -2
285+ ],
286+ "format_match_mode": "FormatAgnostic"
287+ }
288+ ],
289+ "attrs": [
290+ {
291+ "name": "dampening",
292+ "dtype": "float",
293+ "value": 0.0
294+ },
295+ {
296+ "name": "weight_decay",
297+ "dtype": "float",
298+ "value": 0.0
299+ },
300+ {
301+ "name": "nesterov",
302+ "dtype": "bool",
303+ "value": false
304+ }
305+ ]
306+ }
307+ ]
308+}
Aoptim/sgd/op_host/sgd_def.cpp+82-0
@@ -0,0 +1,82 @@
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 sgd_def.cpp
13+ * \brief SGD OpDef —— 仅适配 Ascend950（arch35 / DAV_3510 / regbase）
14+ *
15+ * 支持面对齐 910B/910C 基线（canndev aic-ascend910b-ops-info.ini 的 [SGD] 段），
16+ * 唯一收窄项是 format：A2 大张量支持 NC1HWC0 / NDC1HWC0 / ND / FRACTAL_Z /
17+ * FRACTAL_Z_3D 五种，本算子只做 ND。依据是 ops-nn 仓内 arch35 全族 optim 算子
18+ * （apply_momentum / apply_ftrl / apply_adam_w_v2 / apply_adamax /
19+ * apply_centered_rms_prop）的 def.cpp 一律只声明 ge::FORMAT_ND，无一例声明私有
20+ * format —— 跟随本仓 arch35 既定约定，非自行设计。已于 CP1 批准。
21+ *
22+ * R2（A5 不碰 A2）天然满足：ops-nn 仓内不存在 SGD 算子（optim/fused_sgd 是语义
23+ * 不同的另一个算子，且无 arch35 目录），本算子为净新增，只 AddConfig("ascend950")，
24+ * 不触碰任何 A2 配置（先例：apply_momentum、apply_adamax 亦只有 ascend950 config）。
25+ *
26+ * accum / stat 【不声明为 Output】—— 与 A2 形态一致（canndev proto
27+ * nn_training_ops.h:1431 亦只有一个 OUTPUT），两者靠覆写输入 GM 原地回写。
28+ * 该"proto 与实现不自洽"是本族常态，已在 README 显式声明。
29+ */
30+#include "register/op_def_registry.h"
31+ 
32+namespace ops {
33+ 
34+class SGD : public OpDef {
35+public:
36+ explicit SGD(const char* name) : OpDef(name)
37+ {
38+ this->Input("parameters")
39+ .ParamType(REQUIRED)
40+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
41+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
42+ this->Input("gradient")
43+ .ParamType(REQUIRED)
44+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
45+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
46+ this->Input("learning_rate")
47+ .ParamType(REQUIRED)
48+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
49+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
50+ this->Input("accum")
51+ .ParamType(REQUIRED)
52+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
53+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
54+ this->Input("momentum")
55+ .ParamType(REQUIRED)
56+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
57+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
58+ this->Input("stat")
59+ .ParamType(REQUIRED)
60+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
61+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
62+ this->Output("parameters")
63+ .ParamType(REQUIRED)
64+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT})
65+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
66+ // 属性默认值与 A2 逐字一致（aic-ascend910b-ops-info.ini 的 [SGD] 段）
67+ this->Attr("dampening").AttrType(OPTIONAL).Float(0.0);
68+ this->Attr("weight_decay").AttrType(OPTIONAL).Float(0.0);
69+ this->Attr("nesterov").AttrType(OPTIONAL).Bool(false);
70+ 
71+ OpAICoreConfig aicoreConfig;
72+ aicoreConfig.DynamicCompileStaticFlag(true)
73+ .DynamicRankSupportFlag(true) // 支持 -2（UNKNOWN_RANK）透传
74+ .DynamicShapeSupportFlag(true) // 支持 -1（UNKNOWN_DIM）
75+ .PrecisionReduceFlag(false) // 对齐 A2 precision_reduce.flag=false
76+ .ExtendCfgInfo("opInterface.value", "sgd");
77+ this->AICore().AddConfig("ascend950", aicoreConfig);
78+ }
79+};
80+ 
81+OP_ADD(SGD);
82+} // namespace ops
Aoptim/sgd/op_host/sgd_infershape.cpp+100-0
@@ -0,0 +1,100 @@
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 sgd_infershape.cpp
13+ * \brief SGD InferShape / InferDataType
14+ *
15+ * 行为对齐 A2（canndev nn_training_ops.cc:1906-1973）：
16+ * - parameters_out 的 shape / dtype 等于 parameters；
17+ * - rank 落在 1 ~ 8 之外即 GRAPH_FAILED（故 rank-0 标量被拒）；
18+ * - UNKNOWN_RANK(-2) 透传 —— GE 下 dims = {-2}（DimNum == 1），不触发 rank 拒绝；
19+ * - nesterov == true 时 dampening 必须为 0；weight_decay >= 0。
20+ *
21+ * 注：A2 把这两条属性校验放在 InferShape 里而非 Verify（其 SGDVerify 只调
22+ * CheckSgdDimension），属职责错位（见 01 §6.5）。A5 侧在 InferShape 与 Tiling
23+ * 两个独立入口都做，保证任一入口进来都拦得住。
24+ */
25+#include "log/log.h"
26+#include "register/op_impl_registry.h"
27+#include "op_host/infershape_elewise_util.h"
28+ 
29+using namespace ge;
30+ 
31+namespace ops {
32+namespace {
33+constexpr size_t PARAMETERS_INDEX = 0;
34+constexpr size_t MAX_DIM_NUM = 8;
35+constexpr size_t MIN_DIM_NUM = 1;
36+constexpr size_t ATTR_IDX_DAMPENING = 0;
37+constexpr size_t ATTR_IDX_WEIGHT_DECAY = 1;
38+constexpr size_t ATTR_IDX_NESTEROV = 2;
39+constexpr int64_t UNKNOWN_RANK_DIM = -2;
40+ 
41+// UNKNOWN_RANK 在 GE 下表现为 dims == {-2}，须透传而非按 rank 拒绝。
42+bool IsUnknownRank(const gert::Shape* shape) { return shape->GetDimNum() == 1 && shape->GetDim(0) == UNKNOWN_RANK_DIM; }
43+} // namespace
44+ 
45+static ge::graphStatus InferShapeForSgd(gert::InferShapeContext* context)
46+{
47+ OP_LOGD(context, "InferShapeForSgd begin.");
48+ 
49+ const auto* attrs = context->GetAttrs();
50+ OP_CHECK_IF(attrs == nullptr, OP_LOGE(context, "Get attrs failed."), return ge::GRAPH_FAILED);
51+ const float* dampening = attrs->GetAttrPointer<float>(ATTR_IDX_DAMPENING);
52+ const float* weightDecay = attrs->GetAttrPointer<float>(ATTR_IDX_WEIGHT_DECAY);
53+ const bool* nesterov = attrs->GetAttrPointer<bool>(ATTR_IDX_NESTEROV);
54+ const float dampeningVal = (dampening != nullptr) ? *dampening : 0.0f;
55+ const float weightDecayVal = (weightDecay != nullptr) ? *weightDecay : 0.0f;
56+ const bool nesterovVal = (nesterov != nullptr) ? *nesterov : false;
57+ 
58+ OP_CHECK_IF(
59+ nesterovVal && dampeningVal != 0.0f,
60+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "dampening", std::to_string(dampeningVal).c_str(),
61+ "attr dampening must be 0 when attr nesterov is true"),
62+ return ge::GRAPH_FAILED);
63+ OP_CHECK_IF(weightDecayVal < 0.0f,
64+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context->GetNodeName(), "weight_decay",
65+ std::to_string(weightDecayVal).c_str(),
66+ "attr weight_decay must be more than or equal to 0"),
67+ return ge::GRAPH_FAILED);
68+ 
69+ // 这里守的是「框架没把 shape 递过来」（返回 nullptr），属于内部失败，不是 parameters 取值非法，
70+ // 所以用裸 OP_LOGE 而非 OP_LOGE_FOR_INVALID_*：后者会走 REPORT_PREDEFINED_ERR_MSG 报 EZ0009/EZ0010，
71+ // 对用户提示「你的 shape 不合法」，而此刻并不知道用户给了什么。真正的 shape 取值校验在下面的 dimNum 分支。
72+ const gert::Shape* paramShape = context->GetInputShape(PARAMETERS_INDEX);
73+ OP_CHECK_IF(paramShape == nullptr, OP_LOGE(context, "Get input shape of parameters returns nullptr."),
74+ return ge::GRAPH_FAILED);
75+ if (!IsUnknownRank(paramShape)) {
76+ const size_t dimNum = paramShape->GetDimNum();
77+ OP_CHECK_IF(dimNum < MIN_DIM_NUM || dimNum > MAX_DIM_NUM,
78+ OP_LOGE_FOR_INVALID_SHAPEDIM(context->GetNodeName(), "parameters", std::to_string(dimNum).c_str(),
79+ "1 ~ 8 dims"),
80+ return ge::GRAPH_FAILED);
81+ }
82+ 
83+ ge::graphStatus ret = Ops::Base::InferShape4Elewise(context);
84+ OP_CHECK_IF(ret == ge::GRAPH_FAILED, OP_LOGE(context, "InferShapeForSgd failed."), return ge::GRAPH_FAILED);
85+ OP_LOGD(context, "InferShapeForSgd end.");
86+ return ret;
87+}
88+ 
89+static graphStatus InferDataTypeForSgd(gert::InferDataTypeContext* context)
90+{
91+ // 唯一图输出 parameters 的 dtype 等于输入 parameters。
92+ // accum / stat 不是图输出（靠覆写输入 GM 原地回写），故此处只有一路。
93+ context->SetOutputDataType(PARAMETERS_INDEX, context->GetInputDataType(PARAMETERS_INDEX));
94+ return GRAPH_SUCCESS;
95+}
96+ 
97+// 双挂：InferShape 与 InferDataType 都要注册。漏 InferDataType 会导致 GE 侧输出
98+// dtype 推导缺失/推错，是拷模板时的高频漏项（step3 GATE 明列）。
99+IMPL_OP_INFERSHAPE(SGD).InferShape(InferShapeForSgd).InferDataType(InferDataTypeForSgd);
100+} // namespace ops
Aoptim/sgd/op_kernel/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+add_kernel_sources(
12+ KERNEL_SRC arch35/sgd.cpp
13+ COMPUTE_UNITS ascend950
14+ AUTO_SYNC false
15+)
Aoptim/sgd/op_kernel/arch35/sgd.cpp+141-0
@@ -0,0 +1,141 @@
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 sgd.cpp
13+ * \brief sgd
14+ */
15+ 
16+#include <type_traits>
17+#include "kernel_operator.h"
18+#include "sgd_dag.h"
19+#include "sgd_tiling_key.h"
20+#include "sgd_tiling_data.h"
21+#include "atvoss/elewise/elewise_sch.h"
22+ 
23+using namespace AscendC;
24+using namespace SgdOp;
25+ 
26+namespace {
27+/**
28+ * 读 momentum GM 的第 0 个元素并升到 float32。
29+ *
30+ * fp32 / fp16：直接 GetValue(0) 后 static_cast<float>。
31+ * bf16 ：按位解码。理由是编译器限制，不是风格选择 ——
32+ * optim/apply_centered_rms_prop/op_kernel/arch35/apply_centered_rms_prop.h:184-185
33+ * 明确记录「BF16 fallback: bitwise conversion to avoid unsupported LLVM
34+ * bf16→fp32 scalar cast（"not support bf16 type cast"）」，而其 fp32 / fp16
35+ * 分支照常直接 cast。本函数与该先例逐行同形态。
36+ * 注意受限的只有【标量 cast】：GlobalTensor<bfloat16_t>::GetValue(0) 本身可用，
37+ * 不需要把 GM 指针 reinterpret_cast 成 __gm__ uint16_t*。
38+ *
39+ * 红线：位操作【仅用于解码】，判定发生在解码后的 float 上用 IEEE !=。
40+ * -0.0 的 bf16 位模式 0x8000 解码为 -0.0f，而 -0.0f != 0.0f 为 false
41+ * → 正确判为 0（走"不回写"分支），与 spec `m32 != np.float32(0.0)` 同口径。
42+ * 【不做】位模式比较、【不做】floor/ceil、【不降】fp16 —— canndev sgd.py:142-145
43+ * 那套会让 1e-8 这类极小值在不同芯片上行为相反。
44+ */
45+template <typename T>
46+__aicore__ inline float LoadMomentumScalarF32(GM_ADDR momentumGm)
47+{
48+ GlobalTensor<T> momentumGlobal;
49+ momentumGlobal.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(momentumGm));
50+ T raw = momentumGlobal.GetValue(0);
51+ 
52+ if constexpr (std::is_same<T, bfloat16_t>::value) {
53+ union {
54+ uint16_t u;
55+ bfloat16_t b;
56+ } src;
57+ union {
58+ uint32_t u;
59+ float f;
60+ } dst;
61+ src.b = raw;
62+ dst.u = static_cast<uint32_t>(src.u) << 16; // bf16 → fp32 位扩展，精确无损
63+ return dst.f;
64+ } else {
65+ return static_cast<float>(raw);
66+ }
67+}
68+} // namespace
69+ 
70+/**
71+ * SGD kernel 入口。
72+ *
73+ * 7 个业务 GM_ADDR，顺序与图原型一致。图上只声明 1 个 output（parameters）；
74+ * accum / stat 通过【覆写输入 GM】返回 —— sch.Init 的输出位直接填入 accum / stat
75+ * 这两个输入地址（先例 optim/apply_ftrl/op_kernel/arch35/apply_ftrl.cpp:33）。
76+ *
77+ * 运行期掩码分支：momentum == 0（含 -0.0）时选掩码 DAG，accum / stat 的 GM
78+ * 【不出现在输出位】→ 零写事务 → 逐位保持原值。两个分支共用同一个 TPipe。
79+ */
80+template <uint64_t schMode, uint64_t useNesterov, uint64_t hasWeightDecay, uint64_t hasDampening>
81+__global__ __aicore__ void sgd(GM_ADDR parameters, GM_ADDR gradient, GM_ADDR learning_rate, GM_ADDR accum,
82+ GM_ADDR momentum, GM_ADDR stat, GM_ADDR parameters_out, GM_ADDR workspace,
83+ GM_ADDR tiling)
84+{
85+ REGISTER_TILING_DEFAULT(SgdRegbaseTilingData);
86+ GET_TILING_DATA_WITH_STRUCT(SgdRegbaseTilingData, tilingData, tiling);
87+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
88+ 
89+ constexpr bool kUseNesterov = (static_cast<int>(useNesterov) == 1);
90+ constexpr bool kHasWeightDecay = (static_cast<int>(hasWeightDecay) == 1);
91+ constexpr bool kHasDampening = (static_cast<int>(hasDampening) == 1);
92+ 
93+ // Var 索引必须从 0 起连续无空洞（sch.SetVar 带 static_assert(index < Vars::Size)）：
94+ // weight_decay 恒占 0；dampening 在无 weight_decay 时让位到 0。
95+ constexpr int kVarIdxWeightDecay = SGD_VAR_IDX_WEIGHT_DECAY;
96+ constexpr int kVarIdxDampening = kHasWeightDecay ? 1 : 0;
97+ 
98+ using WritebackDag = SgdOp::SgdDag<DTYPE_PARAMETERS, kUseNesterov, kHasWeightDecay, kHasDampening, true>;
99+ using MaskedDag = SgdOp::SgdDag<DTYPE_PARAMETERS, kUseNesterov, kHasWeightDecay, kHasDampening, false>;
100+ 
101+ const float momentumScalar = LoadMomentumScalarF32<DTYPE_PARAMETERS>(momentum);
102+ 
103+ TPipe pipe; // 单一 TPipe，两个分支共用
104+ 
105+ if (momentumScalar != 0.0f) {
106+ // 回写分支：三路输出。输出位 2、3 填的是【输入】accum / stat 的 GM 地址。
107+ ElementwiseSch<schMode, typename WritebackDag::OpDag> sch(&(tilingData.elewiseTiling), &pipe);
108+ if constexpr (kHasWeightDecay) {
109+ sch.template SetVar<float, kVarIdxWeightDecay>(tilingData.weightDecay);
110+ }
111+ if constexpr (kHasDampening) {
112+ sch.template SetVar<float, kVarIdxDampening>(tilingData.dampening);
113+ }
114+ if constexpr (kHasDampening) {
115+ sch.Init(parameters, gradient, learning_rate, accum, momentum, stat, parameters_out, accum, stat);
116+ } else {
117+ // hasDampening == 0：不读 stat，输入退化为 5 路（In5 不在 DAG 闭包内）。
118+ // 但 stat 仍要被回写为 0，故输出位保留 3 个。
119+ sch.Init(parameters, gradient, learning_rate, accum, momentum, parameters_out, accum, stat);
120+ }
121+ sch.Process();
122+ } else {
123+ // 掩码分支：只回写 parameters。accum / stat 的 GM【从不出现在输出位】。
124+ // ⛔ 输入侧完全不变 —— accum_t 是 parameters_out 的上游必须照常算，
125+ // 且 accum 含 ±inf 时 0 * inf = NaN 须按 IEEE 传播进 parameters_out。
126+ ElementwiseSch<schMode, typename MaskedDag::OpDag> sch(&(tilingData.elewiseTiling), &pipe);
127+ if constexpr (kHasWeightDecay) {
128+ sch.template SetVar<float, kVarIdxWeightDecay>(tilingData.weightDecay);
129+ }
130+ if constexpr (kHasDampening) {
131+ sch.template SetVar<float, kVarIdxDampening>(tilingData.dampening);
132+ }
133+ if constexpr (kHasDampening) {
134+ sch.Init(parameters, gradient, learning_rate, accum, momentum, stat, parameters_out);
135+ } else {
136+ sch.Init(parameters, gradient, learning_rate, accum, momentum, parameters_out);
137+ }
138+ sch.Process();
139+ }
140+ return;
141+}
Aoptim/sgd/op_kernel/arch35/sgd_dag.h+166-0
@@ -0,0 +1,166 @@
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 sgd_dag.h
13+ * \brief SGD 的 ATVOSS DAG 定义（arch35 / Ascend950 / regbase）
14+ */
15+ 
16+#ifndef SGD_DAG_H
17+#define SGD_DAG_H
18+ 
19+#include <type_traits>
20+#include "atvoss/util/dag.h"
21+#include "atvoss/util/vec.h"
22+#include "atvoss/util/placeholder.h"
23+ 
24+namespace SgdOp {
25+using namespace Ops::Base;
26+ 
27+// Var 索引：DAG 中 Var 节点随分支存废，索引必须从 0 起连续无空洞
28+// （sch.SetVar<U, index> 带 static_assert(index < ElemDag::Vars::Size)）。
29+constexpr int SGD_VAR_IDX_WEIGHT_DECAY = 0;
30+ 
31+/**
32+ * SGD DAG。
33+ *
34+ * 计算语义（逐元素，T = float 域；d = dampening、wd = weight_decay、
35+ * lr = learning_rate[0]、m = momentum[0]）：
36+ * 1. grad = gradient + parameters * wd 仅 hasWeightDecay
37+ * 2. accum_t = accum * m + grad 【无条件】
38+ * 3. accum_t -= grad * ((1 - stat) * d) 仅 hasDampening
39+ * 4. parameters_out = nesterov ? p - (grad + accum_t*m) * lr
40+ * : p - accum_t*lr 【无条件写出】
41+ * 5. accum_out = accum_t、stat_out = 0 仅 doWriteback（即 m != 0）
42+ *
43+ * 模板参数 doWriteback 承载 `momentum != 0` 掩码：
44+ * - true —— 回写 DAG，Outputs 三路，sch.Init 输出位 3 个；
45+ * - false —— 掩码 DAG，Outputs 只剩 OpCopyOutParam，accum / stat 的 GM
46+ * 【从不出现在 sch.Init 的输出位】→ 框架 InitOutputArgs 不为其构造
47+ * outGm[]、CopyOut 路径根本不存在 → 零写事务 → 逐位保持输入原值
48+ * 对任意位模式（NaN payload / ±inf / -0.0）平凡成立。
49+ *
50+ * ⛔ 掩码分支【不得】删掉 In3(accum) / In4(momentum) / In5(stat)：
51+ * accum_t 是 parameters_out 的上游，必须照常计算；且 accum 含 ±inf 时
52+ * 0 * inf = NaN 须按 IEEE 传播进 parameters_out。四套布局输入 holder 完全相同。
53+ *
54+ * bf16 的两处特殊处理（先例 optim/apply_ftrl/op_kernel/arch35/apply_ftrl_dag.h）：
55+ * ① 标量输入：非 bf16 走 TensorScalar（Vec::CopyIn + ScalarAttr<true>，不占 UB）；
56+ * bf16 回退 Vec::Duplicate（BufferNum +2）。
57+ * ② 消费算子随之【成对切换】：TensorScalar 产出的是标量，消费者必须是标量变体
58+ * Vec::Muls；Duplicate 产出的是张量，消费者必须是 Vec::Mul。两者不可错配。
59+ *
60+ * ⚠️【TensorScalar 单槽位约束】—— 框架限制，决定了 Step 4 nesterov 的写法：
61+ * ATVOSS 的 ScalarOp 值容器 ScalarOpType（util/node.h: ScalarOpNodes::Export<VarTypeAux>）
62+ * 大小 = ScalarOp 节点【个数】，但 elewise_sch_with_scalar.h 用 【FunList 下标】存取
63+ * （Set<pos> :200/:475、Get<GetFunOutputPos<InputOp>()> :329）；而 placeholder.h:290
64+ * 的单元素特化 VarTypeStruct<T>::Get/Set 直接【忽略 offset】，于是所有越界下标静默
65+ * 塌缩到最后一个槽位 —— 本 DAG 的 lr / momentum 共用同一个物理槽，"最后写入者获胜"。
66+ * ⇒ 约束：某个 TensorScalar 的【全部消费者】必须排在下一个 TensorScalar 的 CopyIn
67+ * 之前（FunList 由 Outputs 反向 DFS、InFuns 从左到右后序生成，顺序是确定的）。
68+ * 本 DAG 满足该约束的方式：lr 全程【只被消费一次】且位于最末（OpAccTMulLr /
69+ * OpNesterovMulLr），momentum 的消费者全在其左子树内。
70+ * ⛔ 因此 nesterov【不得】展开成 p - (grad*lr + accum_t*m*lr)：那样 OpGradMulLr 会
71+ * 把 lr 的 CopyIn 提前到 momentum 之前，末尾的 OpAccTMomMulLr 再读槽位时拿到的是
72+ * momentum，device 实测结果退化为 p - (grad*lr + accum_t*m*m)（bf16 走 Duplicate
73+ * 不入该容器，故只有 fp16/fp32 中招；momentum==0 时 nesterov≡plain 亦不暴露）。
74+ */
75+template <typename U, bool useNesterov, bool hasWeightDecay, bool hasDampening, bool doWriteback, typename T = float>
76+struct SgdDag {
77+ static constexpr bool IS_BF16 = std::is_same<U, bfloat16_t>::value;
78+ // hasWeightDecay 为 false 时 Var<T,0> 让给 dampening，避免索引空洞
79+ static constexpr int VAR_IDX_DAMPENING = hasWeightDecay ? 1 : 0;
80+ 
81+ // ── 输入 holder（顺序与图原型一致：parameters / gradient / learning_rate /
82+ // accum / momentum / stat）────────────────────────────────
83+ using OpCopyInParam = Bind<Vec::CopyIn<U>, Placeholder::In0<U>>;
84+ using OpCopyInGrad = Bind<Vec::CopyIn<U>, Placeholder::In1<U>>;
85+ using OpCopyInLr = std::conditional_t<IS_BF16,
86+ Bind<Vec::Duplicate<U>, Placeholder::In2<U, Placeholder::ScalarAttr<true>>>,
87+ Bind<Vec::CopyIn<U>, Placeholder::In2<U, Placeholder::ScalarAttr<true>>>>;
88+ using OpCopyInAccum = Bind<Vec::CopyIn<U>, Placeholder::In3<U>>;
89+ using OpCopyInMom = std::conditional_t<IS_BF16,
90+ Bind<Vec::Duplicate<U>, Placeholder::In4<U, Placeholder::ScalarAttr<true>>>,
91+ Bind<Vec::CopyIn<U>, Placeholder::In4<U, Placeholder::ScalarAttr<true>>>>;
92+ // In5 仅 hasDampening == true 时进入 DAG 闭包；hasDampening == false 时
93+ // 本 typedef 不被 Outputs 可达，DAGSch 不会收录，输入退化为 5 路（无空洞）。
94+ using OpCopyInStat = Bind<Vec::CopyIn<U>, Placeholder::In5<U>>;
95+ 
96+ // ── 升 float32 域 ───────────────────────────────────────────────────────
97+ using OpParamF = Bind<Vec::Cast<T, U, 0>, OpCopyInParam>;
98+ using OpGradInF = Bind<Vec::Cast<T, U, 0>, OpCopyInGrad>;
99+ using OpLrF = Bind<Vec::Cast<T, U, 0>, OpCopyInLr>;
100+ using OpAccumF = Bind<Vec::Cast<T, U, 0>, OpCopyInAccum>;
101+ using OpMomF = Bind<Vec::Cast<T, U, 0>, OpCopyInMom>;
102+ using OpStatF = Bind<Vec::Cast<T, U, 0>, OpCopyInStat>;
103+ 
104+ // ── Step 1：权重衰减 grad = gradient + parameters * wd ──────────────────
105+ // wd == 0 时【真正跳过】而不是乘 0（spec numerical_stability.skip_zero_branches；
106+ // 且 0 * inf = NaN 会污染结果）。
107+ using VarWeightDecay = Placeholder::Var<T, SGD_VAR_IDX_WEIGHT_DECAY>;
108+ using OpParamMulWd = Bind<Vec::Muls<T>, OpParamF, VarWeightDecay>;
109+ using OpGradWithWd = Bind<Vec::Add<T>, OpGradInF, OpParamMulWd>;
110+ using OpGrad = std::conditional_t<hasWeightDecay, OpGradWithWd, OpGradInF>;
111+ 
112+ // ── Step 2：动量累积 accum_t = accum * m + grad（无条件）────────────────
113+ using OpAccMulMom = std::conditional_t<IS_BF16, Bind<Vec::Mul<T>, OpAccumF, OpMomF>,
114+ Bind<Vec::Muls<T>, OpAccumF, OpMomF>>;
115+ using OpAccumTBase = Bind<Vec::Add<T>, OpAccMulMom, OpGrad>;
116+ 
117+ // ── Step 3：阻尼修正 accum_t -= grad * ((1 - stat) * d) ─────────────────
118+ using ConstNegOne = MAKE_CONST(T, -1);
119+ using ConstOne = MAKE_CONST(T, 1);
120+ using OpStatNeg = Bind<Vec::Muls<T>, OpStatF, ConstNegOne>;
121+ using OpStatAct = Bind<Vec::Adds<T>, OpStatNeg, ConstOne>; // 1 - stat
122+ using VarDampening = Placeholder::Var<T, VAR_IDX_DAMPENING>;
123+ using OpStatActMulD = Bind<Vec::Muls<T>, OpStatAct, VarDampening>;
124+ using OpDampTerm = Bind<Vec::Mul<T>, OpGrad, OpStatActMulD>;
125+ using OpAccumTDamped = Bind<Vec::Sub<T>, OpAccumTBase, OpDampTerm>;
126+ using OpAccumT = std::conditional_t<hasDampening, OpAccumTDamped, OpAccumTBase>;
127+ 
128+ // ── Step 4：权重更新（无条件写出）───────────────────────────────────────
129+ // 非 nesterov：parameters - accum_t * lr
130+ using OpAccTMulLr = std::conditional_t<IS_BF16, Bind<Vec::Mul<T>, OpAccumT, OpLrF>,
131+ Bind<Vec::Muls<T>, OpAccumT, OpLrF>>;
132+ using OpParamPlain = Bind<Vec::Sub<T>, OpParamF, OpAccTMulLr>;
133+ 
134+ // nesterov：parameters - (grad + accum_t * m) * lr
135+ // ⚠️【不得】改写成展开式 p - (grad*lr + accum_t*m*lr)：见文件头「TensorScalar 单槽位」约束。
136+ using OpAccTMulMom = std::conditional_t<IS_BF16, Bind<Vec::Mul<T>, OpAccumT, OpMomF>,
137+ Bind<Vec::Muls<T>, OpAccumT, OpMomF>>;
138+ using OpNesterovSum = Bind<Vec::Add<T>, OpGrad, OpAccTMulMom>; // grad + accum_t * m
139+ using OpNesterovMulLr = std::conditional_t<IS_BF16, Bind<Vec::Mul<T>, OpNesterovSum, OpLrF>,
140+ Bind<Vec::Muls<T>, OpNesterovSum, OpLrF>>;
141+ using OpParamNesterov = Bind<Vec::Sub<T>, OpParamF, OpNesterovMulLr>;
142+ 
143+ using OpParamNew = std::conditional_t<useNesterov, OpParamNesterov, OpParamPlain>;
144+ using OpParamOutCast = Bind<Vec::Cast<U, T, 1>, OpParamNew>; // CAST_RINT：就近偶数舍入
145+ using OpCopyOutParam = Bind<Vec::CopyOut<U>, Placeholder::Out0<U>, OpParamOutCast>; // output: parameters
146+ 
147+ // ── Step 5：momentum != 0 掩码控制的两路回写 ────────────────────────────
148+ using OpAccumOutCast = Bind<Vec::Cast<U, T, 1>, OpAccumT>;
149+ using OpCopyOutAccum = Bind<Vec::CopyOut<U>, Placeholder::Out1<U>, OpAccumOutCast>; // 原地回写 input3: accum
150+ 
151+ using ConstZero = MAKE_CONST(T, 0);
152+ using OpZeroTsr = Bind<Vec::Duplicate<T>, ConstZero>;
153+ using OpStatOutCast = Bind<Vec::Cast<U, T, 1>, OpZeroTsr>;
154+ using OpCopyOutStat = Bind<Vec::CopyOut<U>, Placeholder::Out2<U>, OpStatOutCast>; // 原地回写 input5: stat
155+ 
156+ // Outputs 是两套 DAG 的【唯一】差异行。掩码 DAG 下 OpZeroTsr / OpCopyOutAccum /
157+ // OpCopyOutStat 随 Elems 收缩被 DAGSch 从 Outputs 反向推导时整体裁掉，
158+ // BufferNum 比回写 DAG 小 5~6（Host 按回写 DAG 反解 ubFormer，故不会溢出）。
159+ using Outputs = std::conditional_t<doWriteback, Elems<OpCopyOutParam, OpCopyOutAccum, OpCopyOutStat>,
160+ Elems<OpCopyOutParam>>;
161+ using MemCfg = MemOptCfg<MemLevel::LEVEL_2>;
162+ using OpDag = DAGSch<Outputs, void, MemCfg>;
163+};
164+} // namespace SgdOp
165+ 
166+#endif // SGD_DAG_H
Aoptim/sgd/op_kernel/arch35/sgd_tiling_data.h+31-0
@@ -0,0 +1,31 @@
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+#ifndef SGD_TILING_DATA_H
12+#define SGD_TILING_DATA_H
13+ 
14+/* !
15+ * \file sgd_tiling_data.h
16+ * \brief Host / Kernel 共享的 TilingData（普通 C++ struct，不用废弃宏 BEGIN_TILING_DATA_DEF）
17+ */
18+ 
19+#include "atvoss/elewise/elewise_base_struct.h"
20+ 
21+struct SgdRegbaseTilingData {
22+ Ops::Base::EleBaseTilingDataV2 elewiseTiling; // 框架填充：elemNum / ubFormer / blockNum / scheMode ...
23+ float dampening; // 属性注入，仅 hasDampening == 1 时经 sch.SetVar 下发
24+ float weightDecay; // 属性注入，仅 hasWeightDecay == 1 时经 sch.SetVar 下发
25+};
26+ 
27+// 注：nesterov 是纯编译期分支（TilingKey 维度），不占 TilingData 字段。
28+// learning_rate / momentum 是 Device 侧 [1] 张量，由 Placeholder::ScalarAttr<true>
29+// 从 GM 读取，【不得】在 Host 侧读值放进 TilingData。
30+ 
31+#endif // SGD_TILING_DATA_H
Aoptim/sgd/op_kernel/arch35/sgd_tiling_key.h+63-0
@@ -0,0 +1,63 @@
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 sgd_tiling_key.h
13+ * \brief SGD TilingKey 模板参数定义
14+ *
15+ * 模板参数（四维，全部是编译期分支）：
16+ * - schMode : 框架调度模式，由 ElewiseBaseTiling 自动决定（0 / 1）
17+ * - useNesterov : 属性 nesterov（0 / 1）
18+ * - hasWeightDecay : 属性 weight_decay != 0（0 / 1）
19+ * - hasDampening : 属性 dampening != 0（0 / 1）
20+ *
21+ * 组合数核算：
22+ * 业务 TilingKey = 2(nesterov) × 2(wd) × 2(damp) - 2(非法) = 6 （K0~K5）
23+ * TPL_SEL 展开 = 6 × 2(schMode) = 12
24+ * Kernel binary = 12 × 3(dtype，来自 binary.json) = 36
25+ *
26+ * 非法组合 useNesterov == 1 && hasDampening == 1 由下方两组 ARGS_SEL 剪掉，
27+ * 【不生成对应 binary】；Host 侧 InferShape / Tiling 亦对该组合报
28+ * attribute_value_out_of_range（对齐 A2：nesterov == true 时 dampening 必须为 0）。
29+ *
30+ * ⛔ `momentum == 0` 掩码【不是】TilingKey 维度：momentum 是 Device 侧 [1] 张量，
31+ * Host Tiling 收不到张量数据（$ATV/elewise/elewise_tiling.h:216-245），
32+ * 其值在 Tiling 阶段不可见 → 只能做运行期分支。两套 DAG 同时存在于同一个
33+ * binary 内、由 sgd.cpp 的运行期 if 选择，binary 数量不变。
34+ */
35+ 
36+#ifndef SGD_TILING_KEY_H
37+#define SGD_TILING_KEY_H
38+ 
39+#include "ascendc/host_api/tiling/template_argument.h"
40+ 
41+#define SGD_TPL_FALSE 0
42+#define SGD_TPL_TRUE 1
43+#define SGD_TPL_BIT_WIDTH 1
44+ 
45+ASCENDC_TPL_ARGS_DECL(
46+ SGD, ASCENDC_TPL_UINT_DECL(schMode, SGD_TPL_BIT_WIDTH, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE),
47+ ASCENDC_TPL_UINT_DECL(useNesterov, SGD_TPL_BIT_WIDTH, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE),
48+ ASCENDC_TPL_UINT_DECL(hasWeightDecay, SGD_TPL_BIT_WIDTH, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE),
49+ ASCENDC_TPL_UINT_DECL(hasDampening, SGD_TPL_BIT_WIDTH, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE));
50+ 
51+ASCENDC_TPL_SEL(
52+ // 组一：useNesterov == 0 —— dampening 可 0 可非 0，共 2×2×2 = 8 个组合（K0~K3 × schMode）
53+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(schMode, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE),
54+ ASCENDC_TPL_UINT_SEL(useNesterov, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE),
55+ ASCENDC_TPL_UINT_SEL(hasWeightDecay, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE),
56+ ASCENDC_TPL_UINT_SEL(hasDampening, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE)),
57+ // 组二：useNesterov == 1 —— dampening 必为 0（非法组合已剪），共 2×2 = 4 个组合（K4~K5 × schMode）
58+ ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(schMode, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE),
59+ ASCENDC_TPL_UINT_SEL(useNesterov, ASCENDC_TPL_UI_LIST, SGD_TPL_TRUE),
60+ ASCENDC_TPL_UINT_SEL(hasWeightDecay, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE, SGD_TPL_TRUE),
61+ ASCENDC_TPL_UINT_SEL(hasDampening, ASCENDC_TPL_UI_LIST, SGD_TPL_FALSE)), );
62+ 
63+#endif // SGD_TILING_KEY_H
Aoptim/sgd/tests/CMakeLists.txt+18-0
@@ -0,0 +1,18 @@
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+message(STATUS "=== Debug: start ops.optim.sgd.tests.CMakeLists.txt ")
12+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
13+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
14+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
15+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
16+ add_subdirectory(${SUB_DIR})
17+ endif()
18+endforeach()
Aoptim/sgd/tests/assets/golden.py+154-0
@@ -0,0 +1,154 @@
1+#!/usr/bin/env python3
2+# -*- coding: utf-8 -*-
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+# ----------------------------------------------------------------------------
11+ 
12+"""SGD kernel golden.
13+ 
14+Formula (aligned with docs/spec.yaml math_semantics.formula):
15+ grad = gradient + parameters * weight_decay (only when weight_decay != 0)
16+ accum_t = accum * m + grad (UNCONDITIONAL)
17+ accum_t -= grad * (1 - stat) * dampening (only when dampening != 0)
18+ p_out = p - (grad*lr + accum_t*m*lr) if nesterov else p - accum_t*lr
19+ m != 0 : accum_out = accum_t ; stat_out = 0
20+ m == 0 : accum_out / stat_out keep the input bit pattern (no writeback)
21+ 
22+The arithmetic runs on torch tensors in the float32 domain, matching what the NPU
23+does, and follows the repo convention of lifting bfloat16 / float16 to float32
24+before handing the buffer to torch and casting the result back afterwards.
25+ 
26+Three rules that must not be violated:
27+ 1. accum_t is computed UNCONDITIONALLY. Do not treat "m == 0" as "multiplying by
28+ zero can be skipped": when accum holds +-inf, 0 * inf = NaN and that NaN must
29+ propagate into parameters_out per IEEE 754.
30+ 2. When weight_decay == 0 / dampening == 0 the corresponding step is genuinely
31+ skipped rather than written as "multiply by zero" (same 0 * inf problem).
32+ 3. The m == 0 branch returns the input buffers themselves, not values recomputed
33+ to be equal. NaN payloads, -0.0 and +-inf must survive bit for bit, so that
34+ branch never round-trips through torch.
35+ 
36+Note: learning_rate and momentum share the dtype of parameters, so under float16 /
37+bfloat16 those two scalars are already quantized (lr = 0.1 is not exactly 0.1 in
38+float16). This function consumes the tensor values as given and never rebuilds them
39+from Python float literals.
40+"""
41+ 
42+import numpy as np
43+ 
44+__golden__ = {"kernel": {"sgd": "sgd_golden"}}
45+ 
46+ 
47+def _to_f32_np(arr):
48+ """Lift an input to the float32 compute domain (the NPU works in float32 too).
49+ 
50+ TTK hands over bfloat16 as a real ml_dtypes bfloat16 dtype, which supports
51+ astype(np.float32) directly. Only a raw uint16 bit pattern needs manual shifting.
52+ torch has no bfloat16 numpy bridge, so this widening also makes the buffer
53+ something torch.from_numpy can take.
54+ """
55+ a = np.asarray(arr)
56+ if a.dtype == np.uint16:
57+ return (a.astype(np.uint32) << 16).view(np.float32)
58+ return np.ascontiguousarray(a.astype(np.float32))
59+ 
60+ 
61+def _cast_back(t, ref):
62+ """Cast a torch float32 result back to the reference tensor dtype.
63+ 
64+ Never return a uint16 view here: the harness clamps goldens with
65+ array[array < dtype_min] = -inf, which raises OverflowError on integer arrays.
66+ """
67+ return t.numpy().astype(np.asarray(ref).dtype)
68+ 
69+ 
70+def _scalar(arr):
71+ """Read a [1] / 0-d scalar tensor out as an exactly-representable float32 value.
72+ 
73+ float() widens float32 to double losslessly, and torch casts the Python scalar
74+ back to the tensor dtype when it meets a float32 tensor, so the round trip adds
75+ no rounding of its own.
76+ """
77+ return float(_to_f32_np(arr).reshape(-1)[0])
78+ 
79+ 
80+def sgd_golden(
81+ parameters,
82+ gradient,
83+ learning_rate,
84+ accum,
85+ momentum,
86+ stat,
87+ dampening=0.0,
88+ weight_decay=0.0,
89+ nesterov=False,
90+ **kwargs,
91+):
92+ """Golden function for SGD kernel.
93+ 
94+ Supported dtypes: float32, float16, bfloat16 (all six inputs share one dtype).
95+ 
96+ Args:
97+ parameters: weights to update (numpy.ndarray)
98+ gradient: gradient, same shape and dtype as parameters
99+ learning_rate: scalar tensor of shape [1]
100+ accum: momentum accumulator, same shape and dtype as parameters
101+ momentum: scalar tensor of shape [1]
102+ stat: per-element first-step flag, same shape and dtype as parameters
103+ dampening: float attribute, default 0.0
104+ weight_decay: float attribute, default 0.0
105+ nesterov: bool attribute, default False
106+ 
107+ Returns:
108+ [parameters_out, accum_out, stat_out], each cast back to the input dtype.
109+ All three are in-place writeback slots; accum_out and stat_out return the
110+ input arrays unchanged when momentum == 0 (the same bits, not recomputed
111+ equal values).
112+ """
113+ import torch
114+ 
115+ p32 = torch.from_numpy(_to_f32_np(parameters))
116+ g32 = torch.from_numpy(_to_f32_np(gradient))
117+ a32 = torch.from_numpy(_to_f32_np(accum))
118+ s32 = torch.from_numpy(_to_f32_np(stat))
119+ lr = _scalar(learning_rate)
120+ m = _scalar(momentum)
121+ 
122+ d = float(dampening)
123+ wd = float(weight_decay)
124+ nest = bool(nesterov)
125+ 
126+ # Step 1: genuinely skipped when weight_decay == 0
127+ grad = (g32 + p32 * wd) if wd != 0.0 else g32
128+ 
129+ # Step 2: unconditional
130+ accum_t = a32 * m + grad
131+ 
132+ # Step 3: genuinely skipped when dampening == 0
133+ if d != 0.0:
134+ accum_t = accum_t - grad * ((1.0 - s32) * d)
135+ 
136+ # Step 4: unconditional writeback
137+ if nest:
138+ p_new = p32 - (grad * lr + accum_t * m * lr)
139+ else:
140+ p_new = p32 - accum_t * lr
141+ 
142+ parameters_out = _cast_back(p_new, parameters)
143+ 
144+ # Step 5: momentum != 0 mask. IEEE != is used, so -0.0 counts as zero while
145+ # 1e-8 / 1e-30 count as non-zero.
146+ if m != 0.0:
147+ accum_out = _cast_back(accum_t, accum)
148+ stat_out = _cast_back(torch.zeros_like(s32), stat)
149+ else:
150+ # Keep the input bits: return copies of the inputs, no numeric rebuild.
151+ accum_out = np.asarray(accum).copy()
152+ stat_out = np.asarray(stat).copy()
153+ 
154+ return [parameters_out, accum_out, stat_out]
Aoptim/sgd/tests/ut/CMakeLists.txt+18-0
@@ -0,0 +1,18 @@
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+# 每个目录下需要生成的可执行文件，具体参考：ops/built-in/test/CMakeLists.txt： 50~124
12+file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*)
13+message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ")
14+foreach(SUB_DIR ${CURRENT_SOURCE_DIRS})
15+ if(EXISTS "${SUB_DIR}/CMakeLists.txt")
16+ add_subdirectory(${SUB_DIR})
17+ endif()
18+endforeach()
Aoptim/sgd/tests/ut/op_host/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+message(STATUS "=== Debug: CURRENT_DIRS =${CURRENT_DIRS} ")
11+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+if(UT_TEST_ALL OR OP_HOST_UT)
13+ add_modules_ut_sources(HOSTNAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
14+ add_modules_ut_sources(HOSTNAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
15+endif()
Aoptim/sgd/tests/ut/op_host/arch35/test_sgd_tiling.cpp+382-0
@@ -0,0 +1,382 @@
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 test_sgd_tiling.cpp
13+ * \brief SGD arch35 Tiling 单元测试
14+ *
15+ * 覆盖面：
16+ * ① 6 个业务 TilingKey（K0~K5 = useNesterov × hasWeightDecay × hasDampening 的合法组合）× 3 dtype
17+ * ② 非法组合 nesterov && dampening != 0 被 GetAttr 拦下
18+ * ③ weight_decay < 0 被拦下
19+ * ④ 本算子相对 910B/910C 补齐的校验：大张量不同形、标量 shape != [1]、dtype 不一致、空 tensor
20+ *
21+ * TilingKey 编码（4 位，见 sgd_tiling_key.h）：
22+ * bit0 schMode(框架决定) / bit1 useNesterov / bit2 hasWeightDecay / bit3 hasDampening
23+ * 故期望 key = schMode | (nesterov<<1) | (hasWd<<2) | (hasDamp<<3)。
24+ * schMode 由 ElewiseBaseTiling 依 shape 自行决定，本 UT 用 GetTilingKey() 的高 3 位做断言，
25+ * 不硬编码 schMode，避免与框架实现耦合。
26+ */
27+ 
28+#include <gtest/gtest.h>
29+ 
30+#include <fstream>
31+#include <iostream>
32+#include <vector>
33+ 
34+#include "../../../../op_host/arch35/sgd_tiling.h"
35+#include "ut_op_util.h"
36+#include "exe_graph/runtime/storage_format.h"
37+#include "exe_graph/runtime/storage_shape.h"
38+#include "kernel_run_context_facker.h"
39+#include "test_cube_util.h"
40+ 
41+using namespace ut_util;
42+using namespace std;
43+using namespace ge;
44+ 
45+class TestSgdTiling : public testing::Test {
46+protected:
47+ static void SetUpTestCase() { std::cout << "TestSgdTiling SetUp" << std::endl; }
48+ 
49+ static void TearDownTestCase() { std::cout << "TestSgdTiling TearDown" << std::endl; }
50+};
51+ 
52+namespace {
53+constexpr size_t SGD_INPUT_NUM = 6;
54+constexpr size_t SGD_OUTPUT_NUM = 1;
55+ 
56+void InitPlatForm(fe::PlatFormInfos& platFormInfo, map<string, string>& socInfos, map<string, string>& aicoreSpec,
57+ map<string, string>& intrinsics, map<string, string>& socVersion)
58+{
59+ string compile_info_string = R"({
60+ "hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1",
61+ "Intrinsic_fix_pipe_l0c2out": false,
62+ "Intrinsic_data_move_l12ub": true,
63+ "Intrinsic_data_move_l0c2ub": true,
64+ "Intrinsic_data_move_out2l1_nd2nz": false,
65+ "UB_SIZE": 245760, "L2_SIZE": 33554432, "L1_SIZE": 524288,
66+ "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072,
67+ "CORE_NUM": 64, "socVersion": "Ascend950"}})";
68+ GetPlatFormInfos(compile_info_string.c_str(), socInfos, aicoreSpec, intrinsics, socVersion);
69+ platFormInfo.Init();
70+}
71+ 
72+struct SgdUtCompileInfo {};
73+ 
74+// 通用 tiling 驱动。shapes 允许逐路不同，便于构造"不同形"负用例。
75+ge::graphStatus RunSgdTiling(gert::StorageShape& paramShape, gert::StorageShape& gradShape, gert::StorageShape& lrShape,
76+ gert::StorageShape& accumShape, gert::StorageShape& momentumShape,
77+ gert::StorageShape& statShape, ge::DataType paramDtype, ge::DataType otherDtype,
78+ float dampening, float weightDecay, bool nesterov, uint64_t* outTilingKey)
79+{
80+ fe::PlatFormInfos platFormInfo;
81+ map<string, string> socInfos;
82+ map<string, string> aicoreSpec;
83+ map<string, string> intrinsics;
84+ map<string, string> socVersion = {{"Short_SoC_version", "ASCEND950"}};
85+ InitPlatForm(platFormInfo, socInfos, aicoreSpec, intrinsics, socVersion);
86+ 
87+ std::string opType("SGD");
88+ auto impl = gert::OpImplRegistry::GetInstance().GetOpImpl(opType.c_str());
89+ if (impl == nullptr || impl->tiling == nullptr) {
90+ return ge::GRAPH_FAILED;
91+ }
92+ auto tiling_func = impl->tiling;
93+ 
94+ auto param = gert::TilingData::CreateCap(4096);
95+ auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
96+ auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
97+ if (param == nullptr) {
98+ return ge::GRAPH_FAILED;
99+ }
100+ 
101+ SgdUtCompileInfo compileInfo;
102+ auto inFormat = ge::FORMAT_ND;
103+ 
104+ auto holder = gert::TilingContextFaker()
105+ .SetOpType(opType)
106+ .NodeIoNum(SGD_INPUT_NUM, SGD_OUTPUT_NUM)
107+ .IrInstanceNum({1, 1, 1, 1, 1, 1})
108+ .InputShapes({&paramShape, &gradShape, &lrShape, &accumShape, &momentumShape, &statShape})
109+ .OutputShapes({&paramShape})
110+ .CompileInfo(&compileInfo)
111+ .PlatformInfo(reinterpret_cast<char*>(&platFormInfo))
112+ .NodeInputTd(0, paramDtype, inFormat, inFormat)
113+ .NodeInputTd(1, otherDtype, inFormat, inFormat)
114+ .NodeInputTd(2, otherDtype, inFormat, inFormat)
115+ .NodeInputTd(3, otherDtype, inFormat, inFormat)
116+ .NodeInputTd(4, otherDtype, inFormat, inFormat)
117+ .NodeInputTd(5, otherDtype, inFormat, inFormat)
118+ .NodeOutputTd(0, paramDtype, inFormat, inFormat)
119+ .NodeAttrs({{"dampening", Ops::NN::AnyValue::CreateFrom<float>(dampening)},
120+ {"weight_decay", Ops::NN::AnyValue::CreateFrom<float>(weightDecay)},
121+ {"nesterov", Ops::NN::AnyValue::CreateFrom<bool>(nesterov)}})
122+ .TilingData(param.get())
123+ .Workspace(ws_size)
124+ .Build();
125+ 
126+ gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>();
127+ if (tiling_context->GetPlatformInfo() == nullptr) {
128+ return ge::GRAPH_FAILED;
129+ }
130+ tiling_context->GetPlatformInfo()->SetPlatformRes("SoCInfo", socInfos);
131+ tiling_context->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicoreSpec);
132+ tiling_context->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
133+ tiling_context->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
134+ tiling_context->GetPlatformInfo()->SetPlatformRes("version", socVersion);
135+ 
136+ auto ret = tiling_func(tiling_context);
137+ if (ret == ge::GRAPH_SUCCESS && outTilingKey != nullptr) {
138+ *outTilingKey = tiling_context->GetTilingKey();
139+ }
140+ return ret;
141+}
142+ 
143+// 正常路径便捷封装：6 路同形同 dtype
144+ge::graphStatus RunSgdTilingNormal(std::initializer_list<int64_t> shape, ge::DataType dtype, float dampening,
145+ float weightDecay, bool nesterov, uint64_t* outTilingKey)
146+{
147+ gert::StorageShape big = {shape, shape};
148+ gert::StorageShape one = {{1}, {1}};
149+ return RunSgdTiling(big, big, one, big, one, big, dtype, dtype, dampening, weightDecay, nesterov, outTilingKey);
150+}
151+ 
152+// TilingKey 的业务位（剥掉 bit0 的 schMode）：nesterov | hasWd<<1 | hasDamp<<2
153+uint64_t BizBits(uint64_t tilingKey) { return tilingKey >> 1; }
154+constexpr uint64_t BizExpect(bool nesterov, bool hasWd, bool hasDamp)
155+{
156+ return (nesterov ? 1U : 0U) | (hasWd ? 2U : 0U) | (hasDamp ? 4U : 0U);
157+}
158+} // namespace
159+ 
160+// ───────────────────────── K0 ~ K5 × 3 dtype ─────────────────────────
161+ 
162+TEST_F(TestSgdTiling, sgd_tiling_K0_no_branch_fp32)
163+{
164+ uint64_t key = 0;
165+ ASSERT_EQ(RunSgdTilingNormal({16, 26, 16, 19}, ge::DT_FLOAT, 0.0f, 0.0f, false, &key), ge::GRAPH_SUCCESS);
166+ EXPECT_EQ(BizBits(key), BizExpect(false, false, false));
167+}
168+ 
169+TEST_F(TestSgdTiling, sgd_tiling_K1_dampening_only_fp32)
170+{
171+ uint64_t key = 0;
172+ ASSERT_EQ(RunSgdTilingNormal({16, 26, 16, 19}, ge::DT_FLOAT, 0.5f, 0.0f, false, &key), ge::GRAPH_SUCCESS);
173+ EXPECT_EQ(BizBits(key), BizExpect(false, false, true));
174+}
175+ 
176+TEST_F(TestSgdTiling, sgd_tiling_K2_weight_decay_only_fp16)
177+{
178+ uint64_t key = 0;
179+ ASSERT_EQ(RunSgdTilingNormal({3761, 4, 44, 4}, ge::DT_FLOAT16, 0.0f, 0.01f, false, &key), ge::GRAPH_SUCCESS);
180+ EXPECT_EQ(BizBits(key), BizExpect(false, true, false));
181+}
182+ 
183+TEST_F(TestSgdTiling, sgd_tiling_K3_both_branches_fp16)
184+{
185+ uint64_t key = 0;
186+ ASSERT_EQ(RunSgdTilingNormal({3761, 4, 44, 4}, ge::DT_FLOAT16, 0.5f, 0.01f, false, &key), ge::GRAPH_SUCCESS);
187+ EXPECT_EQ(BizBits(key), BizExpect(false, true, true));
188+}
189+ 
190+TEST_F(TestSgdTiling, sgd_tiling_K4_nesterov_with_weight_decay_bf16)
191+{
192+ uint64_t key = 0;
193+ ASSERT_EQ(RunSgdTilingNormal({7, 2, 7, 8, 10}, ge::DT_BF16, 0.0f, 0.01f, true, &key), ge::GRAPH_SUCCESS);
194+ EXPECT_EQ(BizBits(key), BizExpect(true, true, false));
195+}
196+ 
197+TEST_F(TestSgdTiling, sgd_tiling_K5_nesterov_only_bf16)
198+{
199+ uint64_t key = 0;
200+ ASSERT_EQ(RunSgdTilingNormal({7, 2, 7, 8, 10}, ge::DT_BF16, 0.0f, 0.0f, true, &key), ge::GRAPH_SUCCESS);
201+ EXPECT_EQ(BizBits(key), BizExpect(true, false, false));
202+}
203+ 
204+// 三个 dtype 在同一分支下都要能出 tiling（binary.json 有 3 个 dtype 条目）
205+TEST_F(TestSgdTiling, sgd_tiling_K0_all_three_dtypes)
206+{
207+ uint64_t key = 0;
208+ for (auto dt : {ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) {
209+ ASSERT_EQ(RunSgdTilingNormal({256, 256}, dt, 0.0f, 0.0f, false, &key), ge::GRAPH_SUCCESS);
210+ EXPECT_EQ(BizBits(key), BizExpect(false, false, false));
211+ }
212+}
213+ 
214+// 尾块 / 非 32B 对齐 shape 也要能出 tiling
215+TEST_F(TestSgdTiling, sgd_tiling_unaligned_tail_shape)
216+{
217+ uint64_t key = 0;
218+ ASSERT_EQ(RunSgdTilingNormal({33, 2}, ge::DT_FLOAT, 0.0f, 0.0f, false, &key), ge::GRAPH_SUCCESS);
219+ ASSERT_EQ(RunSgdTilingNormal({7, 7, 33, 2}, ge::DT_FLOAT, 0.0f, 0.0f, false, &key), ge::GRAPH_SUCCESS);
220+}
221+ 
222+// ───────────────────────── 属性非法组合 ─────────────────────────
223+ 
224+TEST_F(TestSgdTiling, sgd_tiling_illegal_nesterov_with_dampening)
225+{
226+ // nesterov && dampening != 0 —— 被 GetAttr 拦下，且该组合不生成 binary
227+ uint64_t key = 0;
228+ ASSERT_EQ(RunSgdTilingNormal({256, 256}, ge::DT_FLOAT, 0.5f, 0.0f, true, &key), ge::GRAPH_FAILED);
229+}
230+ 
231+TEST_F(TestSgdTiling, sgd_tiling_illegal_negative_weight_decay)
232+{
233+ uint64_t key = 0;
234+ ASSERT_EQ(RunSgdTilingNormal({256, 256}, ge::DT_FLOAT, 0.0f, -0.01f, false, &key), ge::GRAPH_FAILED);
235+}
236+ 
237+// ───────────── 本算子相对 910B/910C 补齐的校验（canndev 仅校验 parameters 的 rank）─────────────
238+ 
239+TEST_F(TestSgdTiling, sgd_tiling_reject_rank0)
240+{
241+ gert::StorageShape zeroRank = {{}, {}};
242+ gert::StorageShape one = {{1}, {1}};
243+ ASSERT_EQ(RunSgdTiling(zeroRank, zeroRank, one, zeroRank, one, zeroRank, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f,
244+ false, nullptr),
245+ ge::GRAPH_FAILED);
246+}
247+ 
248+TEST_F(TestSgdTiling, sgd_tiling_reject_rank9)
249+{
250+ std::initializer_list<int64_t> r9 = {2, 2, 2, 2, 2, 2, 2, 2, 2};
251+ gert::StorageShape big = {r9, r9};
252+ gert::StorageShape one = {{1}, {1}};
253+ ASSERT_EQ(RunSgdTiling(big, big, one, big, one, big, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
254+ ge::GRAPH_FAILED);
255+}
256+ 
257+TEST_F(TestSgdTiling, sgd_tiling_reject_empty_tensor)
258+{
259+ // 空 tensor 判非法（不是"空进空出"）—— accum/stat 的原地回写在 numel == 0 下无定义
260+ gert::StorageShape empty = {{0, 3}, {0, 3}};
261+ gert::StorageShape one = {{1}, {1}};
262+ ASSERT_EQ(
263+ RunSgdTiling(empty, empty, one, empty, one, empty, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
264+ ge::GRAPH_FAILED);
265+}
266+ 
267+// ── DFX：空 tensor「全空间」枚举 ────────────────────────────────────────────
268+// changwei-op-dev step5-verify §5.4 要求：空 tensor = **任意一轴或多轴为 0**，
269+// 必须逐形态真跑分类，不能只验一个维度，结论要写成「仅支持哪个轴为空」。
270+// 本算子的契约是【所有形态一律拒绝】（无空进空出语义），故下面每一条都断言 GRAPH_FAILED。
271+// 若将来放开某一轴，本组用例会立刻在该形态上变红，逼迫同步更新契约与文档。
272+//
273+// 覆盖：1-D [0] / 2-D 每轴单独为 0 / 2-D 双轴同时为 0 /
274+// 3-D 每轴单独为 0 / 3-D 多轴组合为 0 / 8-D（rank 上界）末轴为 0
275+struct EmptyShapeCase {
276+ const char* desc;
277+ std::vector<int64_t> dims;
278+};
279+ 
280+TEST_F(TestSgdTiling, sgd_tiling_reject_empty_tensor_full_space)
281+{
282+ const std::vector<EmptyShapeCase> cases = {
283+ {"1D_[0]", {0}},
284+ {"2D_axis0_[0,3]", {0, 3}},
285+ {"2D_axis1_[2,0]", {2, 0}},
286+ {"2D_both_[0,0]", {0, 0}},
287+ {"3D_axis0_[0,2,3]", {0, 2, 3}},
288+ {"3D_axis1_[2,0,3]", {2, 0, 3}},
289+ {"3D_axis2_[2,3,0]", {2, 3, 0}},
290+ {"3D_axis01_[0,0,3]", {0, 0, 3}},
291+ {"3D_axis02_[0,2,0]", {0, 2, 0}},
292+ {"3D_all_[0,0,0]", {0, 0, 0}},
293+ {"8D_lastaxis_[2,2,2,2,2,2,2,0]", {2, 2, 2, 2, 2, 2, 2, 0}},
294+ };
295+ gert::StorageShape one = {{1}, {1}};
296+ for (const auto& c : cases) {
297+ // gert::StorageShape 只能用花括号字面量或逐维 AppendDim 构造，
298+ // 【不能】从 std::vector<int64_t> 隐式转换（会报 could not convert ... to gert::StorageShape）。
299+ gert::StorageShape empty;
300+ for (int64_t d : c.dims) {
301+ empty.MutableOriginShape().AppendDim(d);
302+ empty.MutableStorageShape().AppendDim(d);
303+ }
304+ EXPECT_EQ(
305+ RunSgdTiling(empty, empty, one, empty, one, empty, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
306+ ge::GRAPH_FAILED)
307+ << "空 tensor 形态 " << c.desc << " 未被拒绝；本算子契约为【任意一轴或多轴为 0 一律非法】";
308+ }
309+}
310+ 
311+TEST_F(TestSgdTiling, sgd_tiling_reject_empty_tensor_partial_inputs)
312+{
313+ // 只有【部分】输入为空的形态：parameters 非空但 gradient / accum / stat 为空。
314+ // 这类先撞 CheckSameShape（形状不等）也算拒绝，但必须确认「不崩、返 GRAPH_FAILED」。
315+ gert::StorageShape param = {{2, 3}, {2, 3}};
316+ gert::StorageShape empty = {{0, 3}, {0, 3}};
317+ gert::StorageShape one = {{1}, {1}};
318+ EXPECT_EQ(
319+ RunSgdTiling(param, empty, one, param, one, param, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
320+ ge::GRAPH_FAILED)
321+ << "gradient 为空未被拒绝";
322+ EXPECT_EQ(
323+ RunSgdTiling(param, param, one, empty, one, param, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
324+ ge::GRAPH_FAILED)
325+ << "accum 为空未被拒绝";
326+ EXPECT_EQ(
327+ RunSgdTiling(param, param, one, param, one, empty, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
328+ ge::GRAPH_FAILED)
329+ << "stat 为空未被拒绝";
330+}
331+ 
332+TEST_F(TestSgdTiling, sgd_tiling_reject_shape_mismatch_gradient)
333+{
334+ // 大张量必须严格同形，**不做广播** —— "可广播但不相等"同样非法
335+ gert::StorageShape param = {{2, 3}, {2, 3}};
336+ gert::StorageShape grad = {{1, 3}, {1, 3}};
337+ gert::StorageShape one = {{1}, {1}};
338+ ASSERT_EQ(RunSgdTiling(param, grad, one, param, one, param, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
339+ ge::GRAPH_FAILED);
340+}
341+ 
342+TEST_F(TestSgdTiling, sgd_tiling_reject_shape_mismatch_stat)
343+{
344+ gert::StorageShape param = {{2, 4}, {2, 4}};
345+ gert::StorageShape stat = {{2, 5}, {2, 5}};
346+ gert::StorageShape one = {{1}, {1}};
347+ ASSERT_EQ(RunSgdTiling(param, param, one, param, one, stat, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
348+ ge::GRAPH_FAILED);
349+}
350+ 
351+TEST_F(TestSgdTiling, sgd_tiling_reject_scalar_shape_not_one)
352+{
353+ // learning_rate / momentum 必须是 [1] 或 0D 标量
354+ gert::StorageShape param = {{2, 4}, {2, 4}};
355+ gert::StorageShape two = {{2}, {2}};
356+ gert::StorageShape one = {{1}, {1}};
357+ ASSERT_EQ(
358+ RunSgdTiling(param, param, two, param, one, param, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
359+ ge::GRAPH_FAILED);
360+ ASSERT_EQ(
361+ RunSgdTiling(param, param, one, param, two, param, ge::DT_FLOAT, ge::DT_FLOAT, 0.0f, 0.0f, false, nullptr),
362+ ge::GRAPH_FAILED);
363+}
364+ 
365+TEST_F(TestSgdTiling, sgd_tiling_reject_dtype_mismatch)
366+{
367+ // 9 个张量位必须同 dtype（parameters 为 fp32，其余为 fp16）
368+ gert::StorageShape param = {{2, 4}, {2, 4}};
369+ gert::StorageShape one = {{1}, {1}};
370+ ASSERT_EQ(
371+ RunSgdTiling(param, param, one, param, one, param, ge::DT_FLOAT, ge::DT_FLOAT16, 0.0f, 0.0f, false, nullptr),
372+ ge::GRAPH_FAILED);
373+}
374+ 
375+TEST_F(TestSgdTiling, sgd_tiling_reject_unsupported_dtype)
376+{
377+ gert::StorageShape param = {{2, 4}, {2, 4}};
378+ gert::StorageShape one = {{1}, {1}};
379+ ASSERT_EQ(
380+ RunSgdTiling(param, param, one, param, one, param, ge::DT_INT32, ge::DT_INT32, 0.0f, 0.0f, false, nullptr),
381+ ge::GRAPH_FAILED);
382+}
Aoptim/sgd/tests/ut/op_host/test_sgd_infershape.cpp+162-0
@@ -0,0 +1,162 @@
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 test_sgd_infershape.cpp
13+ * \brief SGD InferShape / InferDataType 单元测试
14+ *
15+ * 覆盖面：
16+ * - 正常推导（rank 2 / rank 1 / rank 8 边界）
17+ * - UNKNOWN_RANK(-2) 透传
18+ * - rank-0 拒绝、rank 9 拒绝（对齐 910B/910C 的 1~8）
19+ * - 属性非法拒绝：nesterov && dampening != 0、weight_decay < 0
20+ * - InferDataType 双挂验证（漏挂会导致 GE 输出 dtype 推导缺失）
21+ */
22+ 
23+#include <gtest/gtest.h> // NOLINT
24+#include <iostream>
25+#include <vector>
26+#include "infershape_test_util.h" // NOLINT
27+#include "ut_op_common.h"
28+#include "../../../op_graph/sgd_proto.h"
29+ 
30+class SGD : public testing::Test {
31+protected:
32+ static void SetUpTestCase() { std::cout << "SGD SetUp" << std::endl; }
33+ 
34+ static void TearDownTestCase() { std::cout << "SGD TearDown" << std::endl; }
35+};
36+ 
37+namespace {
38+constexpr size_t SGD_INPUT_NUM = 6;
39+constexpr size_t SGD_OUTPUT_NUM = 1;
40+ 
41+// 按图原型顺序：parameters / gradient / learning_rate / accum / momentum / stat
42+ge::graphStatus RunSgdInferShape(gert::StorageShape& bigShape, gert::StorageShape& scalarShape, float dampening,
43+ float weightDecay, bool nesterov)
44+{
45+ auto inferShapeFunc = gert::OpImplRegistry::GetInstance().GetOpImpl("SGD")->infer_shape;
46+ if (inferShapeFunc == nullptr) {
47+ return ge::GRAPH_FAILED;
48+ }
49+ auto holder = gert::InferShapeContextFaker()
50+ .NodeIoNum(SGD_INPUT_NUM, SGD_OUTPUT_NUM)
51+ .IrInstanceNum({1, 1, 1, 1, 1, 1})
52+ .InputShapes({&bigShape, &bigShape, &scalarShape, &bigShape, &scalarShape, &bigShape})
53+ .OutputShapes({&bigShape})
54+ .NodeAttrs({{"dampening", Ops::NN::AnyValue::CreateFrom<float>(dampening)},
55+ {"weight_decay", Ops::NN::AnyValue::CreateFrom<float>(weightDecay)},
56+ {"nesterov", Ops::NN::AnyValue::CreateFrom<bool>(nesterov)}})
57+ .Build();
58+ return inferShapeFunc(holder.GetContext<gert::InferShapeContext>());
59+}
60+} // namespace
61+ 
62+TEST_F(SGD, sgd_infershape_registered)
63+{
64+ // op type 必须是全大写 SGD —— 与 canndev 的 GE op type 及 ini 段名 [SGD] 一致
65+ ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl("SGD"), nullptr);
66+ auto impl = gert::OpImplRegistry::GetInstance().GetOpImpl("SGD");
67+ ASSERT_NE(impl->infer_shape, nullptr);
68+ // 双挂校验：InferDataType 必须一并注册，漏挂会让 GE 侧输出 dtype 推导缺失
69+ ASSERT_NE(impl->infer_datatype, nullptr);
70+}
71+ 
72+TEST_F(SGD, sgd_infershape_normal_rank2)
73+{
74+ gert::StorageShape bigShape = {{96, 256}, {96, 256}};
75+ gert::StorageShape scalarShape = {{1}, {1}};
76+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, 0.0f, false), ge::GRAPH_SUCCESS);
77+}
78+ 
79+TEST_F(SGD, sgd_infershape_normal_rank1_min_boundary)
80+{
81+ gert::StorageShape bigShape = {{33}, {33}};
82+ gert::StorageShape scalarShape = {{1}, {1}};
83+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, 0.0f, false), ge::GRAPH_SUCCESS);
84+}
85+ 
86+TEST_F(SGD, sgd_infershape_normal_rank8_max_boundary)
87+{
88+ gert::StorageShape bigShape = {{2, 2, 2, 2, 2, 2, 2, 2}, {2, 2, 2, 2, 2, 2, 2, 2}};
89+ gert::StorageShape scalarShape = {{1}, {1}};
90+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, 0.0f, false), ge::GRAPH_SUCCESS);
91+}
92+ 
93+TEST_F(SGD, sgd_infershape_unknown_rank_passthrough)
94+{
95+ // UNKNOWN_RANK 在 GE 下表现为 dims == {-2}，必须【透传】而非按 rank 拒绝
96+ gert::StorageShape bigShape = {{-2}, {-2}};
97+ gert::StorageShape scalarShape = {{1}, {1}};
98+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, 0.0f, false), ge::GRAPH_SUCCESS);
99+}
100+ 
101+TEST_F(SGD, sgd_infershape_rank0_rejected)
102+{
103+ // rank-0 标量被拒 —— 对齐 910B/910C（canndev var_dims.size() == 0 判非法）
104+ gert::StorageShape bigShape = {{}, {}};
105+ gert::StorageShape scalarShape = {{1}, {1}};
106+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, 0.0f, false), ge::GRAPH_FAILED);
107+}
108+ 
109+TEST_F(SGD, sgd_infershape_rank9_rejected)
110+{
111+ // rank > 8 被拒（kMaxDimNum = 8）
112+ gert::StorageShape bigShape = {{2, 2, 2, 2, 2, 2, 2, 2, 2}, {2, 2, 2, 2, 2, 2, 2, 2, 2}};
113+ gert::StorageShape scalarShape = {{1}, {1}};
114+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, 0.0f, false), ge::GRAPH_FAILED);
115+}
116+ 
117+TEST_F(SGD, sgd_infershape_nesterov_with_nonzero_dampening_rejected)
118+{
119+ // nesterov == true 时 dampening 必须为 0
120+ gert::StorageShape bigShape = {{96, 256}, {96, 256}};
121+ gert::StorageShape scalarShape = {{1}, {1}};
122+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.5f, 0.0f, true), ge::GRAPH_FAILED);
123+}
124+ 
125+TEST_F(SGD, sgd_infershape_nesterov_with_zero_dampening_ok)
126+{
127+ gert::StorageShape bigShape = {{96, 256}, {96, 256}};
128+ gert::StorageShape scalarShape = {{1}, {1}};
129+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, 0.01f, true), ge::GRAPH_SUCCESS);
130+}
131+ 
132+TEST_F(SGD, sgd_infershape_negative_weight_decay_rejected)
133+{
134+ // weight_decay 必须 >= 0
135+ gert::StorageShape bigShape = {{96, 256}, {96, 256}};
136+ gert::StorageShape scalarShape = {{1}, {1}};
137+ ASSERT_EQ(RunSgdInferShape(bigShape, scalarShape, 0.0f, -0.01f, false), ge::GRAPH_FAILED);
138+}
139+ 
140+TEST_F(SGD, sgd_inferdatatype_follows_parameters)
141+{
142+ auto impl = gert::OpImplRegistry::GetInstance().GetOpImpl("SGD");
143+ ASSERT_NE(impl, nullptr);
144+ auto inferDataTypeFunc = impl->infer_datatype;
145+ ASSERT_NE(inferDataTypeFunc, nullptr);
146+ 
147+ // 唯一图输出 parameters 的 dtype 等于输入 parameters；accum / stat 不是图输出
148+ auto holder = gert::InferDataTypeContextFaker()
149+ .NodeIoNum(SGD_INPUT_NUM, SGD_OUTPUT_NUM)
150+ .IrInstanceNum({1, 1, 1, 1, 1, 1})
151+ .NodeInputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
152+ .NodeInputTd(1, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
153+ .NodeInputTd(2, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
154+ .NodeInputTd(3, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
155+ .NodeInputTd(4, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
156+ .NodeInputTd(5, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
157+ .NodeOutputTd(0, ge::DT_BF16, ge::FORMAT_ND, ge::FORMAT_ND)
158+ .Build();
159+ auto context = holder.GetContext<gert::InferDataTypeContext>();
160+ ASSERT_EQ(inferDataTypeFunc(context), ge::GRAPH_SUCCESS);
161+ ASSERT_EQ(context->GetOutputDataType(0), ge::DT_BF16);
162+}