Star1020
1931
代码
代码
Issues198
Pull Requests343
Actions
项目讨论
Wiki
分析
Star1020
1931
已合并
[CANNbot]generate ApplyCenteredRMSProp operator for Ascend950 #4187
gxx_gitcode创建于 4月24日
[CANNbot]generate ApplyCenteredRMSProp operator for Ascend950 #4187
已合并
gxx_gitcode创建于 4月24日
从gxx_gitcode/ops-nnclaude_applycenteredrmsprop_experimental合入到cann/ops-nnmaster
共 11 个文件变更+1531-0
Aexperimental/optim/apply_centered_rms_prop/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+# NOTE: Portions of this code were AI-generated and have been
11+# technically reviewed for functional accuracy and security
12+ 
13+# 设置算子定义时支持的芯片类型
14+set(SUPPORT_COMPUTE_UNIT "ascend950")
15+# 平铺后 op_host 下不再有 arch 子目录，tiling 文件直接位于 op_host/ 下
16+set(SUPPORT_TILING_DIR "${CMAKE_CURRENT_SOURCE_DIR}/op_host")
17+ 
18+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE apply_centered_rms_prop ACLNNTYPE aclnn COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
Aexperimental/optim/apply_centered_rms_prop/README.md+168-0
@@ -0,0 +1,168 @@
1+# ApplyCenteredRMSProp
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+## 功能说明
15+ 
16+- 算子功能：ApplyCenteredRMSProp 是带"中心化"修正的 RMSProp 优化器算子，功能对标 `tf.raw_ops.ApplyCenteredRMSProp`。在 RMSProp 基础上，额外维护一阶梯度指数移动平均 `mg`，并以 `ms - mg^2` 作为方差估计参与归一化，从而获得更稳定的步长。`var`/`mg`/`ms`/`mom` 均为 Ref Tensor，算子执行后**原地更新**。
17+- 计算公式：
18+ 
19+ $$
20+ \begin{aligned}
21+ mg_t &= \rho \cdot mg_{t-1} + (1 - \rho) \cdot \text{grad}_t \\
22+ ms_t &= \rho \cdot ms_{t-1} + (1 - \rho) \cdot \text{grad}_t^2 \\
23+ denom_t &= \sqrt{ms_t - mg_t^2 + \epsilon} \\
24+ mom_t &= \text{momentum} \cdot mom_{t-1} + \text{lr} \cdot \frac{\text{grad}_t}{denom_t} \\
25+ var_t &= var_{t-1} - mom_t
26+ \end{aligned}
27+ $$
28+ 
29+- 说明：
30+ - `var`/`mg`/`ms`/`mom` 为 Ref Tensor，与对应的 `*_out` 输出共享存储以实现 inplace 更新。
31+ - `lr`、`rho`、`momentum`、`epsilon` 为 0-D 或 1 元素 1-D 标量 Tensor。
32+ - 逐元素独立计算，无跨元素/跨核依赖。
33+ 
34+## 参数说明
35+ 
36+<table style="table-layout: fixed; width: 1576px"><colgroup>
37+<col style="width: 150px">
38+<col style="width: 150px">
39+<col style="width: 420px">
40+<col style="width: 200px">
41+<col style="width: 140px">
42+</colgroup>
43+<thead>
44+ <tr>
45+ <th>参数名</th>
46+ <th>输入/输出/属性</th>
47+ <th>描述</th>
48+ <th>数据类型</th>
49+ <th>数据格式</th>
50+ </tr></thead>
51+<tbody>
52+ <tr>
53+ <td>var</td>
54+ <td>输入</td>
55+ <td>公式中的 var，待更新的模型参数（Ref Tensor，原地更新）。shape 与 mg/ms/mom/grad 一致。</td>
56+ <td>FLOAT16, FLOAT</td>
57+ <td>ND</td>
58+ </tr>
59+ <tr>
60+ <td>mg</td>
61+ <td>输入</td>
62+ <td>公式中的 mg，一阶梯度指数移动平均（Ref Tensor，原地更新）。shape 与 var/ms/mom/grad 一致。</td>
63+ <td>FLOAT16, FLOAT</td>
64+ <td>ND</td>
65+ </tr>
66+ <tr>
67+ <td>ms</td>
68+ <td>输入</td>
69+ <td>公式中的 ms，二阶梯度指数移动平均（Ref Tensor，原地更新）。shape 与 var/mg/mom/grad 一致。</td>
70+ <td>FLOAT16, FLOAT</td>
71+ <td>ND</td>
72+ </tr>
73+ <tr>
74+ <td>mom</td>
75+ <td>输入</td>
76+ <td>公式中的 mom，动量项（Ref Tensor，原地更新）。shape 与 var/mg/ms/grad 一致。</td>
77+ <td>FLOAT16, FLOAT</td>
78+ <td>ND</td>
79+ </tr>
80+ <tr>
81+ <td>lr</td>
82+ <td>输入</td>
83+ <td>公式中的 lr，学习率。0-D 或 1 元素 1-D Tensor。</td>
84+ <td>FLOAT16, FLOAT</td>
85+ <td>ND</td>
86+ </tr>
87+ <tr>
88+ <td>rho</td>
89+ <td>输入</td>
90+ <td>公式中的 rho，指数衰减系数。0-D 或 1 元素 1-D Tensor。</td>
91+ <td>FLOAT16, FLOAT</td>
92+ <td>ND</td>
93+ </tr>
94+ <tr>
95+ <td>momentum</td>
96+ <td>输入</td>
97+ <td>公式中的 momentum，动量系数。0-D 或 1 元素 1-D Tensor。</td>
98+ <td>FLOAT16, FLOAT</td>
99+ <td>ND</td>
100+ </tr>
101+ <tr>
102+ <td>epsilon</td>
103+ <td>输入</td>
104+ <td>公式中的 epsilon，数值稳定项（&gt; 0）。0-D 或 1 元素 1-D Tensor。</td>
105+ <td>FLOAT16, FLOAT</td>
106+ <td>ND</td>
107+ </tr>
108+ <tr>
109+ <td>grad</td>
110+ <td>输入</td>
111+ <td>公式中的 grad，当前步的梯度张量。shape 与 var/mg/ms/mom 一致。</td>
112+ <td>FLOAT16, FLOAT</td>
113+ <td>ND</td>
114+ </tr>
115+ <tr>
116+ <td>var_out</td>
117+ <td>输出</td>
118+ <td>更新后的参数，与输入 var 共享存储（inplace 更新）。</td>
119+ <td>FLOAT16, FLOAT</td>
120+ <td>ND</td>
121+ </tr>
122+ <tr>
123+ <td>mg_out</td>
124+ <td>输出</td>
125+ <td>更新后的一阶梯度均值，与输入 mg 共享存储（inplace 更新）。</td>
126+ <td>FLOAT16, FLOAT</td>
127+ <td>ND</td>
128+ </tr>
129+ <tr>
130+ <td>ms_out</td>
131+ <td>输出</td>
132+ <td>更新后的二阶梯度均值，与输入 ms 共享存储（inplace 更新）。</td>
133+ <td>FLOAT16, FLOAT</td>
134+ <td>ND</td>
135+ </tr>
136+ <tr>
137+ <td>mom_out</td>
138+ <td>输出</td>
139+ <td>更新后的动量项，与输入 mom 共享存储（inplace 更新）。</td>
140+ <td>FLOAT16, FLOAT</td>
141+ <td>ND</td>
142+ </tr>
143+</tbody></table>
144+ 
145+## 约束说明
146+ 
147+- 仅支持 <term>Ascend 950PR/Ascend 950DT</term>（arch35 / DAV_3510），不适配其他芯片代际。
148+- 支持 `float16` 与 `float32` 数据类型，所有输入 Tensor 的 dtype 必须一致。
149+- `var`、`mg`、`ms`、`mom`、`grad` 五者 shape 必须完全一致，且均为连续排布的 ND Tensor。
150+- `lr`、`rho`、`momentum`、`epsilon` 必须为 0-D 或 1 元素 1-D 的标量 Tensor。
151+- 调用方需保证 `epsilon > 0`、`ms - mg^2 + epsilon > 0`；当 `denom == 0` 时 `rsqrt` 输出 Inf/NaN，行为与 PyTorch / TensorFlow 原生实现一致，需由上游调用方规避。
152+- `var`/`mg`/`ms`/`mom` 为 Ref Tensor，Host aclnn 侧必须显式构造四个占位输出 Tensor（`var_out`/`mg_out`/`ms_out`/`mom_out`），并与各自输入共享 Device 地址以保证 inplace 语义。
153+ 
154+## 调用说明
155+ 
156+<table><thead>
157+ <tr>
158+ <th>调用方式</th>
159+ <th>调用样例</th>
160+ <th>说明</th>
161+ </tr></thead>
162+<tbody>
163+ <tr>
164+ <td>aclnn 调用</td>
165+ <td><a href="./examples/arch35/test_aclnn_apply_centered_rms_prop.cpp">test_aclnn_apply_centered_rms_prop</a></td>
166+ <td>Ascend 950 上通过 aclnn 两段式接口 <code>aclnnApplyCenteredRMSPropGetWorkspaceSize</code> → <code>aclnnApplyCenteredRMSProp</code> 调用。<code>var_out</code>/<code>mg_out</code>/<code>ms_out</code>/<code>mom_out</code> 需与对应 Ref 输入共享 Device 地址以保证 inplace 更新。</td>
167+ </tr>
168+</tbody></table>
Aexperimental/optim/apply_centered_rms_prop/examples/arch35/test_aclnn_apply_centered_rms_prop.cpp+216-0
@@ -0,0 +1,216 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+// Minimal two-phase ACLNN call demo for aclnnApplyCenteredRMSProp.
17+//
18+// Scenario: shape=[16], dtype=float32. Performs one Centered RMSProp step
19+// in-place on (var, mg, ms, mom) and verifies against a CPU golden.
20+//
21+// Build & run via examples/run.sh.
22+//
23+// NOTE: The autogen ACLNN signature requires 4 placeholder "out" tensors
24+// (varOutOut/mgOutOut/msOutOut/momOutOut). To observe the in-place update,
25+// each placeholder MUST be created on the SAME device buffer as its Ref input.
26+ 
27+#include <cmath>
28+#include <cstdio>
29+#include <cstdlib>
30+#include <cstring>
31+#include <vector>
32+ 
33+#include "acl/acl.h"
34+#include "aclnn/acl_meta.h"
35+#include "aclnn_apply_centered_rms_prop.h"
36+ 
37+namespace {
38+ 
39+#define CHECK_ACL(expr) \
40+ do { \
41+ aclError __err = (expr); \
42+ if (__err != ACL_SUCCESS) { \
43+ std::fprintf(stderr, "ACL error %d at %s:%d: %s\n", __err, \
44+ __FILE__, __LINE__, #expr); \
45+ std::exit(1); \
46+ } \
47+ } while (0)
48+ 
49+#define CHECK_ACLNN(expr) \
50+ do { \
51+ aclnnStatus __s = (expr); \
52+ if (__s != ACL_SUCCESS) { \
53+ std::fprintf(stderr, "ACLNN error %d at %s:%d: %s\n", __s, \
54+ __FILE__, __LINE__, #expr); \
55+ std::exit(1); \
56+ } \
57+ } while (0)
58+ 
59+void* DevAlloc(size_t bytes) {
60+ void* p = nullptr;
61+ CHECK_ACL(aclrtMalloc(&p, bytes, ACL_MEM_MALLOC_HUGE_FIRST));
62+ return p;
63+}
64+ 
65+aclTensor* MakeTensor(void* ptr, const std::vector<int64_t>& shape,
66+ aclDataType dt) {
67+ std::vector<int64_t> strides(shape.size(), 1);
68+ for (int64_t i = static_cast<int64_t>(shape.size()) - 2; i >= 0; --i) {
69+ strides[i] = strides[i + 1] * shape[i + 1];
70+ }
71+ return aclCreateTensor(
72+ shape.empty() ? nullptr : shape.data(), shape.size(), dt,
73+ strides.empty() ? nullptr : strides.data(), 0, ACL_FORMAT_ND,
74+ shape.empty() ? nullptr : shape.data(), shape.size(), ptr);
75+}
76+ 
77+// CPU golden for one Centered RMSProp step (fp32).
78+struct Golden {
79+ std::vector<float> var, mg, ms, mom;
80+};
81+Golden ComputeGolden(std::vector<float> var, std::vector<float> mg,
82+ std::vector<float> ms, std::vector<float> mom,
83+ const std::vector<float>& grad,
84+ float lr, float rho, float momentum, float epsilon) {
85+ size_t n = var.size();
86+ for (size_t i = 0; i < n; ++i) {
87+ mg[i] = rho * mg[i] + (1.0f - rho) * grad[i];
88+ ms[i] = rho * ms[i] + (1.0f - rho) * grad[i] * grad[i];
89+ float denom = std::sqrt(ms[i] - mg[i] * mg[i] + epsilon);
90+ mom[i] = momentum * mom[i] + lr * grad[i] / denom;
91+ var[i] = var[i] - mom[i];
92+ }
93+ return {std::move(var), std::move(mg), std::move(ms), std::move(mom)};
94+}
95+ 
96+bool Compare(const std::vector<float>& actual, const std::vector<float>& expect,
97+ const char* name) {
98+ constexpr float kRtol = 1e-4f;
99+ constexpr float kAtol = 1e-5f;
100+ int passed = 0;
101+ for (size_t i = 0; i < actual.size(); ++i) {
102+ float diff = std::fabs(actual[i] - expect[i]);
103+ float thr = kAtol + kRtol * std::fabs(expect[i]);
104+ if (diff <= thr) ++passed;
105+ }
106+ std::printf("Result (%s): %d/%zu passed\n", name, passed, actual.size());
107+ return passed == static_cast<int>(actual.size());
108+}
109+ 
110+} // namespace
111+ 
112+int main() {
113+ // ---- 1. ACL init / device / stream ----
114+ CHECK_ACL(aclInit(nullptr));
115+ CHECK_ACL(aclrtSetDevice(0));
116+ aclrtStream stream = nullptr;
117+ CHECK_ACL(aclrtCreateStream(&stream));
118+ 
119+ // ---- 2. Prepare host data ----
120+ const int64_t N = 16;
121+ std::vector<int64_t> shape = {N};
122+ 
123+ std::vector<float> var(N), mg(N), ms(N), mom(N), grad(N);
124+ for (int64_t i = 0; i < N; ++i) {
125+ var[i] = 0.10f * (i + 1);
126+ mg[i] = 0.01f * (i + 1);
127+ ms[i] = 0.20f + 0.01f * i; // > 0
128+ mom[i] = 0.05f * (i + 1);
129+ grad[i] = 0.02f * ((i % 5) - 2); // small mixed-sign
130+ }
131+ float lr = 1e-2f, rho = 0.9f, momentum = 0.9f, epsilon = 1e-6f;
132+ 
133+ // CPU golden BEFORE H2D (var/mg/ms/mom are in-place updated on device).
134+ auto golden = ComputeGolden(var, mg, ms, mom, grad, lr, rho, momentum, epsilon);
135+ 
136+ // ---- 3. Device buffers ----
137+ size_t main_bytes = N * sizeof(float);
138+ size_t scal_bytes = sizeof(float);
139+ void* d_var = DevAlloc(main_bytes);
140+ void* d_mg = DevAlloc(main_bytes);
141+ void* d_ms = DevAlloc(main_bytes);
142+ void* d_mom = DevAlloc(main_bytes);
143+ void* d_grad = DevAlloc(main_bytes);
144+ void* d_lr = DevAlloc(scal_bytes);
145+ void* d_rho = DevAlloc(scal_bytes);
146+ void* d_momc = DevAlloc(scal_bytes);
147+ void* d_eps = DevAlloc(scal_bytes);
148+ 
149+ CHECK_ACL(aclrtMemcpy(d_var, main_bytes, var.data(), main_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
150+ CHECK_ACL(aclrtMemcpy(d_mg, main_bytes, mg.data(), main_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
151+ CHECK_ACL(aclrtMemcpy(d_ms, main_bytes, ms.data(), main_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
152+ CHECK_ACL(aclrtMemcpy(d_mom, main_bytes, mom.data(), main_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
153+ CHECK_ACL(aclrtMemcpy(d_grad, main_bytes, grad.data(), main_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
154+ CHECK_ACL(aclrtMemcpy(d_lr, scal_bytes, &lr, scal_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
155+ CHECK_ACL(aclrtMemcpy(d_rho, scal_bytes, &rho, scal_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
156+ CHECK_ACL(aclrtMemcpy(d_momc, scal_bytes, &momentum, scal_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
157+ CHECK_ACL(aclrtMemcpy(d_eps, scal_bytes, &epsilon, scal_bytes, ACL_MEMCPY_HOST_TO_DEVICE));
158+ 
159+ // ---- 4. Build aclTensors ----
160+ std::vector<int64_t> scalar_shape; // 0-D
161+ aclTensor* t_var = MakeTensor(d_var, shape, ACL_FLOAT);
162+ aclTensor* t_mg = MakeTensor(d_mg, shape, ACL_FLOAT);
163+ aclTensor* t_ms = MakeTensor(d_ms, shape, ACL_FLOAT);
164+ aclTensor* t_mom = MakeTensor(d_mom, shape, ACL_FLOAT);
165+ aclTensor* t_lr = MakeTensor(d_lr, scalar_shape, ACL_FLOAT);
166+ aclTensor* t_rho = MakeTensor(d_rho, scalar_shape, ACL_FLOAT);
167+ aclTensor* t_momc = MakeTensor(d_momc, scalar_shape, ACL_FLOAT);
168+ aclTensor* t_eps = MakeTensor(d_eps, scalar_shape, ACL_FLOAT);
169+ aclTensor* t_grad = MakeTensor(d_grad, shape, ACL_FLOAT);
170+ // Out placeholders alias SAME device storage (in-place semantics).
171+ aclTensor* t_var_o = MakeTensor(d_var, shape, ACL_FLOAT);
172+ aclTensor* t_mg_o = MakeTensor(d_mg, shape, ACL_FLOAT);
173+ aclTensor* t_ms_o = MakeTensor(d_ms, shape, ACL_FLOAT);
174+ aclTensor* t_mom_o = MakeTensor(d_mom, shape, ACL_FLOAT);
175+ 
176+ // ---- 5. Two-phase ACLNN call ----
177+ uint64_t ws_size = 0;
178+ aclOpExecutor* executor = nullptr;
179+ CHECK_ACLNN(aclnnApplyCenteredRMSPropGetWorkspaceSize(
180+ t_var, t_mg, t_ms, t_mom, t_lr, t_rho, t_momc, t_eps, t_grad,
181+ t_var_o, t_mg_o, t_ms_o, t_mom_o, &ws_size, &executor));
182+ 
183+ void* ws_ptr = nullptr;
184+ if (ws_size > 0) ws_ptr = DevAlloc(ws_size);
185+ 
186+ CHECK_ACLNN(aclnnApplyCenteredRMSProp(ws_ptr, ws_size, executor, stream));
187+ CHECK_ACL(aclrtSynchronizeStream(stream));
188+ 
189+ // ---- 6. D2H + compare ----
190+ std::vector<float> out_var(N), out_mg(N), out_ms(N), out_mom(N);
191+ CHECK_ACL(aclrtMemcpy(out_var.data(), main_bytes, d_var, main_bytes, ACL_MEMCPY_DEVICE_TO_HOST));
192+ CHECK_ACL(aclrtMemcpy(out_mg.data(), main_bytes, d_mg, main_bytes, ACL_MEMCPY_DEVICE_TO_HOST));
193+ CHECK_ACL(aclrtMemcpy(out_ms.data(), main_bytes, d_ms, main_bytes, ACL_MEMCPY_DEVICE_TO_HOST));
194+ CHECK_ACL(aclrtMemcpy(out_mom.data(), main_bytes, d_mom, main_bytes, ACL_MEMCPY_DEVICE_TO_HOST));
195+ 
196+ bool ok = true;
197+ ok &= Compare(out_var, golden.var, "var");
198+ ok &= Compare(out_mg, golden.mg, "mg");
199+ ok &= Compare(out_ms, golden.ms, "ms");
200+ ok &= Compare(out_mom, golden.mom, "mom");
201+ std::printf("%s\n", ok ? "ALL PASS" : "FAIL");
202+ 
203+ // ---- 7. Cleanup ----
204+ for (aclTensor* t : {t_var, t_mg, t_ms, t_mom, t_lr, t_rho, t_momc, t_eps,
205+ t_grad, t_var_o, t_mg_o, t_ms_o, t_mom_o}) {
206+ if (t) aclDestroyTensor(t);
207+ }
208+ if (ws_ptr) aclrtFree(ws_ptr);
209+ for (void* p : {d_var, d_mg, d_ms, d_mom, d_grad, d_lr, d_rho, d_momc, d_eps}) {
210+ aclrtFree(p);
211+ }
212+ aclrtDestroyStream(stream);
213+ aclrtResetDevice(0);
214+ aclFinalize();
215+ return ok ? 0 : 1;
216+}
Aexperimental/optim/apply_centered_rms_prop/op_host/apply_centered_rms_prop_def.cpp+139-0
@@ -0,0 +1,139 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+/**
17+ * \file apply_centered_rms_prop_def.cpp
18+ * \brief ApplyCenteredRMSProp operator definition (iteration-1 skeleton).
19+ *
20+ * Optimizer-style operator: 9 inputs (4 Ref + 4 scalar + 1 grad), 4 inplace outputs.
21+ *
22+ * Inputs (in order):
23+ * var (Tensor, float16/float32, ND, shape == mg/ms/mom/grad) -- inplace
24+ * mg (Tensor, float16/float32, ND, shape == var/ms/mom/grad) -- inplace
25+ * ms (Tensor, float16/float32, ND, shape == var/mg/mom/grad) -- inplace
26+ * mom (Tensor, float16/float32, ND, shape == var/mg/ms/grad) -- inplace
27+ * lr (Tensor, float16/float32, ND, 0-D or 1-element 1-D) -- scalar
28+ * rho (Tensor, float16/float32, ND, 0-D or 1-element 1-D) -- scalar
29+ * momentum (Tensor, float16/float32, ND, 0-D or 1-element 1-D) -- scalar
30+ * epsilon (Tensor, float16/float32, ND, 0-D or 1-element 1-D) -- scalar
31+ * grad (Tensor, float16/float32, ND, shape == var/mg/ms/mom)
32+ *
33+ * Outputs (4, inplace via shared storage at L2 layer):
34+ * var_out / mg_out / ms_out / mom_out -- shares storage with input ref tensors
35+ *
36+ * Target: Ascend950 (arch35 / DAV_3510) only.
37+ *
38+ * Iteration-1 skeleton: registers float16 + float32 dtype pair (fp16 path
39+ * implemented first, fp32 path stubbed out in kernel for iteration-2).
40+ */
41+#include "register/op_def_registry.h"
42+ 
43+namespace ops {
44+class ApplyCenteredRMSProp : public OpDef {
45+public:
46+ explicit ApplyCenteredRMSProp(const char* name) : OpDef(name)
47+ {
48+ // --- Inputs (9) --------------------------------------------------
49+ this->Input("var")
50+ .ParamType(REQUIRED)
51+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
52+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
53+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
54+ .AutoContiguous();
55+ this->Input("mg")
56+ .ParamType(REQUIRED)
57+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
58+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
59+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
60+ .AutoContiguous();
61+ this->Input("ms")
62+ .ParamType(REQUIRED)
63+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
64+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
65+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
66+ .AutoContiguous();
67+ this->Input("mom")
68+ .ParamType(REQUIRED)
69+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
70+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
71+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
72+ .AutoContiguous();
73+ this->Input("lr")
74+ .ParamType(REQUIRED)
75+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
76+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
77+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
78+ this->Input("rho")
79+ .ParamType(REQUIRED)
80+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
81+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
82+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
83+ this->Input("momentum")
84+ .ParamType(REQUIRED)
85+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
86+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
87+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
88+ this->Input("epsilon")
89+ .ParamType(REQUIRED)
90+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
91+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
92+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
93+ this->Input("grad")
94+ .ParamType(REQUIRED)
95+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
96+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
97+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
98+ .AutoContiguous();
99+ 
100+ // --- Outputs (4, inplace via shared storage at L2 layer) ----------
101+ this->Output("var_out")
102+ .ParamType(REQUIRED)
103+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
104+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
105+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
106+ .AutoContiguous();
107+ this->Output("mg_out")
108+ .ParamType(REQUIRED)
109+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
110+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
111+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
112+ .AutoContiguous();
113+ this->Output("ms_out")
114+ .ParamType(REQUIRED)
115+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
116+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
117+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
118+ .AutoContiguous();
119+ this->Output("mom_out")
120+ .ParamType(REQUIRED)
121+ .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
122+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
123+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND})
124+ .AutoContiguous();
125+ 
126+ // --- Ascend950 (arch35) AI Core configuration ---------------------
127+ OpAICoreConfig aiCoreConfig;
128+ aiCoreConfig.DynamicCompileStaticFlag(true)
129+ .DynamicFormatFlag(false)
130+ .DynamicRankSupportFlag(true)
131+ .DynamicShapeSupportFlag(true)
132+ .NeedCheckSupportFlag(false)
133+ .PrecisionReduceFlag(false)
134+ .ExtendCfgInfo("opFile.value", "apply_centered_rms_prop");
135+ this->AICore().AddConfig("ascend950", aiCoreConfig);
136+ }
137+};
138+OP_ADD(ApplyCenteredRMSProp);
139+} // namespace ops
Aexperimental/optim/apply_centered_rms_prop/op_host/apply_centered_rms_prop_infershape.cpp+88-0
@@ -0,0 +1,88 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+/**
17+ * \file apply_centered_rms_prop_infershape.cpp
18+ * \brief ApplyCenteredRMSProp shape / dtype inference.
19+ *
20+ * - var_out.shape = var.shape
21+ * - mg_out.shape = mg.shape
22+ * - ms_out.shape = ms.shape
23+ * - mom_out.shape = mom.shape
24+ * - outputs dtype follow inputs (float16 or float32)
25+ */
26+ 
27+#include "register/op_impl_registry.h"
28+#include "exe_graph/runtime/infer_shape_context.h"
29+#include "exe_graph/runtime/infer_datatype_context.h"
30+ 
31+using namespace ge;
32+ 
33+namespace ops {
34+ 
35+static ge::graphStatus InferShape4ApplyCenteredRMSProp(gert::InferShapeContext* context)
36+{
37+ // Input 0 = var -> Output 0 = var_out
38+ const gert::Shape* varShape = context->GetInputShape(0);
39+ if (varShape == nullptr) {
40+ return ge::GRAPH_FAILED;
41+ }
42+ // Input 1 = mg -> Output 1 = mg_out
43+ const gert::Shape* mgShape = context->GetInputShape(1);
44+ if (mgShape == nullptr) {
45+ return ge::GRAPH_FAILED;
46+ }
47+ // Input 2 = ms -> Output 2 = ms_out
48+ const gert::Shape* msShape = context->GetInputShape(2);
49+ if (msShape == nullptr) {
50+ return ge::GRAPH_FAILED;
51+ }
52+ // Input 3 = mom -> Output 3 = mom_out
53+ const gert::Shape* momShape = context->GetInputShape(3);
54+ if (momShape == nullptr) {
55+ return ge::GRAPH_FAILED;
56+ }
57+ 
58+ gert::Shape* varOutShape = context->GetOutputShape(0);
59+ gert::Shape* mgOutShape = context->GetOutputShape(1);
60+ gert::Shape* msOutShape = context->GetOutputShape(2);
61+ gert::Shape* momOutShape = context->GetOutputShape(3);
62+ if (varOutShape == nullptr || mgOutShape == nullptr ||
63+ msOutShape == nullptr || momOutShape == nullptr) {
64+ return ge::GRAPH_FAILED;
65+ }
66+ 
67+ *varOutShape = *varShape;
68+ *mgOutShape = *mgShape;
69+ *msOutShape = *msShape;
70+ *momOutShape = *momShape;
71+ return ge::GRAPH_SUCCESS;
72+}
73+ 
74+static ge::graphStatus InferDataType4ApplyCenteredRMSProp(gert::InferDataTypeContext* context)
75+{
76+ // Outputs dtype follow their corresponding Ref input.
77+ context->SetOutputDataType(0, context->GetInputDataType(0));
78+ context->SetOutputDataType(1, context->GetInputDataType(1));
79+ context->SetOutputDataType(2, context->GetInputDataType(2));
80+ context->SetOutputDataType(3, context->GetInputDataType(3));
81+ return ge::GRAPH_SUCCESS;
82+}
83+ 
84+IMPL_OP_INFERSHAPE(ApplyCenteredRMSProp)
85+ .InferShape(InferShape4ApplyCenteredRMSProp)
86+ .InferDataType(InferDataType4ApplyCenteredRMSProp);
87+ 
88+} // namespace ops
Aexperimental/optim/apply_centered_rms_prop/op_host/apply_centered_rms_prop_tiling.cpp+278-0
@@ -0,0 +1,278 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+/**
17+ * \file apply_centered_rms_prop_tiling.cpp
18+ * \brief ApplyCenteredRMSProp tiling (arch35).
19+ *
20+ * Tiling strategy (DESIGN §5):
21+ * 1. Multi-core: split total elements evenly across AI vector cores.
22+ * blockFactor = CeilAlign(CeilDiv(total, coreNum), ubBlockSize)
23+ * so neighbouring cores do not trample each other's GM output.
24+ * 2. UB: ubFactor = floor((UB_size - reserve) / (N * compute_dtype_size))
25+ * where N depends on dtype path (DESIGN §5.2):
26+ * - fp32: 8 fp32 buffers
27+ * - fp16: 12 mixed buffers (5 fp16 queues x 2B + 7 fp32 tmp x 4B)
28+ * -> Conservatively use 12 * 4B = 48B / element (upper bound).
29+ * Rounded down to 32-byte granularity.
30+ *
31+ * TilingKey (DESIGN §5.4):
32+ * dtype fp32 -> TilingKey 1
33+ * dtype fp16 -> TilingKey 2
34+ * selected via ASCENDC_TPL_SEL_PARAM(context, dTypeVar).
35+ */
36+ 
37+#include "register/op_def_registry.h"
38+#include "op_common/log/log.h"
39+#include "op_common/op_host/util/math_util.h"
40+#include "op_common/op_host/util/platform_util.h"
41+#include "../op_kernel/apply_centered_rms_prop_tiling_data.h"
42+#include "../op_kernel/apply_centered_rms_prop_tiling_key.h"
43+ 
44+namespace optiling {
45+ 
46+using Ops::Base::CeilDiv;
47+using Ops::Base::CeilAlign;
48+using Ops::Base::FloorAlign;
49+using Ops::Base::FloorDiv;
50+using Ops::Base::GetUbBlockSize;
51+ 
52+constexpr uint32_t WS_SYS_SIZE = 0U;
53+// Per-tile target element count -- matches apply_proximal_adagrad's 2048
54+// ceiling so the first iteration stays deterministic while leaving UB
55+// headroom for the fp16 12-buffer layout.
56+constexpr int64_t TILE_ELEM_NUM_TARGET = 2048;
57+// UB co-resident buffer accounting (per element, in 4-byte units).
58+//
59+// fp16 path (Kernel InitBuffer): 9 queues x 2 DB x 2B + 8 fp32 TBufs x 4B
60+// = 36B + 32B = 68B / elem -> 17 fp32-units / elem.
61+// fp32 path (Kernel InitBuffer, Suggestion-1 conditional):
62+// 9 queues x 2 DB x 4B + 3 fp32 TBufs x 4B (denom/tmp1/tmp2 only;
63+// the 5 cast-buffers are NOT allocated on fp32 path)
64+// = 72B + 12B = 84B / elem -> 21 fp32-units / elem.
65+// We round up to 21 for the fp32 path to match the kernel InitBuffer footprint.
66+constexpr int64_t UB_BUFFER_COUNT_FP32 = 21;
67+constexpr int64_t UB_BUFFER_COUNT_FP16 = 17;
68+constexpr int64_t SIZE_FP32 = 4;
69+constexpr int64_t SIZE_FP16 = 2;
70+ 
71+static const gert::Shape g_vec_1_shape = {1};
72+ 
73+static inline const gert::Shape EnsureNotScalar(const gert::Shape& in_shape)
74+{
75+ if (in_shape.GetDimNum() == 0) {
76+ return g_vec_1_shape;
77+ }
78+ return in_shape;
79+}
80+ 
81+static ge::graphStatus GetPlatformInfo(gert::TilingContext* context,
82+ uint64_t& ubSize, int64_t& coreNum)
83+{
84+ fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo();
85+ OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr);
86+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
87+ coreNum = ascendcPlatform.GetCoreNumAiv();
88+ OP_CHECK_IF(coreNum == 0, OP_LOGE(context, "coreNum is 0"),
89+ return ge::GRAPH_FAILED);
90+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
91+ OP_CHECK_IF(ubSize == 0, OP_LOGE(context, "ubSize is 0"),
92+ return ge::GRAPH_FAILED);
93+ return ge::GRAPH_SUCCESS;
94+}
95+ 
96+static ge::graphStatus GetShapeInfo(gert::TilingContext* context,
97+ int64_t& totalElements,
98+ ge::DataType& dataType)
99+{
100+ // Input 0 = var; its shape is canonical (var/mg/ms/mom/grad must match).
101+ auto inputVar = context->GetInputShape(0);
102+ OP_CHECK_NULL_WITH_CONTEXT(context, inputVar);
103+ auto varShape = EnsureNotScalar(inputVar->GetStorageShape());
104+ totalElements = varShape.GetShapeSize();
105+ 
106+ auto inputDesc = context->GetInputDesc(0);
107+ OP_CHECK_NULL_WITH_CONTEXT(context, inputDesc);
108+ dataType = inputDesc->GetDataType();
109+ OP_CHECK_IF(dataType != ge::DT_FLOAT && dataType != ge::DT_FLOAT16,
110+ OP_LOGE(context,
111+ "ApplyCenteredRMSProp: only float32 / float16 are "
112+ "supported, got %d",
113+ static_cast<int>(dataType)),
114+ return ge::GRAPH_FAILED);
115+ 
116+ // Tensor-shape consistency: mg/ms/mom/grad must equal var.shape.
117+ // (Issue/Suggestion-4 closure: shape_mismatch must be rejected at host
118+ // tiling instead of relying on autogen L2 checker, which only catches
119+ // dtype mismatches.)
120+ constexpr size_t kMainTensorIdx[] = {1U, 2U, 3U, 8U}; // mg, ms, mom, grad
121+ const char* kMainTensorName[] = {"mg", "ms", "mom", "grad"};
122+ for (size_t i = 0; i < sizeof(kMainTensorIdx) / sizeof(kMainTensorIdx[0]); ++i) {
123+ auto inShape = context->GetInputShape(kMainTensorIdx[i]);
124+ OP_CHECK_NULL_WITH_CONTEXT(context, inShape);
125+ auto s = EnsureNotScalar(inShape->GetStorageShape());
126+ OP_CHECK_IF(
127+ s.GetShapeSize() != totalElements,
128+ OP_LOGE(context,
129+ "ApplyCenteredRMSProp: %s.numel=%ld must equal var.numel=%ld",
130+ kMainTensorName[i],
131+ static_cast<long>(s.GetShapeSize()),
132+ static_cast<long>(totalElements)),
133+ return ge::GRAPH_FAILED);
134+ }
135+ 
136+ // Scalar shape self-defense (Issue-1): lr/rho/momentum/epsilon must be
137+ // 0-D or single-element 1-D (numel ∈ {0, 1}). Empty scalar is rejected
138+ // because LoadScalar reads element[0]. (numel == 0 untreated would lead
139+ // to OOB GetValue in kernel.)
140+ constexpr size_t kScalarIdx[] = {4U, 5U, 6U, 7U}; // lr, rho, momentum, epsilon
141+ const char* kScalarName[] = {"lr", "rho", "momentum", "epsilon"};
142+ for (size_t i = 0; i < sizeof(kScalarIdx) / sizeof(kScalarIdx[0]); ++i) {
143+ auto inShape = context->GetInputShape(kScalarIdx[i]);
144+ OP_CHECK_NULL_WITH_CONTEXT(context, inShape);
145+ const auto& rawShape = inShape->GetStorageShape();
146+ int64_t numel = (rawShape.GetDimNum() == 0) ? 1 : rawShape.GetShapeSize();
147+ OP_CHECK_IF(
148+ numel != 1,
149+ OP_LOGE(context,
150+ "ApplyCenteredRMSProp: scalar %s must be 0-D or 1-element 1-D, "
151+ "got numel=%ld",
152+ kScalarName[i], static_cast<long>(numel)),
153+ return ge::GRAPH_FAILED);
154+ }
155+ 
156+ return ge::GRAPH_SUCCESS;
157+}
158+ 
159+static ge::graphStatus GetWorkspaceSize(gert::TilingContext* context)
160+{
161+ size_t* currentWorkspace = context->GetWorkspaceSizes(1);
162+ OP_CHECK_NULL_WITH_CONTEXT(context, currentWorkspace);
163+ currentWorkspace[0] = WS_SYS_SIZE;
164+ return ge::GRAPH_SUCCESS;
165+}
166+ 
167+static ge::graphStatus ApplyCenteredRMSPropTilingFunc(gert::TilingContext* context)
168+{
169+ // 1. Platform info
170+ uint64_t ubSize = 0;
171+ int64_t coreNum = 0;
172+ OP_CHECK_IF(GetPlatformInfo(context, ubSize, coreNum) != ge::GRAPH_SUCCESS,
173+ OP_LOGE(context, "GetPlatformInfo error"),
174+ return ge::GRAPH_FAILED);
175+ 
176+ // 2. Shape / dtype info
177+ int64_t totalElements = 0;
178+ ge::DataType dataType = ge::DT_FLOAT;
179+ OP_CHECK_IF(GetShapeInfo(context, totalElements, dataType) != ge::GRAPH_SUCCESS,
180+ OP_LOGE(context, "GetShapeInfo error"),
181+ return ge::GRAPH_FAILED);
182+ 
183+ // 3. Workspace
184+ OP_CHECK_IF(GetWorkspaceSize(context) != ge::GRAPH_SUCCESS,
185+ OP_LOGE(context, "GetWorkspaceSize error"),
186+ return ge::GRAPH_FAILED);
187+ 
188+ // 4. Fill TilingData
189+ ApplyCenteredRMSPropTilingData* tiling =
190+ context->GetTilingData<ApplyCenteredRMSPropTilingData>();
191+ OP_CHECK_NULL_WITH_CONTEXT(context, tiling);
192+ OP_CHECK_IF(
193+ memset_s(tiling, sizeof(ApplyCenteredRMSPropTilingData), 0,
194+ sizeof(ApplyCenteredRMSPropTilingData)) != EOK,
195+ OP_LOGE(context, "set tiling data error"),
196+ return ge::GRAPH_FAILED);
197+ 
198+ tiling->totalElements = totalElements;
199+ 
200+ // Empty tensor: run a single idle core so the launcher still succeeds.
201+ if (totalElements == 0) {
202+ tiling->blockFactor = 0;
203+ tiling->ubFactor = 0;
204+ context->SetBlockDim(1);
205+ uint32_t dTypeVar = static_cast<uint32_t>(dataType);
206+ ASCENDC_TPL_SEL_PARAM(context, dTypeVar);
207+ return ge::GRAPH_SUCCESS;
208+ }
209+ 
210+ // ubBlockSize = 32B / sizeof(T).
211+ int64_t ubBlockSize = Ops::Base::GetUbBlockSize(context);
212+ OP_CHECK_IF(ubBlockSize <= 0,
213+ OP_LOGE(context, "invalid ubBlockSize=%ld", ubBlockSize),
214+ return ge::GRAPH_FAILED);
215+ 
216+ // Multi-core split: ceil-aligned to DMA granularity.
217+ int64_t blockFactor = CeilAlign(CeilDiv(totalElements, coreNum), ubBlockSize);
218+ int64_t usedCoreNum = CeilDiv(totalElements, blockFactor);
219+ 
220+ // UB split per dtype path. perElemBytes = sizeof(fp32) for both paths;
221+ // ubBufCount captures the actual per-element 4-byte-unit footprint
222+ // (fp32 path: 26 units; fp16 path: 17 units).
223+ int64_t ubBufCount;
224+ int64_t perElemBytes = SIZE_FP32;
225+ if (dataType == ge::DT_FLOAT) {
226+ ubBufCount = UB_BUFFER_COUNT_FP32;
227+ } else {
228+ ubBufCount = UB_BUFFER_COUNT_FP16;
229+ }
230+ 
231+ int64_t ubCapacityElem =
232+ FloorAlign(FloorDiv(static_cast<int64_t>(ubSize) / perElemBytes,
233+ ubBufCount),
234+ ubBlockSize);
235+ OP_CHECK_IF(ubCapacityElem <= 0,
236+ OP_LOGE(context, "UB too small: ubCapacityElem=%ld",
237+ ubCapacityElem),
238+ return ge::GRAPH_FAILED);
239+ 
240+ int64_t ubFactor = (TILE_ELEM_NUM_TARGET < ubCapacityElem)
241+ ? TILE_ELEM_NUM_TARGET
242+ : ubCapacityElem;
243+ 
244+ // Also cap by blockFactor so a single core does not allocate more UB
245+ // space than it will ever use.
246+ if (ubFactor > blockFactor) {
247+ ubFactor = FloorAlign(blockFactor, ubBlockSize);
248+ if (ubFactor <= 0) {
249+ ubFactor = ubBlockSize;
250+ }
251+ }
252+ 
253+ tiling->blockFactor = blockFactor;
254+ tiling->ubFactor = ubFactor;
255+ 
256+ context->SetBlockDim(usedCoreNum);
257+ 
258+ // 5. TilingKey via ASCENDC_TPL_SEL_PARAM (template-argument mechanism).
259+ // Key is fully encoded by dtype (fp16=KEY2, fp32=KEY1).
260+ uint32_t dTypeVar = static_cast<uint32_t>(dataType);
261+ ASCENDC_TPL_SEL_PARAM(context, dTypeVar);
262+ 
263+ return ge::GRAPH_SUCCESS;
264+}
265+ 
266+static ge::graphStatus TilingParseForApplyCenteredRMSProp(
267+ [[maybe_unused]] gert::TilingParseContext* context)
268+{
269+ return ge::GRAPH_SUCCESS;
270+}
271+ 
272+struct ApplyCenteredRMSPropCompileInfo {};
273+ 
274+IMPL_OP_OPTILING(ApplyCenteredRMSProp)
275+ .Tiling(ApplyCenteredRMSPropTilingFunc)
276+ .TilingParse<ApplyCenteredRMSPropCompileInfo>(TilingParseForApplyCenteredRMSProp);
277+ 
278+} // namespace optiling
Aexperimental/optim/apply_centered_rms_prop/op_kernel/apply_centered_rms_prop.cpp+50-0
@@ -0,0 +1,50 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+/**
17+ * \file apply_centered_rms_prop_arch35.cpp
18+ * \brief ApplyCenteredRMSProp kernel entry (arch35 / Ascend950).
19+ *
20+ * Signature follows registry-invoke convention:
21+ * all inputs (4 Ref + 4 scalar + 1 grad) ->
22+ * all outputs (var_out / mg_out / ms_out / mom_out) ->
23+ * workspace -> tiling
24+ *
25+ * Template parameter:
26+ * - D_T_VAR : input dtype (C_DT_FLOAT16 -> TilingKey 2 / fp16 path,
27+ * C_DT_FLOAT -> TilingKey 1 / fp32 path).
28+ *
29+ * Iteration-1 skeleton implements the fp16 path. The fp32 path is registered
30+ * in the template-arg declaration so the autogen pipeline produces both
31+ * binary slots, but the kernel body for fp32 will land in iteration-2.
32+ */
33+ 
34+#include "apply_centered_rms_prop.h"
35+ 
36+template <typename D_T_VAR>
37+__global__ __aicore__ void apply_centered_rms_prop(
38+ GM_ADDR var, GM_ADDR mg, GM_ADDR ms, GM_ADDR mom,
39+ GM_ADDR lr, GM_ADDR rho, GM_ADDR momentum, GM_ADDR epsilon,
40+ GM_ADDR grad,
41+ GM_ADDR var_out, GM_ADDR mg_out, GM_ADDR ms_out, GM_ADDR mom_out,
42+ GM_ADDR workspace, GM_ADDR tiling)
43+{
44+ REGISTER_TILING_DEFAULT(ApplyCenteredRMSPropTilingData);
45+ GET_TILING_DATA_WITH_STRUCT(ApplyCenteredRMSPropTilingData, tilingData, tiling);
46+ NsApplyCenteredRMSProp::ApplyCenteredRMSProp<D_T_VAR> op;
47+ op.Init(var, mg, ms, mom, lr, rho, momentum, epsilon, grad,
48+ var_out, mg_out, ms_out, mom_out, &tilingData);
49+ op.Process();
50+}
Aexperimental/optim/apply_centered_rms_prop/op_kernel/apply_centered_rms_prop.h+479-0
@@ -0,0 +1,479 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+/**
17+ * \file apply_centered_rms_prop.h
18+ * \brief ApplyCenteredRMSProp kernel class (arch35 / Ascend950).
19+ *
20+ * Implements the Centered RMSProp optimizer step per element:
21+ * mg_new = rho * mg + (1 - rho) * grad
22+ * ms_new = rho * ms + (1 - rho) * grad * grad
23+ * denom = sqrt(ms_new - mg_new * mg_new + epsilon)
24+ * mom_new = momentum * mom + lr * grad / denom
25+ * var_new = var - mom_new
26+ *
27+ * Routes:
28+ * - D_T_VAR = half (TilingKey 2): fp16 path. Cast(fp16->fp32) on CopyIn,
29+ * compute in fp32, Cast(fp32->fp16) on CopyOut.
30+ * - D_T_VAR = float (TilingKey 1): fp32 path. Direct in-place fp32 compute
31+ * (no Cast); reuses denom/tmp1/tmp2 fp32 scratch buffers. Maxs(...,0.0f)
32+ * clamp shared with fp16 path for sqrt-of-negative numerical safety.
33+ */
34+ 
35+#ifndef APPLY_CENTERED_RMS_PROP_H
36+#define APPLY_CENTERED_RMS_PROP_H
37+ 
38+#include "kernel_operator.h"
39+#include "kernel_tiling/kernel_tiling.h"
40+#include "apply_centered_rms_prop_tiling_data.h"
41+#include "apply_centered_rms_prop_tiling_key.h"
42+ 
43+namespace NsApplyCenteredRMSProp {
44+ 
45+using AscendC::TPipe;
46+using AscendC::TQue;
47+using AscendC::TBuf;
48+using AscendC::QuePosition;
49+using AscendC::GlobalTensor;
50+using AscendC::LocalTensor;
51+using AscendC::DataCopyExtParams;
52+using AscendC::DataCopyPad;
53+using AscendC::DataCopyPadExtParams;
54+using AscendC::GetBlockIdx;
55+using AscendC::Add;
56+using AscendC::Sub;
57+using AscendC::Mul;
58+using AscendC::Muls;
59+using AscendC::Adds;
60+using AscendC::Maxs;
61+using AscendC::Div;
62+using AscendC::Sqrt;
63+using AscendC::Cast;
64+using AscendC::RoundMode;
65+ 
66+template <typename T>
67+class ApplyCenteredRMSProp {
68+public:
69+ __aicore__ inline ApplyCenteredRMSProp() {}
70+ 
71+ __aicore__ inline void Init(GM_ADDR var, GM_ADDR mg, GM_ADDR ms, GM_ADDR mom,
72+ GM_ADDR lr, GM_ADDR rho, GM_ADDR momentum,
73+ GM_ADDR epsilon, GM_ADDR grad,
74+ GM_ADDR varOut, GM_ADDR mgOut,
75+ GM_ADDR msOut, GM_ADDR momOut,
76+ const ApplyCenteredRMSPropTilingData* tilingData);
77+ __aicore__ inline void Process();
78+ 
79+private:
80+ __aicore__ inline void CopyInTile(int64_t gmOffset, int64_t currentNum);
81+ __aicore__ inline void Compute(int64_t currentNum);
82+ __aicore__ inline void CopyOutTile(int64_t gmOffset, int64_t currentNum);
83+ 
84+ __aicore__ inline float LoadScalar(const GlobalTensor<T>& src) const;
85+ 
86+ // Suggestion-3: shared DataCopyExtParams builder for tile DMA blocks.
87+ static __aicore__ inline DataCopyExtParams MakeTileCopyParams(int64_t currentNum)
88+ {
89+ DataCopyExtParams p;
90+ p.blockCount = 1;
91+ p.blockLen = static_cast<uint32_t>(currentNum * sizeof(T));
92+ p.srcStride = 0;
93+ p.dstStride = 0;
94+ return p;
95+ }
96+ 
97+private:
98+ TPipe pipe_;
99+ // IN queues: var / mg / ms / mom / grad, double-buffered.
100+ TQue<QuePosition::VECIN, 2> varInQue_;
101+ TQue<QuePosition::VECIN, 2> mgInQue_;
102+ TQue<QuePosition::VECIN, 2> msInQue_;
103+ TQue<QuePosition::VECIN, 2> momInQue_;
104+ TQue<QuePosition::VECIN, 2> gradInQue_;
105+ // OUT queues: var / mg / ms / mom, double-buffered.
106+ TQue<QuePosition::VECOUT, 2> varOutQue_;
107+ TQue<QuePosition::VECOUT, 2> mgOutQue_;
108+ TQue<QuePosition::VECOUT, 2> msOutQue_;
109+ TQue<QuePosition::VECOUT, 2> momOutQue_;
110+ // fp32 scratch buffers (only used on the fp16 path for Cast intermediates,
111+ // and on both paths for compute scratch).
112+ TBuf<QuePosition::VECCALC> varF32Buf_;
113+ TBuf<QuePosition::VECCALC> mgF32Buf_;
114+ TBuf<QuePosition::VECCALC> msF32Buf_;
115+ TBuf<QuePosition::VECCALC> momF32Buf_;
116+ TBuf<QuePosition::VECCALC> gradF32Buf_;
117+ TBuf<QuePosition::VECCALC> denomBuf_;
118+ TBuf<QuePosition::VECCALC> tmp1Buf_;
119+ TBuf<QuePosition::VECCALC> tmp2Buf_;
120+ 
121+ GlobalTensor<T> varGm_;
122+ GlobalTensor<T> mgGm_;
123+ GlobalTensor<T> msGm_;
124+ GlobalTensor<T> momGm_;
125+ GlobalTensor<T> gradGm_;
126+ GlobalTensor<T> lrGm_;
127+ GlobalTensor<T> rhoGm_;
128+ GlobalTensor<T> momentumGm_;
129+ GlobalTensor<T> epsilonGm_;
130+ GlobalTensor<T> varOutGm_;
131+ GlobalTensor<T> mgOutGm_;
132+ GlobalTensor<T> msOutGm_;
133+ GlobalTensor<T> momOutGm_;
134+ 
135+ // Scalars loaded from GM at Init time (always promoted to float for
136+ // intra-kernel use, regardless of dtype path).
137+ float lrScalar_ = 0.0f;
138+ float rhoScalar_ = 0.0f;
139+ float momentumScalar_ = 0.0f;
140+ float epsilonScalar_ = 0.0f;
141+ // Derived: oneMinusRho_ = 1.0f - rho.
142+ float oneMinusRho_ = 0.0f;
143+ 
144+ int64_t blockOffset_ = 0;
145+ int64_t blockLen_ = 0;
146+ int64_t ubFactor_ = 0;
147+};
148+ 
149+// =============================================================================
150+// LoadScalar: read element[0] of a 1-element GM tensor as float.
151+// =============================================================================
152+template <typename T>
153+__aicore__ inline float ApplyCenteredRMSProp<T>::LoadScalar(
154+ const GlobalTensor<T>& src) const
155+{
156+ return static_cast<float>(src.GetValue(0));
157+}
158+ 
159+// =============================================================================
160+// Init
161+// =============================================================================
162+template <typename T>
163+__aicore__ inline void ApplyCenteredRMSProp<T>::Init(
164+ GM_ADDR var, GM_ADDR mg, GM_ADDR ms, GM_ADDR mom,
165+ GM_ADDR lr, GM_ADDR rho, GM_ADDR momentum, GM_ADDR epsilon,
166+ GM_ADDR grad,
167+ GM_ADDR varOut, GM_ADDR mgOut, GM_ADDR msOut, GM_ADDR momOut,
168+ const ApplyCenteredRMSPropTilingData* tilingData)
169+{
170+ ubFactor_ = tilingData->ubFactor;
171+ 
172+ // Empty-tensor / degenerate tiling -> just record zero length; Process()
173+ // will short-circuit before any compute.
174+ if (tilingData->totalElements == 0 || tilingData->blockFactor == 0) {
175+ blockOffset_ = 0;
176+ blockLen_ = 0;
177+ return;
178+ }
179+ 
180+ blockOffset_ = tilingData->blockFactor * static_cast<int64_t>(GetBlockIdx());
181+ int64_t remaining = tilingData->totalElements - blockOffset_;
182+ if (remaining <= 0) {
183+ blockLen_ = 0;
184+ return;
185+ }
186+ blockLen_ = (remaining > tilingData->blockFactor)
187+ ? tilingData->blockFactor
188+ : remaining;
189+ 
190+ // Main vectorised tensors -- slice each core's view.
191+ varGm_.SetGlobalBuffer((__gm__ T*)var + blockOffset_, blockLen_);
192+ mgGm_.SetGlobalBuffer((__gm__ T*)mg + blockOffset_, blockLen_);
193+ msGm_.SetGlobalBuffer((__gm__ T*)ms + blockOffset_, blockLen_);
194+ momGm_.SetGlobalBuffer((__gm__ T*)mom + blockOffset_, blockLen_);
195+ gradGm_.SetGlobalBuffer((__gm__ T*)grad + blockOffset_, blockLen_);
196+ varOutGm_.SetGlobalBuffer((__gm__ T*)varOut + blockOffset_, blockLen_);
197+ mgOutGm_.SetGlobalBuffer((__gm__ T*)mgOut + blockOffset_, blockLen_);
198+ msOutGm_.SetGlobalBuffer((__gm__ T*)msOut + blockOffset_, blockLen_);
199+ momOutGm_.SetGlobalBuffer((__gm__ T*)momOut + blockOffset_, blockLen_);
200+ 
201+ // Scalar tensors: single-element broadcast-read.
202+ lrGm_.SetGlobalBuffer((__gm__ T*)lr, 1);
203+ rhoGm_.SetGlobalBuffer((__gm__ T*)rho, 1);
204+ momentumGm_.SetGlobalBuffer((__gm__ T*)momentum, 1);
205+ epsilonGm_.SetGlobalBuffer((__gm__ T*)epsilon, 1);
206+ lrScalar_ = LoadScalar(lrGm_);
207+ rhoScalar_ = LoadScalar(rhoGm_);
208+ momentumScalar_ = LoadScalar(momentumGm_);
209+ epsilonScalar_ = LoadScalar(epsilonGm_);
210+ oneMinusRho_ = 1.0f - rhoScalar_;
211+ 
212+ // UB buffer allocation -- mixed dtype layout for fp16 path.
213+ pipe_.InitBuffer(varInQue_, 2, ubFactor_ * sizeof(T));
214+ pipe_.InitBuffer(mgInQue_, 2, ubFactor_ * sizeof(T));
215+ pipe_.InitBuffer(msInQue_, 2, ubFactor_ * sizeof(T));
216+ pipe_.InitBuffer(momInQue_, 2, ubFactor_ * sizeof(T));
217+ pipe_.InitBuffer(gradInQue_, 2, ubFactor_ * sizeof(T));
218+ pipe_.InitBuffer(varOutQue_, 2, ubFactor_ * sizeof(T));
219+ pipe_.InitBuffer(mgOutQue_, 2, ubFactor_ * sizeof(T));
220+ pipe_.InitBuffer(msOutQue_, 2, ubFactor_ * sizeof(T));
221+ pipe_.InitBuffer(momOutQue_, 2, ubFactor_ * sizeof(T));
222+ 
223+ // fp32 scratch (also reused on fp32 path -- single buffer each, no DB).
224+ // Suggestion-1 (UB layout optimisation): the 5 fp32 cast-buffers
225+ // (var/mg/ms/mom/grad F32) are only needed on the fp16 path; on the
226+ // fp32 path Compute() reuses denom/tmp1/tmp2 as scratch. Conditionally
227+ // allocate to free ~9 fp32-units / elem for the fp32 path (Tiling sets
228+ // UB_BUFFER_COUNT_FP32 = 17 to match).
229+ if constexpr (std::is_same_v<T, half>) {
230+ pipe_.InitBuffer(varF32Buf_, ubFactor_ * sizeof(float));
231+ pipe_.InitBuffer(mgF32Buf_, ubFactor_ * sizeof(float));
232+ pipe_.InitBuffer(msF32Buf_, ubFactor_ * sizeof(float));
233+ pipe_.InitBuffer(momF32Buf_, ubFactor_ * sizeof(float));
234+ pipe_.InitBuffer(gradF32Buf_, ubFactor_ * sizeof(float));
235+ }
236+ pipe_.InitBuffer(denomBuf_, ubFactor_ * sizeof(float));
237+ pipe_.InitBuffer(tmp1Buf_, ubFactor_ * sizeof(float));
238+ pipe_.InitBuffer(tmp2Buf_, ubFactor_ * sizeof(float));
239+}
240+ 
241+// =============================================================================
242+// CopyInTile: pad-aware DataCopyPad of var / mg / ms / mom / grad.
243+// =============================================================================
244+template <typename T>
245+__aicore__ inline void ApplyCenteredRMSProp<T>::CopyInTile(
246+ int64_t gmOffset, int64_t currentNum)
247+{
248+ LocalTensor<T> varLocal = varInQue_.template AllocTensor<T>();
249+ LocalTensor<T> mgLocal = mgInQue_.template AllocTensor<T>();
250+ LocalTensor<T> msLocal = msInQue_.template AllocTensor<T>();
251+ LocalTensor<T> momLocal = momInQue_.template AllocTensor<T>();
252+ LocalTensor<T> gradLocal = gradInQue_.template AllocTensor<T>();
253+ 
254+ DataCopyExtParams copyParams = MakeTileCopyParams(currentNum);
255+ 
256+ // Right-pad tail lanes ([currentNum, alignedNum)) with safe values:
257+ // - var/mg/mom/grad: pad with 0 (Add/Mul with 0 preserves accumulators).
258+ // - ms: pad with 1 (sqrt(1 - 0 + eps) > 0, avoids div-by-0).
259+ constexpr int64_t kAlignBlock = 32 / sizeof(T);
260+ int64_t alignedNum =
261+ ((currentNum + kAlignBlock - 1) / kAlignBlock) * kAlignBlock;
262+ uint8_t rightPadCount = static_cast<uint8_t>(alignedNum - currentNum);
263+ DataCopyPadExtParams<T> padZero{true, 0, rightPadCount, static_cast<T>(0)};
264+ DataCopyPadExtParams<T> padOne{true, 0, rightPadCount, static_cast<T>(1.0f)};
265+ 
266+ DataCopyPad(varLocal, varGm_[gmOffset], copyParams, padZero);
267+ DataCopyPad(mgLocal, mgGm_[gmOffset], copyParams, padZero);
268+ DataCopyPad(msLocal, msGm_[gmOffset], copyParams, padOne);
269+ DataCopyPad(momLocal, momGm_[gmOffset], copyParams, padZero);
270+ DataCopyPad(gradLocal, gradGm_[gmOffset], copyParams, padZero);
271+ 
272+ varInQue_.EnQue(varLocal);
273+ mgInQue_.EnQue(mgLocal);
274+ msInQue_.EnQue(msLocal);
275+ momInQue_.EnQue(momLocal);
276+ gradInQue_.EnQue(gradLocal);
277+}
278+ 
279+// =============================================================================
280+// CopyOutTile: write back var / mg / ms / mom to their inplace GM slots.
281+// =============================================================================
282+template <typename T>
283+__aicore__ inline void ApplyCenteredRMSProp<T>::CopyOutTile(
284+ int64_t gmOffset, int64_t currentNum)
285+{
286+ LocalTensor<T> varOutLocal = varOutQue_.template DeQue<T>();
287+ LocalTensor<T> mgOutLocal = mgOutQue_.template DeQue<T>();
288+ LocalTensor<T> msOutLocal = msOutQue_.template DeQue<T>();
289+ LocalTensor<T> momOutLocal = momOutQue_.template DeQue<T>();
290+ 
291+ DataCopyExtParams copyParams = MakeTileCopyParams(currentNum);
292+ 
293+ DataCopyPad(varOutGm_[gmOffset], varOutLocal, copyParams);
294+ DataCopyPad(mgOutGm_[gmOffset], mgOutLocal, copyParams);
295+ DataCopyPad(msOutGm_[gmOffset], msOutLocal, copyParams);
296+ DataCopyPad(momOutGm_[gmOffset], momOutLocal, copyParams);
297+ 
298+ varOutQue_.FreeTensor(varOutLocal);
299+ mgOutQue_.FreeTensor(mgOutLocal);
300+ msOutQue_.FreeTensor(msOutLocal);
301+ momOutQue_.FreeTensor(momOutLocal);
302+}
303+ 
304+// =============================================================================
305+// Compute: per-tile compute body.
306+//
307+// fp16 path (TilingKey 2):
308+// 1. Cast(varF32, varLocal_fp16) ... same for mg/ms/mom/grad
309+// 2. mg_new = rho * mg + (1-rho) * grad
310+// 3. ms_new = rho * ms + (1-rho) * grad * grad
311+// 4. tmp = ms_new - mg_new * mg_new + epsilon
312+// 5. denom = sqrt(tmp)
313+// 6. mom_new = momentum * mom + lr * grad / denom
314+// 7. var_new = var - mom_new
315+// 8. Cast(varOutLocal_fp16, varF32) ... same for mg/ms/mom
316+//
317+// fp32 path (TilingKey 1): direct in-place fp32 compute (no Cast, reuses
318+// denom/tmp1/tmp2 scratch buffers).
319+// =============================================================================
320+template <typename T>
321+__aicore__ inline void ApplyCenteredRMSProp<T>::Compute(int64_t currentNum)
322+{
323+ LocalTensor<T> varLocal = varInQue_.template DeQue<T>();
324+ LocalTensor<T> mgLocal = mgInQue_.template DeQue<T>();
325+ LocalTensor<T> msLocal = msInQue_.template DeQue<T>();
326+ LocalTensor<T> momLocal = momInQue_.template DeQue<T>();
327+ LocalTensor<T> gradLocal = gradInQue_.template DeQue<T>();
328+ 
329+ LocalTensor<T> varOutLocal = varOutQue_.template AllocTensor<T>();
330+ LocalTensor<T> mgOutLocal = mgOutQue_.template AllocTensor<T>();
331+ LocalTensor<T> msOutLocal = msOutQue_.template AllocTensor<T>();
332+ LocalTensor<T> momOutLocal = momOutQue_.template AllocTensor<T>();
333+ 
334+ // 32-byte align the work count (matches ubBlockSize for any dtype).
335+ constexpr int64_t kAlignBlock = 32 / sizeof(T);
336+ int64_t alignedNum =
337+ ((currentNum + kAlignBlock - 1) / kAlignBlock) * kAlignBlock;
338+ int32_t n = static_cast<int32_t>(alignedNum);
339+ 
340+ if constexpr (std::is_same_v<T, half>) {
341+ // ---------- fp16 path (TilingKey 2) -----------------------------------
342+ LocalTensor<float> varF32 = varF32Buf_.template Get<float>();
343+ LocalTensor<float> mgF32 = mgF32Buf_.template Get<float>();
344+ LocalTensor<float> msF32 = msF32Buf_.template Get<float>();
345+ LocalTensor<float> momF32 = momF32Buf_.template Get<float>();
346+ LocalTensor<float> gradF32 = gradF32Buf_.template Get<float>();
347+ LocalTensor<float> denom = denomBuf_.template Get<float>();
348+ LocalTensor<float> tmp1 = tmp1Buf_.template Get<float>();
349+ LocalTensor<float> tmp2 = tmp2Buf_.template Get<float>();
350+ 
351+ // 1. Cast inputs fp16 -> fp32 (RoundMode::CAST_NONE for half->float).
352+ Cast(varF32, varLocal, RoundMode::CAST_NONE, n);
353+ Cast(mgF32, mgLocal, RoundMode::CAST_NONE, n);
354+ Cast(msF32, msLocal, RoundMode::CAST_NONE, n);
355+ Cast(momF32, momLocal, RoundMode::CAST_NONE, n);
356+ Cast(gradF32, gradLocal, RoundMode::CAST_NONE, n);
357+ 
358+ // 2. mg_new = rho * mg + (1-rho) * grad
359+ Muls(tmp1, mgF32, rhoScalar_, n); // tmp1 = rho * mg
360+ Muls(tmp2, gradF32, oneMinusRho_, n); // tmp2 = (1-rho) * grad
361+ Add (mgF32, tmp1, tmp2, n); // mgF32 = mg_new
362+ 
363+ // 3. ms_new = rho * ms + (1-rho) * grad*grad
364+ Mul (tmp1, gradF32, gradF32, n); // tmp1 = grad*grad
365+ Muls(tmp1, tmp1, oneMinusRho_, n); // tmp1 = (1-rho)*grad*grad
366+ Muls(tmp2, msF32, rhoScalar_, n); // tmp2 = rho * ms
367+ Add (msF32, tmp1, tmp2, n); // msF32 = ms_new
368+ 
369+ // 4. tmp = ms_new - mg_new*mg_new + epsilon
370+ // Clamp inner to >= 0 to mirror CPU golden's numerical-safety guard
371+ // (ms_new < mg_new^2 can happen due to fp16 truncation, leading to
372+ // Sqrt(NaN) and downstream NaN propagation). See
373+ // issues/issue_20260423_sqrt_negative_clamp_1.md.
374+ Mul (tmp1, mgF32, mgF32, n); // tmp1 = mg_new^2
375+ Sub (tmp2, msF32, tmp1, n); // tmp2 = ms_new - mg_new^2
376+ Maxs(tmp2, tmp2, 0.0f, n); // tmp2 = max(tmp2, 0) -- clamp
377+ Adds(tmp2, tmp2, epsilonScalar_, n); // tmp2 = ... + epsilon
378+ 
379+ // 5. denom = sqrt(tmp)
380+ Sqrt(denom, tmp2, n);
381+ 
382+ // 6. mom_new = momentum * mom + lr * grad / denom
383+ Div (tmp1, gradF32, denom, n); // tmp1 = grad / denom
384+ Muls(tmp1, tmp1, lrScalar_, n); // tmp1 = lr * grad / denom
385+ Muls(tmp2, momF32, momentumScalar_, n); // tmp2 = momentum * mom
386+ Add (momF32, tmp1, tmp2, n); // momF32 = mom_new
387+ 
388+ // 7. var_new = var - mom_new
389+ Sub (varF32, varF32, momF32, n);
390+ 
391+ // 8. Cast back fp32 -> fp16 (RoundMode::CAST_RINT for float->half).
392+ Cast(varOutLocal, varF32, RoundMode::CAST_RINT, n);
393+ Cast(mgOutLocal, mgF32, RoundMode::CAST_RINT, n);
394+ Cast(msOutLocal, msF32, RoundMode::CAST_RINT, n);
395+ Cast(momOutLocal, momF32, RoundMode::CAST_RINT, n);
396+ } else {
397+ // ---------- fp32 path (TilingKey 1) ----------------------------------
398+ // Direct fp32 compute: no Cast, operate on the fp16 queue tensors as
399+ // LocalTensor<float> aliases. Reuse denomBuf_ / tmp1Buf_ / tmp2Buf_
400+ // as fp32 scratch (mgF32/msF32/momF32/gradF32 buffers are left unused
401+ // on this path; UB budget is computed accordingly in Tiling).
402+ //
403+ // Computation matches the fp16 path (DESIGN §6.1 fp32 route):
404+ // mg_new = rho * mg + (1 - rho) * grad
405+ // ms_new = rho * ms + (1 - rho) * grad * grad
406+ // denom = sqrt(max(ms_new - mg_new^2, 0) + epsilon)
407+ // mom_new = momentum * mom + lr * grad / denom
408+ // var_new = var - mom_new
409+ // Maxs(..., 0.0f) clamp is shared with fp16 path for numerical safety.
410+ LocalTensor<float> denom = denomBuf_.template Get<float>();
411+ LocalTensor<float> tmp1 = tmp1Buf_.template Get<float>();
412+ LocalTensor<float> tmp2 = tmp2Buf_.template Get<float>();
413+ 
414+ // 2. mg_new = rho * mg + (1-rho) * grad
415+ Muls(tmp1, mgLocal, static_cast<T>(rhoScalar_), n);
416+ Muls(tmp2, gradLocal, static_cast<T>(oneMinusRho_), n);
417+ Add (mgOutLocal, tmp1, tmp2, n); // mgOut = mg_new
418+ 
419+ // 3. ms_new = rho * ms + (1-rho) * grad*grad
420+ Mul (tmp1, gradLocal, gradLocal, n);
421+ Muls(tmp1, tmp1, static_cast<T>(oneMinusRho_), n);
422+ Muls(tmp2, msLocal, static_cast<T>(rhoScalar_), n);
423+ Add (msOutLocal, tmp1, tmp2, n); // msOut = ms_new
424+ 
425+ // 4. inner = max(ms_new - mg_new^2, 0) + epsilon
426+ Mul (tmp1, mgOutLocal, mgOutLocal, n); // tmp1 = mg_new^2
427+ Sub (tmp2, msOutLocal, tmp1, n); // tmp2 = ms_new - mg_new^2
428+ Maxs(tmp2, tmp2, 0.0f, n); // clamp >= 0
429+ Adds(tmp2, tmp2, static_cast<T>(epsilonScalar_), n);
430+ 
431+ // 5. denom = sqrt(inner)
432+ Sqrt(denom, tmp2, n);
433+ 
434+ // 6. mom_new = momentum * mom + lr * grad / denom
435+ Div (tmp1, gradLocal, denom, n);
436+ Muls(tmp1, tmp1, static_cast<T>(lrScalar_), n);
437+ Muls(tmp2, momLocal, static_cast<T>(momentumScalar_), n);
438+ Add (momOutLocal, tmp1, tmp2, n); // momOut = mom_new
439+ 
440+ // 7. var_new = var - mom_new
441+ Sub (varOutLocal, varLocal, momOutLocal, n);
442+ }
443+ 
444+ varOutQue_.template EnQue<T>(varOutLocal);
445+ mgOutQue_.template EnQue<T>(mgOutLocal);
446+ msOutQue_.template EnQue<T>(msOutLocal);
447+ momOutQue_.template EnQue<T>(momOutLocal);
448+ 
449+ varInQue_.FreeTensor(varLocal);
450+ mgInQue_.FreeTensor(mgLocal);
451+ msInQue_.FreeTensor(msLocal);
452+ momInQue_.FreeTensor(momLocal);
453+ gradInQue_.FreeTensor(gradLocal);
454+}
455+ 
456+// =============================================================================
457+// Process: main loop over UB-sized chunks.
458+// =============================================================================
459+template <typename T>
460+__aicore__ inline void ApplyCenteredRMSProp<T>::Process()
461+{
462+ if (blockLen_ <= 0) {
463+ return;
464+ }
465+ int64_t loopCount = (blockLen_ + ubFactor_ - 1) / ubFactor_;
466+ for (int64_t i = 0; i < loopCount; i++) {
467+ int64_t gmOffset = i * ubFactor_;
468+ int64_t currentNum = (i == (loopCount - 1))
469+ ? (blockLen_ - gmOffset)
470+ : ubFactor_;
471+ CopyInTile(gmOffset, currentNum);
472+ Compute(currentNum);
473+ CopyOutTile(gmOffset, currentNum);
474+ }
475+}
476+ 
477+} // namespace NsApplyCenteredRMSProp
478+ 
479+#endif // APPLY_CENTERED_RMS_PROP_H
Aexperimental/optim/apply_centered_rms_prop/op_kernel/apply_centered_rms_prop_tiling_data.h+40-0
@@ -0,0 +1,40 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+/**
17+ * \file apply_centered_rms_prop_tiling_data.h
18+ * \brief ApplyCenteredRMSProp TilingData structure (arch35).
19+ *
20+ * Standard C++ struct form (per DESIGN.md §5.3 "not using BEGIN_TILING_DATA_DEF").
21+ *
22+ * Iteration-1 skeleton: fp16 path only. The struct itself is dtype-agnostic;
23+ * fp32 path will be added in iteration-2 and reuse the same TilingData layout.
24+ */
25+ 
26+#ifndef _APPLY_CENTERED_RMS_PROP_TILING_DATA_H_
27+#define _APPLY_CENTERED_RMS_PROP_TILING_DATA_H_
28+ 
29+#include <cstdint>
30+ 
31+struct ApplyCenteredRMSPropTilingData {
32+ // Total number of elements in var/mg/ms/mom/grad (all share the same shape).
33+ int64_t totalElements = 0;
34+ // Number of elements per-core (aligned up to ubBlockSize to keep DMA safe).
35+ int64_t blockFactor = 0;
36+ // Number of elements processed per UB iteration.
37+ int64_t ubFactor = 0;
38+};
39+ 
40+#endif // _APPLY_CENTERED_RMS_PROP_TILING_DATA_H_
Aexperimental/optim/apply_centered_rms_prop/op_kernel/apply_centered_rms_prop_tiling_key.h+55-0
@@ -0,0 +1,55 @@
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+ * NOTE: Portions of this code were AI-generated and have been
13+ * technically reviewed for functional accuracy and security
14+ */
15+ 
16+/**
17+ * \file apply_centered_rms_prop_tiling_key.h
18+ * \brief ApplyCenteredRMSProp TilingKey template-argument declaration (arch35).
19+ *
20+ * DESIGN §5.4:
21+ * TilingKey = 1 -> float32 path (direct fp32 compute, no Cast)
22+ * TilingKey = 2 -> float16 path (Cast fp16->fp32, compute, Cast fp32->fp16)
23+ *
24+ * They are encoded via D_T_VAR (dtype template arg).
25+ * D_T_VAR = C_DT_FLOAT -> TilingKey 1
26+ * D_T_VAR = C_DT_FLOAT16 -> TilingKey 2
27+ *
28+ * Iteration-1 skeleton registers both fp16 and fp32 in the template-argument
29+ * declaration so the autogen pipeline produces both binary slots. The fp16
30+ * branch is fully implemented; the fp32 branch is a stub dispatched via
31+ * if constexpr in the kernel entry (fp32 body will land in iteration-2).
32+ *
33+ * Uses ASCENDC_TPL_ARGS_DECL template-argument mechanism.
34+ * TILING_KEY_IS macro is forbidden per DESIGN §5.4.
35+ */
36+ 
37+#ifndef __APPLY_CENTERED_RMS_PROP_TILING_KEY_H__
38+#define __APPLY_CENTERED_RMS_PROP_TILING_KEY_H__
39+ 
40+#include "ascendc/host_api/tiling/template_argument.h"
41+ 
42+ASCENDC_TPL_ARGS_DECL(ApplyCenteredRMSProp,
43+ ASCENDC_TPL_DATATYPE_DECL(D_T_VAR, C_DT_FLOAT16, C_DT_FLOAT, ASCENDC_TPL_INPUT(0))
44+);
45+ 
46+ASCENDC_TPL_SEL(
47+ ASCENDC_TPL_ARGS_SEL(
48+ ASCENDC_TPL_DATATYPE_SEL(D_T_VAR, C_DT_FLOAT16)
49+ ),
50+ ASCENDC_TPL_ARGS_SEL(
51+ ASCENDC_TPL_DATATYPE_SEL(D_T_VAR, C_DT_FLOAT)
52+ ),
53+);
54+ 
55+#endif // __APPLY_CENTERED_RMS_PROP_TILING_KEY_H__