Star906
1825
代码
代码
Issues226
Pull Requests303
Actions
项目讨论
Wiki
分析
Star906
1825
已合并
[CANNbot] add apply_proximal_adagrad operator for ascend950 #4101
gxx_gitcode创建于 4月22日
[CANNbot] add apply_proximal_adagrad operator for ascend950 #4101
已合并
gxx_gitcode创建于 4月22日
从gxx_gitcode/ops-nnclaude_applyproximaladagrad_experimental合入到cann/ops-nnmaster
共 11 个文件变更+1426-0
Aexperimental/optim/apply_proximal_adagrad/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_proximal_adagrad ACLNNTYPE aclnn COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR} DISABLE_IN_OPP TRUE)
Aexperimental/optim/apply_proximal_adagrad/README.md+139-0
@@ -0,0 +1,139 @@
1+# ApplyProximalAdagrad
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+- 算子功能：ApplyProximalAdagrad 是结合 Adagrad 自适应学习率与 FOBOS（Forward-Backward Splitting）Proximal 近端算法的优化器算子，功能对标 `tf.raw_ops.ApplyProximalAdagrad`。基于梯度平方累加器自适应调整学习率，并通过软阈值（L1 正则化）与缩放（L2 正则化）对模型参数进行原地更新。
17+- 计算公式：
18+ 
19+ $$
20+ \begin{aligned}
21+ \text{accum}_t &= \text{accum}_{t-1} + \text{grad}_t^2 \\
22+ \eta_t &= \frac{\text{lr}}{\sqrt{\text{accum}_t}} \\
23+ \text{prox}_t &= \text{var}_{t-1} - \eta_t \cdot \text{grad}_t \\
24+ \text{var}_t &= \frac{\text{sign}(\text{prox}_t)}{1 + \eta_t \cdot \text{l2}} \cdot \max\!\left(|\text{prox}_t| - \eta_t \cdot \text{l1},\ 0\right)
25+ \end{aligned}
26+ $$
27+ 
28+ 当 L1 = 0 时简化为：
29+ 
30+ $$
31+ \text{var}_t = \frac{\text{prox}_t}{1 + \eta_t \cdot \text{l2}}
32+ $$
33+ 
34+- 说明：
35+ - `var`（参数）与 `accum`（梯度平方累加器）均为 Ref Tensor，算子执行后**原地更新**。
36+ - `lr`、`l1`、`l2` 为 0-D 标量 Tensor，分别要求 `lr > 0`、`l1 ≥ 0`、`l2 ≥ 0`。
37+ - 逐元素独立计算，天然确定性，无跨元素/跨核依赖。
38+ 
39+## 参数说明
40+ 
41+<table style="table-layout: fixed; width: 1576px"><colgroup>
42+<col style="width: 150px">
43+<col style="width: 150px">
44+<col style="width: 420px">
45+<col style="width: 140px">
46+<col style="width: 140px">
47+</colgroup>
48+<thead>
49+ <tr>
50+ <th>参数名</th>
51+ <th>输入/输出/属性</th>
52+ <th>描述</th>
53+ <th>数据类型</th>
54+ <th>数据格式</th>
55+ </tr></thead>
56+<tbody>
57+ <tr>
58+ <td>var</td>
59+ <td>输入</td>
60+ <td>公式中的 var，待更新的模型参数（Ref Tensor，原地更新）。shape 与 accum/grad 一致。</td>
61+ <td>FLOAT</td>
62+ <td>ND</td>
63+ </tr>
64+ <tr>
65+ <td>accum</td>
66+ <td>输入</td>
67+ <td>公式中的 accum，梯度平方累加器（Ref Tensor，原地更新）。shape 与 var/grad 一致，要求各元素非负。</td>
68+ <td>FLOAT</td>
69+ <td>ND</td>
70+ </tr>
71+ <tr>
72+ <td>lr</td>
73+ <td>输入</td>
74+ <td>公式中的 lr，学习率。0-D 或 1 元素 1-D Tensor，要求 lr &gt; 0。</td>
75+ <td>FLOAT</td>
76+ <td>ND</td>
77+ </tr>
78+ <tr>
79+ <td>l1</td>
80+ <td>输入</td>
81+ <td>公式中的 l1，L1 正则化强度。0-D 或 1 元素 1-D Tensor，要求 l1 ≥ 0。</td>
82+ <td>FLOAT</td>
83+ <td>ND</td>
84+ </tr>
85+ <tr>
86+ <td>l2</td>
87+ <td>输入</td>
88+ <td>公式中的 l2，L2 正则化强度。0-D 或 1 元素 1-D Tensor，要求 l2 ≥ 0。</td>
89+ <td>FLOAT</td>
90+ <td>ND</td>
91+ </tr>
92+ <tr>
93+ <td>grad</td>
94+ <td>输入</td>
95+ <td>公式中的 grad，当前步的梯度张量。shape 与 var/accum 一致。</td>
96+ <td>FLOAT</td>
97+ <td>ND</td>
98+ </tr>
99+ <tr>
100+ <td>var_out</td>
101+ <td>输出</td>
102+ <td>更新后的参数，与输入 var 共享存储（inplace 更新）。</td>
103+ <td>FLOAT</td>
104+ <td>ND</td>
105+ </tr>
106+ <tr>
107+ <td>accum_out</td>
108+ <td>输出</td>
109+ <td>更新后的累加器，与输入 accum 共享存储（inplace 更新）。</td>
110+ <td>FLOAT</td>
111+ <td>ND</td>
112+ </tr>
113+</tbody></table>
114+ 
115+## 约束说明
116+ 
117+- 仅支持 <term>Ascend 950PR/Ascend 950DT</term>（arch35 / DAV_3510），不适配其他芯片代际。
118+- 仅支持 `float32` 数据类型。
119+- `var`、`accum`、`grad` 三者 shape 必须完全一致，且均为连续排布的 ND Tensor。
120+- `lr`、`l1`、`l2` 必须为 0-D 或 1 元素 1-D 的标量 Tensor。
121+- 调用方需保证 `accum ≥ 0`、`lr > 0`、`l1 ≥ 0`、`l2 ≥ 0`；算子内部不做运行时值域校验。
122+- `accum + grad^2 == 0` 时 `rsqrt` 输出 Inf/NaN，行为与 PyTorch / TensorFlow 原生实现一致，需由上游调用方规避。
123+- `var` 与 `accum` 为 Ref Tensor，Host aclnn 侧必须显式构造两个占位输出 Tensor（`var_out`、`accum_out`），并与各自输入共享 Device 地址以保证 inplace 语义。
124+ 
125+## 调用说明
126+ 
127+<table><thead>
128+ <tr>
129+ <th>调用方式</th>
130+ <th>调用样例</th>
131+ <th>说明</th>
132+ </tr></thead>
133+<tbody>
134+ <tr>
135+ <td>aclnn 调用</td>
136+ <td><a href="./examples/arch35/test_aclnn_apply_proximal_adagrad.cpp">test_aclnn_apply_proximal_adagrad</a></td>
137+ <td>Ascend 950 上通过 aclnn 两段式接口 <code>aclnnApplyProximalAdagradGetWorkspaceSize</code> → <code>aclnnApplyProximalAdagrad</code> 调用。<code>var_out</code> / <code>accum_out</code> 需与 <code>var</code> / <code>accum</code> 共享 Device 地址以保证 inplace 更新。</td>
138+ </tr>
139+</tbody></table>
Aexperimental/optim/apply_proximal_adagrad/examples/arch35/test_aclnn_apply_proximal_adagrad.cpp+279-0
@@ -0,0 +1,279 @@
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 test_aclnn_apply_proximal_adagrad.cpp
18+ * @brief ApplyProximalAdagrad 算子 ACLNN 调用示例（两段式接口）
19+ */
20+ 
21+#include <iostream>
22+#include <iomanip>
23+#include <cstdint>
24+#include <cstdio>
25+#include <cstring>
26+#include <cmath>
27+#include <vector>
28+ 
29+#include "acl/acl.h"
30+#include "aclnn_apply_proximal_adagrad.h"
31+ 
32+#define CHECK_ACL(expr) \
33+ do { \
34+ auto _ret = (expr); \
35+ if (_ret != ACL_SUCCESS) { \
36+ std::cerr << "ACL Error: " << #expr << " returned " << _ret \
37+ << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
38+ goto cleanup; \
39+ } \
40+ } while (0)
41+ 
42+namespace {
43+ 
44+// ---------------------------------------------------------------------------
45+// CPU Golden: 与算子语义完全一致
46+// accum' = accum + grad^2
47+// eta = lr / sqrt(accum')
48+// prox = var - eta * grad
49+// var' = (l1==0) ? prox / (1 + eta*l2)
50+// : sign(prox) / (1 + eta*l2) * max(|prox| - eta*l1, 0)
51+// ---------------------------------------------------------------------------
52+inline void CpuGoldenStep(float& var, float& accum, float grad, float lr, float l1, float l2)
53+{
54+ accum = accum + grad * grad;
55+ float eta = lr / std::sqrt(accum);
56+ float prox = var - eta * grad;
57+ float denom = 1.0f + eta * l2;
58+ if (l1 == 0.0f) {
59+ var = prox / denom;
60+ } else {
61+ float sgn = (prox > 0.0f) ? 1.0f : ((prox < 0.0f) ? -1.0f : 0.0f);
62+ float mag = std::fabs(prox) - eta * l1;
63+ if (mag < 0.0f) mag = 0.0f;
64+ var = sgn / denom * mag;
65+ }
66+}
67+ 
68+inline bool ApproxEq(float a, float b, float atol = 1e-4f, float rtol = 1e-4f)
69+{
70+ float diff = std::fabs(a - b);
71+ return diff <= atol + rtol * std::fabs(b);
72+}
73+ 
74+} // namespace
75+ 
76+int main()
77+{
78+ // =========================================================================
79+ // 1. 参数设置：shape=[16]、dtype=float32
80+ // =========================================================================
81+ constexpr int64_t ELEM_COUNT = 16;
82+ const int64_t shape[] = {ELEM_COUNT};
83+ const int64_t strides[] = {1};
84+ constexpr int64_t ndim = 1;
85+ 
86+ // 标量 (lr, l1, l2) 用 numel=1 的 1-D Tensor 表示
87+ const int64_t scalarShape[] = {1};
88+ const int64_t scalarStrides[] = {1};
89+ 
90+ float hostVar[ELEM_COUNT] = {
91+ // 覆盖 +/- / 0 / 较大值
92+ 1.00f, -1.00f, 0.50f, -0.50f,
93+ 0.10f, -0.10f, 0.00f, 2.00f,
94+ -2.00f, 0.05f, -0.05f, 3.50f,
95+ -3.50f, 0.01f, -0.01f, 0.25f,
96+ };
97+ float hostAccum[ELEM_COUNT] = {
98+ // 必须非负（调用方保证）。覆盖小值 / 中等值。
99+ 0.10f, 0.10f, 0.20f, 0.20f,
100+ 0.05f, 0.05f, 0.10f, 1.00f,
101+ 1.00f, 0.01f, 0.01f, 0.50f,
102+ 0.50f, 0.30f, 0.30f, 0.40f,
103+ };
104+ float hostGrad[ELEM_COUNT] = {
105+ 0.10f, 0.10f, -0.20f, 0.20f,
106+ 0.30f, -0.30f, 0.10f, -0.50f,
107+ 0.50f, 0.05f, -0.05f, 1.00f,
108+ -1.00f, 0.40f, -0.40f, 0.15f,
109+ };
110+ float hostLr[1] = {0.1f};
111+ float hostL1[1] = {0.01f}; // l1 > 0 → 走 HAS_L1=1 分支（含 abs/max 阈值收缩）
112+ float hostL2[1] = {0.0f};
113+ 
114+ // CPU Golden（独立副本，避免污染输入）
115+ float goldVar[ELEM_COUNT];
116+ float goldAccum[ELEM_COUNT];
117+ std::memcpy(goldVar, hostVar, sizeof(hostVar));
118+ std::memcpy(goldAccum, hostAccum, sizeof(hostAccum));
119+ for (int i = 0; i < ELEM_COUNT; ++i) {
120+ CpuGoldenStep(goldVar[i], goldAccum[i], hostGrad[i], hostLr[0], hostL1[0], hostL2[0]);
121+ }
122+ 
123+ // =========================================================================
124+ // 2. ACL 初始化
125+ // =========================================================================
126+ int32_t ret = 1;
127+ aclrtStream stream = nullptr;
128+ void *devVar = nullptr, *devAccum = nullptr;
129+ void *devLr = nullptr, *devL1 = nullptr, *devL2 = nullptr;
130+ void *devGrad = nullptr;
131+ void *workspace = nullptr;
132+ aclTensor *tVar = nullptr, *tAccum = nullptr;
133+ aclTensor *tLr = nullptr, *tL1 = nullptr, *tL2 = nullptr;
134+ aclTensor *tGrad = nullptr;
135+ aclTensor *tVarOut = nullptr, *tAccumOut = nullptr;
136+ 
137+ CHECK_ACL(aclInit(nullptr));
138+ CHECK_ACL(aclrtSetDevice(0));
139+ CHECK_ACL(aclrtCreateStream(&stream));
140+ 
141+ // =========================================================================
142+ // 3. 设备内存分配 & 输入数据拷贝 Host→Device
143+ // =========================================================================
144+ {
145+ const size_t vecBytes = ELEM_COUNT * sizeof(float);
146+ const size_t scalarBytes = 1 * sizeof(float);
147+ 
148+ CHECK_ACL(aclrtMalloc(&devVar, vecBytes, ACL_MEM_MALLOC_HUGE_FIRST));
149+ CHECK_ACL(aclrtMalloc(&devAccum, vecBytes, ACL_MEM_MALLOC_HUGE_FIRST));
150+ CHECK_ACL(aclrtMalloc(&devGrad, vecBytes, ACL_MEM_MALLOC_HUGE_FIRST));
151+ CHECK_ACL(aclrtMalloc(&devLr, scalarBytes, ACL_MEM_MALLOC_HUGE_FIRST));
152+ CHECK_ACL(aclrtMalloc(&devL1, scalarBytes, ACL_MEM_MALLOC_HUGE_FIRST));
153+ CHECK_ACL(aclrtMalloc(&devL2, scalarBytes, ACL_MEM_MALLOC_HUGE_FIRST));
154+ 
155+ CHECK_ACL(aclrtMemcpy(devVar, vecBytes, hostVar, vecBytes, ACL_MEMCPY_HOST_TO_DEVICE));
156+ CHECK_ACL(aclrtMemcpy(devAccum, vecBytes, hostAccum, vecBytes, ACL_MEMCPY_HOST_TO_DEVICE));
157+ CHECK_ACL(aclrtMemcpy(devGrad, vecBytes, hostGrad, vecBytes, ACL_MEMCPY_HOST_TO_DEVICE));
158+ CHECK_ACL(aclrtMemcpy(devLr, scalarBytes, hostLr, scalarBytes, ACL_MEMCPY_HOST_TO_DEVICE));
159+ CHECK_ACL(aclrtMemcpy(devL1, scalarBytes, hostL1, scalarBytes, ACL_MEMCPY_HOST_TO_DEVICE));
160+ CHECK_ACL(aclrtMemcpy(devL2, scalarBytes, hostL2, scalarBytes, ACL_MEMCPY_HOST_TO_DEVICE));
161+ 
162+ // =====================================================================
163+ // 4. 创建 aclTensor
164+ // var/accum/grad: FLOAT, ND, [16]
165+ // lr/l1/l2: FLOAT, ND, [1] (numel=1 1-D 标量)
166+ // varOutOut / accumOutOut: 绑定与 var / accum 同一 Device 存储 → 观察 inplace
167+ // =====================================================================
168+ tVar = aclCreateTensor(shape, ndim, ACL_FLOAT, strides, 0,
169+ ACL_FORMAT_ND, shape, ndim, devVar);
170+ tAccum = aclCreateTensor(shape, ndim, ACL_FLOAT, strides, 0,
171+ ACL_FORMAT_ND, shape, ndim, devAccum);
172+ tGrad = aclCreateTensor(shape, ndim, ACL_FLOAT, strides, 0,
173+ ACL_FORMAT_ND, shape, ndim, devGrad);
174+ tLr = aclCreateTensor(scalarShape, 1, ACL_FLOAT, scalarStrides, 0,
175+ ACL_FORMAT_ND, scalarShape, 1, devLr);
176+ tL1 = aclCreateTensor(scalarShape, 1, ACL_FLOAT, scalarStrides, 0,
177+ ACL_FORMAT_ND, scalarShape, 1, devL1);
178+ tL2 = aclCreateTensor(scalarShape, 1, ACL_FLOAT, scalarStrides, 0,
179+ ACL_FORMAT_ND, scalarShape, 1, devL2);
180+ // 占位输出复用同一 Device buffer，调用结束后通过 devVar / devAccum 直接读取
181+ tVarOut = aclCreateTensor(shape, ndim, ACL_FLOAT, strides, 0,
182+ ACL_FORMAT_ND, shape, ndim, devVar);
183+ tAccumOut = aclCreateTensor(shape, ndim, ACL_FLOAT, strides, 0,
184+ ACL_FORMAT_ND, shape, ndim, devAccum);
185+ 
186+ if (!tVar || !tAccum || !tLr || !tL1 || !tL2 || !tGrad || !tVarOut || !tAccumOut) {
187+ std::cerr << "aclCreateTensor failed" << std::endl;
188+ goto cleanup;
189+ }
190+ 
191+ // =====================================================================
192+ // 5. 调用 aclnnApplyProximalAdagrad（两段式接口）
193+ // =====================================================================
194+ uint64_t workspaceSize = 0;
195+ aclOpExecutor *executor = nullptr;
196+ 
197+ CHECK_ACL(aclnnApplyProximalAdagradGetWorkspaceSize(
198+ tVar, tAccum, tLr, tL1, tL2, tGrad, tVarOut, tAccumOut,
199+ &workspaceSize, &executor));
200+ 
201+ if (workspaceSize > 0) {
202+ CHECK_ACL(aclrtMalloc(&workspace, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST));
203+ }
204+ 
205+ CHECK_ACL(aclnnApplyProximalAdagrad(workspace, workspaceSize, executor, stream));
206+ CHECK_ACL(aclrtSynchronizeStream(stream));
207+ 
208+ // =====================================================================
209+ // 6. 读取 inplace 结果（var / accum 已被更新）& 精度验证
210+ // =====================================================================
211+ float npuVar[ELEM_COUNT] = {};
212+ float npuAccum[ELEM_COUNT] = {};
213+ const size_t vecBytesOut = ELEM_COUNT * sizeof(float);
214+ CHECK_ACL(aclrtMemcpy(npuVar, vecBytesOut, devVar, vecBytesOut, ACL_MEMCPY_DEVICE_TO_HOST));
215+ CHECK_ACL(aclrtMemcpy(npuAccum, vecBytesOut, devAccum, vecBytesOut, ACL_MEMCPY_DEVICE_TO_HOST));
216+ 
217+ std::cout << "ApplyProximalAdagrad Example (shape: [16], dtype: float32)" << std::endl;
218+ std::cout << " scalars: lr=" << hostLr[0]
219+ << " l1=" << hostL1[0]
220+ << " l2=" << hostL2[0] << std::endl;
221+ std::cout << " formula: accum' = accum + grad^2" << std::endl;
222+ std::cout << " eta = lr / sqrt(accum')" << std::endl;
223+ std::cout << " prox = var - eta * grad" << std::endl;
224+ std::cout << " var' = sign(prox) / (1 + eta*l2) * max(|prox| - eta*l1, 0)" << std::endl;
225+ std::cout << "-------------------------------------------------------------------------------" << std::endl;
226+ std::printf(" %4s | %10s | %10s | %10s | %10s | %10s | %s\n",
227+ "Idx", "var", "accum", "grad", "var_npu", "var_gold", "Status");
228+ std::cout << "-------------------------------------------------------------------------------" << std::endl;
229+ 
230+ int passVar = 0, passAccum = 0;
231+ for (int i = 0; i < ELEM_COUNT; ++i) {
232+ bool okV = ApproxEq(npuVar[i], goldVar[i]);
233+ bool okA = ApproxEq(npuAccum[i], goldAccum[i]);
234+ if (okV) ++passVar;
235+ if (okA) ++passAccum;
236+ std::printf(" %4d | %10.5f | %10.5f | %10.5f | %10.5f | %10.5f | %s/%s\n",
237+ i, hostVar[i], hostAccum[i], hostGrad[i],
238+ npuVar[i], goldVar[i],
239+ okV ? "PASS" : "FAIL",
240+ okA ? "PASS" : "FAIL");
241+ }
242+ 
243+ std::cout << "-------------------------------------------------------------------------------" << std::endl;
244+ std::cout << "Result (var): " << passVar << "/" << ELEM_COUNT << " passed" << std::endl;
245+ std::cout << "Result (accum): " << passAccum << "/" << ELEM_COUNT << " passed" << std::endl;
246+ if (passVar == ELEM_COUNT && passAccum == ELEM_COUNT) {
247+ std::cout << "ALL PASS" << std::endl;
248+ ret = 0;
249+ } else {
250+ std::cout << "FAILED" << std::endl;
251+ ret = 1;
252+ }
253+ }
254+ 
255+ // =========================================================================
256+ // 7. 资源释放
257+ // =========================================================================
258+cleanup:
259+ if (tVar) aclDestroyTensor(tVar);
260+ if (tAccum) aclDestroyTensor(tAccum);
261+ if (tLr) aclDestroyTensor(tLr);
262+ if (tL1) aclDestroyTensor(tL1);
263+ if (tL2) aclDestroyTensor(tL2);
264+ if (tGrad) aclDestroyTensor(tGrad);
265+ if (tVarOut) aclDestroyTensor(tVarOut);
266+ if (tAccumOut) aclDestroyTensor(tAccumOut);
267+ if (workspace) aclrtFree(workspace);
268+ if (devVar) aclrtFree(devVar);
269+ if (devAccum) aclrtFree(devAccum);
270+ if (devGrad) aclrtFree(devGrad);
271+ if (devLr) aclrtFree(devLr);
272+ if (devL1) aclrtFree(devL1);
273+ if (devL2) aclrtFree(devL2);
274+ if (stream) aclrtDestroyStream(stream);
275+ aclrtResetDevice(0);
276+ aclFinalize();
277+ 
278+ return ret;
279+}
Aexperimental/optim/apply_proximal_adagrad/op_host/apply_proximal_adagrad_def.cpp+105-0
@@ -0,0 +1,105 @@
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_proximal_adagrad_def.cpp
18+ * \brief ApplyProximalAdagrad operator definition
19+ *
20+ * Optimizer-style operator: 6 inputs, 2 inplace outputs.
21+ *
22+ * Inputs (in order):
23+ * var (Tensor, float32, ND, shape == accum == grad) -- inplace updated
24+ * accum (Tensor, float32, ND, shape == var == grad) -- inplace updated
25+ * lr (Tensor, float32, ND, 0-D or 1-element 1-D) -- scalar
26+ * l1 (Tensor, float32, ND, 0-D or 1-element 1-D) -- scalar
27+ * l2 (Tensor, float32, ND, 0-D or 1-element 1-D) -- scalar
28+ * grad (Tensor, float32, ND, shape == var == accum)
29+ *
30+ * Outputs:
31+ * var_out (Tensor, float32, ND) -- shares storage with input var
32+ * accum_out (Tensor, float32, ND) -- shares storage with input accum
33+ *
34+ * Target: Ascend950 (arch35 / DAV_3510) only.
35+ */
36+#include "register/op_def_registry.h"
37+ 
38+namespace ops {
39+class ApplyProximalAdagrad : public OpDef {
40+public:
41+ explicit ApplyProximalAdagrad(const char* name) : OpDef(name)
42+ {
43+ // --- Inputs (6) ---------------------------------------------------
44+ this->Input("var")
45+ .ParamType(REQUIRED)
46+ .DataType({ge::DT_FLOAT})
47+ .Format({ge::FORMAT_ND})
48+ .UnknownShapeFormat({ge::FORMAT_ND})
49+ .AutoContiguous();
50+ this->Input("accum")
51+ .ParamType(REQUIRED)
52+ .DataType({ge::DT_FLOAT})
53+ .Format({ge::FORMAT_ND})
54+ .UnknownShapeFormat({ge::FORMAT_ND})
55+ .AutoContiguous();
56+ this->Input("lr")
57+ .ParamType(REQUIRED)
58+ .DataType({ge::DT_FLOAT})
59+ .Format({ge::FORMAT_ND})
60+ .UnknownShapeFormat({ge::FORMAT_ND});
61+ this->Input("l1")
62+ .ParamType(REQUIRED)
63+ .DataType({ge::DT_FLOAT})
64+ .Format({ge::FORMAT_ND})
65+ .UnknownShapeFormat({ge::FORMAT_ND});
66+ this->Input("l2")
67+ .ParamType(REQUIRED)
68+ .DataType({ge::DT_FLOAT})
69+ .Format({ge::FORMAT_ND})
70+ .UnknownShapeFormat({ge::FORMAT_ND});
71+ this->Input("grad")
72+ .ParamType(REQUIRED)
73+ .DataType({ge::DT_FLOAT})
74+ .Format({ge::FORMAT_ND})
75+ .UnknownShapeFormat({ge::FORMAT_ND})
76+ .AutoContiguous();
77+ 
78+ // --- Outputs (2, inplace via shared storage at L2 layer) ----------
79+ this->Output("var_out")
80+ .ParamType(REQUIRED)
81+ .DataType({ge::DT_FLOAT})
82+ .Format({ge::FORMAT_ND})
83+ .UnknownShapeFormat({ge::FORMAT_ND})
84+ .AutoContiguous();
85+ this->Output("accum_out")
86+ .ParamType(REQUIRED)
87+ .DataType({ge::DT_FLOAT})
88+ .Format({ge::FORMAT_ND})
89+ .UnknownShapeFormat({ge::FORMAT_ND})
90+ .AutoContiguous();
91+ 
92+ // --- Ascend950 (arch35) AI Core configuration ---------------------
93+ OpAICoreConfig aiCoreConfig;
94+ aiCoreConfig.DynamicCompileStaticFlag(true)
95+ .DynamicFormatFlag(false)
96+ .DynamicRankSupportFlag(true)
97+ .DynamicShapeSupportFlag(true)
98+ .NeedCheckSupportFlag(false)
99+ .PrecisionReduceFlag(false)
100+ .ExtendCfgInfo("opFile.value", "apply_proximal_adagrad");
101+ this->AICore().AddConfig("ascend950", aiCoreConfig);
102+ }
103+};
104+OP_ADD(ApplyProximalAdagrad);
105+} // namespace ops
Aexperimental/optim/apply_proximal_adagrad/op_host/apply_proximal_adagrad_infershape.cpp+72-0
@@ -0,0 +1,72 @@
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_proximal_adagrad_infershape.cpp
18+ * \brief ApplyProximalAdagrad shape / dtype inference.
19+ *
20+ * - var_out.shape = var.shape
21+ * - accum_out.shape = accum.shape
22+ * - both outputs are float32
23+ */
24+ 
25+#include "register/op_impl_registry.h"
26+#include "exe_graph/runtime/infer_shape_context.h"
27+#include "exe_graph/runtime/infer_datatype_context.h"
28+ 
29+using namespace ge;
30+ 
31+namespace ops {
32+ 
33+static ge::graphStatus InferShape4ApplyProximalAdagrad(gert::InferShapeContext* context)
34+{
35+ // Input 0 = var, output 0 = var_out (inplace)
36+ const gert::Shape* varShape = context->GetInputShape(0);
37+ if (varShape == nullptr) {
38+ return ge::GRAPH_FAILED;
39+ }
40+ // Input 1 = accum, output 1 = accum_out (inplace)
41+ const gert::Shape* accumShape = context->GetInputShape(1);
42+ if (accumShape == nullptr) {
43+ return ge::GRAPH_FAILED;
44+ }
45+ 
46+ gert::Shape* varOutShape = context->GetOutputShape(0);
47+ if (varOutShape == nullptr) {
48+ return ge::GRAPH_FAILED;
49+ }
50+ gert::Shape* accumOutShape = context->GetOutputShape(1);
51+ if (accumOutShape == nullptr) {
52+ return ge::GRAPH_FAILED;
53+ }
54+ 
55+ *varOutShape = *varShape;
56+ *accumOutShape = *accumShape;
57+ return ge::GRAPH_SUCCESS;
58+}
59+ 
60+static ge::graphStatus InferDataType4ApplyProximalAdagrad(gert::InferDataTypeContext* context)
61+{
62+ // Outputs dtype follow var/accum (both must be float32 per spec).
63+ context->SetOutputDataType(0, context->GetInputDataType(0));
64+ context->SetOutputDataType(1, context->GetInputDataType(1));
65+ return ge::GRAPH_SUCCESS;
66+}
67+ 
68+IMPL_OP_INFERSHAPE(ApplyProximalAdagrad)
69+ .InferShape(InferShape4ApplyProximalAdagrad)
70+ .InferDataType(InferDataType4ApplyProximalAdagrad);
71+ 
72+} // namespace ops
Aexperimental/optim/apply_proximal_adagrad/op_host/apply_proximal_adagrad_tiling.cpp+231-0
@@ -0,0 +1,231 @@
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_proximal_adagrad_tiling.cpp
18+ * \brief ApplyProximalAdagrad tiling (arch35).
19+ *
20+ * Tiling strategy (iteration-1 skeleton, single TilingKey path):
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: allocate the per-tile size (ubFactor) based on an 8-buffer layout
25+ * that covers var / accum / grad IN queues (x 2 double-buffered),
26+ * var / accum OUT queues (x 2 double-buffered), and a few tmp scratch
27+ * buffers (eta / prox / thresh / scale / sign) kept in float32.
28+ * Total co-resident float32 tensors per tile ~= 3*2 (in) + 2*2 (out) +
29+ * 4 (tmp) = 14 buffers. ubFactor = floor(UB / (14 * 4B)) aligned down
30+ * to 32-byte granularity; an absolute ceiling of 2048 elements keeps
31+ * the first iteration deterministic (see DESIGN s5.2, TILE_ELEM_NUM).
32+ */
33+ 
34+#include "register/op_def_registry.h"
35+#include "op_common/log/log.h"
36+#include "op_common/op_host/util/math_util.h"
37+#include "op_common/op_host/util/platform_util.h"
38+#include "../op_kernel/apply_proximal_adagrad_tiling_data.h"
39+#include "../op_kernel/apply_proximal_adagrad_tiling_key.h"
40+ 
41+namespace optiling {
42+ 
43+using Ops::Base::CeilDiv;
44+using Ops::Base::CeilAlign;
45+using Ops::Base::FloorAlign;
46+using Ops::Base::FloorDiv;
47+using Ops::Base::GetUbBlockSize;
48+ 
49+constexpr uint32_t WS_SYS_SIZE = 0U;
50+constexpr int64_t TYPE_SIZE = 4; // sizeof(float) -- fp32 only.
51+// Per-tile target element count (float32 * 2048 = 8KB).
52+// Iteration-1 uses this conservative tile; later iterations may make it
53+// platform-derived.
54+constexpr int64_t TILE_ELEM_NUM_TARGET = 2048;
55+// Co-resident fp32 UB tensors per tile: var/accum/grad IN (double buffer) +
56+// var/accum OUT (double buffer) + 4 tmp (eta / prox / thresh / scale).
57+// SUG-002: This MUST stay in sync with the kernel-side constant
58+// `kUbResidentFp32TensorCount` defined at the top of
59+// `op_kernel/apply_proximal_adagrad.h`. Any change to the kernel's UB
60+// buffer layout (adding/removing TBuf or TQue, changing double-buffer depth)
61+// MUST update both constants together; see the kernel-side static_assert and
62+// breakdown comment.
63+constexpr int64_t UB_BUFFER_COUNT = 14;
64+ 
65+static const gert::Shape g_vec_1_shape = {1};
66+ 
67+static inline const gert::Shape EnsureNotScalar(const gert::Shape& in_shape)
68+{
69+ if (in_shape.GetDimNum() == 0) {
70+ return g_vec_1_shape;
71+ }
72+ return in_shape;
73+}
74+ 
75+static ge::graphStatus GetPlatformInfo(gert::TilingContext* context,
76+ uint64_t& ubSize, int64_t& coreNum)
77+{
78+ fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo();
79+ OP_CHECK_NULL_WITH_CONTEXT(context, platformInfoPtr);
80+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
81+ coreNum = ascendcPlatform.GetCoreNumAiv();
82+ OP_CHECK_IF(coreNum == 0, OP_LOGE(context, "coreNum is 0"),
83+ return ge::GRAPH_FAILED);
84+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
85+ OP_CHECK_IF(ubSize == 0, OP_LOGE(context, "ubSize is 0"),
86+ return ge::GRAPH_FAILED);
87+ return ge::GRAPH_SUCCESS;
88+}
89+ 
90+static ge::graphStatus GetShapeInfo(gert::TilingContext* context,
91+ int64_t& totalElements,
92+ ge::DataType& dataType)
93+{
94+ // Input 0 = var; its shape is canonical (var/accum/grad must match).
95+ auto inputVar = context->GetInputShape(0);
96+ OP_CHECK_NULL_WITH_CONTEXT(context, inputVar);
97+ auto varShape = EnsureNotScalar(inputVar->GetStorageShape());
98+ totalElements = varShape.GetShapeSize();
99+ 
100+ auto inputDesc = context->GetInputDesc(0);
101+ OP_CHECK_NULL_WITH_CONTEXT(context, inputDesc);
102+ dataType = inputDesc->GetDataType();
103+ OP_CHECK_IF(dataType != ge::DT_FLOAT,
104+ OP_LOGE(context,
105+ "ApplyProximalAdagrad: only float32 is supported, got %d",
106+ static_cast<int>(dataType)),
107+ return ge::GRAPH_FAILED);
108+ return ge::GRAPH_SUCCESS;
109+}
110+ 
111+static ge::graphStatus GetWorkspaceSize(gert::TilingContext* context)
112+{
113+ size_t* currentWorkspace = context->GetWorkspaceSizes(1);
114+ OP_CHECK_NULL_WITH_CONTEXT(context, currentWorkspace);
115+ currentWorkspace[0] = WS_SYS_SIZE;
116+ return ge::GRAPH_SUCCESS;
117+}
118+ 
119+static ge::graphStatus ApplyProximalAdagradTilingFunc(gert::TilingContext* context)
120+{
121+ // 1. Platform info
122+ uint64_t ubSize = 0;
123+ int64_t coreNum = 0;
124+ OP_CHECK_IF(GetPlatformInfo(context, ubSize, coreNum) != ge::GRAPH_SUCCESS,
125+ OP_LOGE(context, "GetPlatformInfo error"),
126+ return ge::GRAPH_FAILED);
127+ 
128+ // 2. Shape / dtype info
129+ int64_t totalElements = 0;
130+ ge::DataType dataType = ge::DT_FLOAT;
131+ OP_CHECK_IF(GetShapeInfo(context, totalElements, dataType) != ge::GRAPH_SUCCESS,
132+ OP_LOGE(context, "GetShapeInfo error"),
133+ return ge::GRAPH_FAILED);
134+ 
135+ // 3. Workspace
136+ OP_CHECK_IF(GetWorkspaceSize(context) != ge::GRAPH_SUCCESS,
137+ OP_LOGE(context, "GetWorkspaceSize error"),
138+ return ge::GRAPH_FAILED);
139+ 
140+ // 4. Fill TilingData
141+ ApplyProximalAdagradTilingData* tiling =
142+ context->GetTilingData<ApplyProximalAdagradTilingData>();
143+ OP_CHECK_NULL_WITH_CONTEXT(context, tiling);
144+ OP_CHECK_IF(
145+ memset_s(tiling, sizeof(ApplyProximalAdagradTilingData), 0,
146+ sizeof(ApplyProximalAdagradTilingData)) != EOK,
147+ OP_LOGE(context, "set tiling data error"),
148+ return ge::GRAPH_FAILED);
149+ 
150+ tiling->totalElements = totalElements;
151+ 
152+ // Empty tensor: run a single idle core so the launcher still succeeds.
153+ if (totalElements == 0) {
154+ tiling->blockFactor = 0;
155+ tiling->ubFactor = 0;
156+ context->SetBlockDim(1);
157+ uint32_t dTypeVar = static_cast<uint32_t>(dataType);
158+ // Pick the simplest binary (PAD_TAIL=0, HAS_L1=1) for the empty path -
159+ // Process() short-circuits before doing any compute.
160+ uint32_t padTail = 0U;
161+ uint32_t hasL1 = 1U;
162+ ASCENDC_TPL_SEL_PARAM(context, dTypeVar, padTail, hasL1);
163+ return ge::GRAPH_SUCCESS;
164+ }
165+ 
166+ // ubBlockSize = 32B / sizeof(T), with T = fp32 -> 8 elements.
167+ int64_t ubBlockSize = Ops::Base::GetUbBlockSize(context);
168+ OP_CHECK_IF(ubBlockSize <= 0,
169+ OP_LOGE(context, "invalid ubBlockSize=%ld", ubBlockSize),
170+ return ge::GRAPH_FAILED);
171+ 
172+ // Multi-core split: ceil-aligned to DMA granularity.
173+ int64_t blockFactor = CeilAlign(CeilDiv(totalElements, coreNum), ubBlockSize);
174+ int64_t usedCoreNum = CeilDiv(totalElements, blockFactor);
175+ 
176+ // UB split. Cap by platform UB, target TILE_ELEM_NUM_TARGET (2048).
177+ int64_t ubCapacityElem =
178+ FloorAlign(FloorDiv(static_cast<int64_t>(ubSize) / TYPE_SIZE,
179+ UB_BUFFER_COUNT),
180+ ubBlockSize);
181+ OP_CHECK_IF(ubCapacityElem <= 0,
182+ OP_LOGE(context, "UB too small: ubCapacityElem=%ld",
183+ ubCapacityElem),
184+ return ge::GRAPH_FAILED);
185+ 
186+ int64_t ubFactor = (TILE_ELEM_NUM_TARGET < ubCapacityElem)
187+ ? TILE_ELEM_NUM_TARGET
188+ : ubCapacityElem;
189+ 
190+ // Also cap by blockFactor so a single core does not allocate more UB
191+ // space than it will ever use.
192+ if (ubFactor > blockFactor) {
193+ ubFactor = FloorAlign(blockFactor, ubBlockSize);
194+ if (ubFactor <= 0) {
195+ ubFactor = ubBlockSize;
196+ }
197+ }
198+ 
199+ tiling->blockFactor = blockFactor;
200+ tiling->ubFactor = ubFactor;
201+ 
202+ context->SetBlockDim(usedCoreNum);
203+ 
204+ // 5. TilingKey via ASCENDC_TPL_SEL_PARAM (template-argument mechanism).
205+ // Iteration-2: derive PAD_TAIL from shape alignment. HAS_L1 cannot be
206+ // determined from host-side tiling without a Host<->Device sync (lr/l1/l2
207+ // are aclTensor inputs whose values live in Device GM), so we default it
208+ // to 1 and rely on the kernel to take a runtime fast-path when l1 == 0.
209+ // The HAS_L1 = 0 binary is still produced so UT can drive TilingKey
210+ // 10003 directly (and a future L0-API hint can flip this from host).
211+ uint32_t dTypeVar = static_cast<uint32_t>(dataType);
212+ uint32_t padTail = ((totalElements % ubBlockSize) != 0) ? 1U : 0U;
213+ uint32_t hasL1 = 1U;
214+ ASCENDC_TPL_SEL_PARAM(context, dTypeVar, padTail, hasL1);
215+ 
216+ return ge::GRAPH_SUCCESS;
217+}
218+ 
219+static ge::graphStatus TilingParseForApplyProximalAdagrad(
220+ [[maybe_unused]] gert::TilingParseContext* context)
221+{
222+ return ge::GRAPH_SUCCESS;
223+}
224+ 
225+struct ApplyProximalAdagradCompileInfo {};
226+ 
227+IMPL_OP_OPTILING(ApplyProximalAdagrad)
228+ .Tiling(ApplyProximalAdagradTilingFunc)
229+ .TilingParse<ApplyProximalAdagradCompileInfo>(TilingParseForApplyProximalAdagrad);
230+ 
231+} // namespace optiling
Aexperimental/optim/apply_proximal_adagrad/op_kernel/apply_proximal_adagrad.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_proximal_adagrad.cpp
18+ * \brief ApplyProximalAdagrad kernel entry (arch35 / Ascend950).
19+ *
20+ * Signature follows registry-invoke convention:
21+ * all inputs -> all outputs -> workspace -> tiling
22+ *
23+ * Template parameters are driven by ASCENDC_TPL_SEL_PARAM on the host:
24+ * - D_T_VAR : input dtype (currently C_DT_FLOAT only).
25+ * - PAD_TAIL : whether the per-core tail is non-32B-aligned (0/1).
26+ * In iteration-2 the kernel always uses DataCopyPad on the
27+ * trailing tile, so the parameter is reserved for future
28+ * aligned-only fast-path optimisations and serves as the
29+ * TilingKey 10001 vs 10002 discriminator.
30+ * - HAS_L1 : compile-time hint that l1 may be non-zero. When 0, the
31+ * kernel statically drops the sign + soft-threshold branch
32+ * (TilingKey 10003 fast path).
33+ */
34+ 
35+#include "apply_proximal_adagrad.h"
36+ 
37+template <typename D_T_VAR, uint32_t PAD_TAIL, uint32_t HAS_L1>
38+__global__ __aicore__ void apply_proximal_adagrad(
39+ GM_ADDR var, GM_ADDR accum,
40+ GM_ADDR lr, GM_ADDR l1, GM_ADDR l2,
41+ GM_ADDR grad,
42+ GM_ADDR var_out, GM_ADDR accum_out,
43+ GM_ADDR workspace, GM_ADDR tiling)
44+{
45+ REGISTER_TILING_DEFAULT(ApplyProximalAdagradTilingData);
46+ GET_TILING_DATA_WITH_STRUCT(ApplyProximalAdagradTilingData, tilingData, tiling);
47+ NsApplyProximalAdagrad::ApplyProximalAdagrad<D_T_VAR, (PAD_TAIL != 0U), (HAS_L1 != 0U)> op;
48+ op.Init(var, accum, lr, l1, l2, grad, var_out, accum_out, &tilingData);
49+ op.Process();
50+}
Aexperimental/optim/apply_proximal_adagrad/op_kernel/apply_proximal_adagrad.h+435-0
@@ -0,0 +1,435 @@
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_proximal_adagrad.h
18+ * \brief ApplyProximalAdagrad kernel class (arch35 / Ascend950).
19+ *
20+ * Implements the Proximal Adagrad update per element:
21+ *
22+ * accum = accum + grad * grad
23+ * eta = lr * rsqrt(accum)
24+ * prox = var - eta * grad
25+ * if (l1 > 0):
26+ * var = sign(prox) * max(|prox| - eta*l1, 0) / (1 + eta*l2)
27+ * else:
28+ * var = prox / (1 + eta*l2)
29+ *
30+ * Iteration-2 routes:
31+ * - PAD_TAIL: 0 (aligned tail) / 1 (non-aligned tail). Both branches use
32+ * DataCopyPad which is alignment-tolerant; the flag is preserved so future
33+ * iterations can swap in a faster pure-DataCopy fast path on PAD_TAIL=0.
34+ * - HAS_L1: 1 (general path with sign + soft-threshold). The kernel also
35+ * applies a runtime fast-path when l1Scalar_ == 0 inside this branch (this
36+ * is the path taken when host cannot read l1 from GM at tiling time).
37+ * - HAS_L1: 0 (TilingKey 10003 dedicated fast path - sign / soft-threshold
38+ * statically dropped). Reachable via host-side hint or UT.
39+ */
40+ 
41+#ifndef APPLY_PROXIMAL_ADAGRAD_H
42+#define APPLY_PROXIMAL_ADAGRAD_H
43+ 
44+#include "kernel_operator.h"
45+#include "kernel_tiling/kernel_tiling.h"
46+#include "apply_proximal_adagrad_tiling_data.h"
47+#include "apply_proximal_adagrad_tiling_key.h"
48+ 
49+namespace NsApplyProximalAdagrad {
50+ 
51+// SUG-002 fix: encode the UB resident-tensor count as a kernel-side constant so
52+// that any change to the buffer layout below MUST update this value. The host
53+// tiling (apply_proximal_adagrad_tiling.cpp::UB_BUFFER_COUNT) sizes the per-tile
54+// UB capacity using the same constant; the two values MUST be kept in sync, or
55+// the tile size computed by the host may exceed actual UB usage and risk
56+// overflow. The breakdown is:
57+// - 3 IN queues (var / accum / grad) x double-buffer (2) = 6
58+// - 2 OUT queues (var / accum) x double-buffer (2) = 4
59+// - 4 scratch buffers (eta / prox / thresh / scale), single-buffered = 4
60+// total = 14
61+constexpr int64_t kUbResidentFp32TensorCount = 14;
62+static_assert(kUbResidentFp32TensorCount == 14,
63+ "UB resident tensor count must stay in sync with host tiling "
64+ "constant UB_BUFFER_COUNT in apply_proximal_adagrad_tiling.cpp; "
65+ "update both together when adding/removing UB buffers.");
66+ 
67+using AscendC::TPipe;
68+using AscendC::TQue;
69+using AscendC::TBuf;
70+using AscendC::QuePosition;
71+using AscendC::GlobalTensor;
72+using AscendC::LocalTensor;
73+using AscendC::DataCopyParams;
74+using AscendC::DataCopyPad;
75+using AscendC::DataCopyExtParams;
76+using AscendC::DataCopyPadExtParams;
77+using AscendC::GetBlockIdx;
78+using AscendC::Add;
79+using AscendC::Sub;
80+using AscendC::Mul;
81+using AscendC::Muls;
82+using AscendC::Adds;
83+using AscendC::Abs;
84+using AscendC::Div;
85+using AscendC::Rsqrt;
86+using AscendC::Maxs;
87+using AscendC::Compare;
88+using AscendC::Compares;
89+using AscendC::CompareScalar;
90+using AscendC::Select;
91+using AscendC::Duplicate;
92+using AscendC::CMPMODE;
93+using AscendC::SELMODE;
94+ 
95+template <typename T, bool PAD_TAIL = true, bool HAS_L1 = true>
96+class ApplyProximalAdagrad {
97+public:
98+ __aicore__ inline ApplyProximalAdagrad() {}
99+ 
100+ __aicore__ inline void Init(GM_ADDR var, GM_ADDR accum,
101+ GM_ADDR lr, GM_ADDR l1, GM_ADDR l2,
102+ GM_ADDR grad,
103+ GM_ADDR varOut, GM_ADDR accumOut,
104+ const ApplyProximalAdagradTilingData* tilingData);
105+ __aicore__ inline void Process();
106+ 
107+private:
108+ __aicore__ inline void CopyInTile(int64_t gmOffset, int64_t currentNum);
109+ __aicore__ inline void Compute(int64_t currentNum);
110+ __aicore__ inline void CopyOutTile(int64_t gmOffset, int64_t currentNum);
111+ 
112+ __aicore__ inline float LoadScalar(const GlobalTensor<T>& src) const;
113+ 
114+private:
115+ TPipe pipe_;
116+ // IN queues: var / accum / grad, double-buffered.
117+ TQue<QuePosition::VECIN, 2> varInQue_;
118+ TQue<QuePosition::VECIN, 2> accumInQue_;
119+ TQue<QuePosition::VECIN, 2> gradInQue_;
120+ // OUT queues: var / accum, double-buffered.
121+ TQue<QuePosition::VECOUT, 2> varOutQue_;
122+ TQue<QuePosition::VECOUT, 2> accumOutQue_;
123+ // Scratch compute buffers (VECCALC, not queue-synchronised).
124+ TBuf<QuePosition::VECCALC> etaBuf_; // lr * rsqrt(accum)
125+ TBuf<QuePosition::VECCALC> proxBuf_; // var - eta * grad
126+ TBuf<QuePosition::VECCALC> threshBuf_; // max(|prox| - eta*l1, 0) and helpers
127+ TBuf<QuePosition::VECCALC> scaleBuf_; // 1 + eta * l2
128+ 
129+ GlobalTensor<T> varGm_;
130+ GlobalTensor<T> accumGm_;
131+ GlobalTensor<T> gradGm_;
132+ GlobalTensor<T> lrGm_;
133+ GlobalTensor<T> l1Gm_;
134+ GlobalTensor<T> l2Gm_;
135+ GlobalTensor<T> varOutGm_;
136+ GlobalTensor<T> accumOutGm_;
137+ 
138+ // Scalars loaded from GM at Init time.
139+ float lrScalar_ = 0.0f;
140+ float l1Scalar_ = 0.0f;
141+ float l2Scalar_ = 0.0f;
142+ 
143+ int64_t blockOffset_ = 0;
144+ int64_t blockLen_ = 0;
145+ int64_t ubFactor_ = 0;
146+};
147+ 
148+// =============================================================================
149+// LoadScalar: pick the first element of a 1-element GM tensor into a register.
150+// =============================================================================
151+template <typename T, bool PAD_TAIL, bool HAS_L1>
152+__aicore__ inline float ApplyProximalAdagrad<T, PAD_TAIL, HAS_L1>::LoadScalar(
153+ const GlobalTensor<T>& src) const
154+{
155+ return static_cast<float>(src.GetValue(0));
156+}
157+ 
158+// =============================================================================
159+// Init
160+// =============================================================================
161+template <typename T, bool PAD_TAIL, bool HAS_L1>
162+__aicore__ inline void ApplyProximalAdagrad<T, PAD_TAIL, HAS_L1>::Init(
163+ GM_ADDR var, GM_ADDR accum,
164+ GM_ADDR lr, GM_ADDR l1, GM_ADDR l2,
165+ GM_ADDR grad,
166+ GM_ADDR varOut, GM_ADDR accumOut,
167+ const ApplyProximalAdagradTilingData* tilingData)
168+{
169+ ubFactor_ = tilingData->ubFactor;
170+ 
171+ // Empty-tensor / degenerate tiling -> just record zero length; Process()
172+ // will short-circuit.
173+ if (tilingData->totalElements == 0 || tilingData->blockFactor == 0) {
174+ blockOffset_ = 0;
175+ blockLen_ = 0;
176+ return;
177+ }
178+ 
179+ blockOffset_ = tilingData->blockFactor * static_cast<int64_t>(GetBlockIdx());
180+ int64_t remaining = tilingData->totalElements - blockOffset_;
181+ if (remaining <= 0) {
182+ blockLen_ = 0;
183+ return;
184+ }
185+ blockLen_ = (remaining > tilingData->blockFactor)
186+ ? tilingData->blockFactor
187+ : remaining;
188+ 
189+ // Main vectorised tensors -- slice each core's view.
190+ varGm_.SetGlobalBuffer((__gm__ T*)var + blockOffset_, blockLen_);
191+ accumGm_.SetGlobalBuffer((__gm__ T*)accum + blockOffset_, blockLen_);
192+ gradGm_.SetGlobalBuffer((__gm__ T*)grad + blockOffset_, blockLen_);
193+ varOutGm_.SetGlobalBuffer((__gm__ T*)varOut + blockOffset_, blockLen_);
194+ accumOutGm_.SetGlobalBuffer((__gm__ T*)accumOut + blockOffset_, blockLen_);
195+ 
196+ // Scalar tensors: single element broadcast-read.
197+ lrGm_.SetGlobalBuffer((__gm__ T*)lr, 1);
198+ l1Gm_.SetGlobalBuffer((__gm__ T*)l1, 1);
199+ l2Gm_.SetGlobalBuffer((__gm__ T*)l2, 1);
200+ lrScalar_ = LoadScalar(lrGm_);
201+ l1Scalar_ = LoadScalar(l1Gm_);
202+ l2Scalar_ = LoadScalar(l2Gm_);
203+ 
204+ // UB buffer allocation.
205+ pipe_.InitBuffer(varInQue_, 2, ubFactor_ * sizeof(T));
206+ pipe_.InitBuffer(accumInQue_, 2, ubFactor_ * sizeof(T));
207+ pipe_.InitBuffer(gradInQue_, 2, ubFactor_ * sizeof(T));
208+ pipe_.InitBuffer(varOutQue_, 2, ubFactor_ * sizeof(T));
209+ pipe_.InitBuffer(accumOutQue_, 2, ubFactor_ * sizeof(T));
210+ 
211+ pipe_.InitBuffer(etaBuf_, ubFactor_ * sizeof(float));
212+ pipe_.InitBuffer(proxBuf_, ubFactor_ * sizeof(float));
213+ pipe_.InitBuffer(threshBuf_, ubFactor_ * sizeof(float));
214+ pipe_.InitBuffer(scaleBuf_, ubFactor_ * sizeof(float));
215+}
216+ 
217+// =============================================================================
218+// CopyInTile: pad-aware DataCopyPad of var / accum / grad.
219+// DataCopyPad transparently handles aligned and non-aligned blockLen, so the
220+// same primitive serves both PAD_TAIL=0 (aligned tile) and PAD_TAIL=1
221+// (non-aligned tail) routes.
222+// =============================================================================
223+template <typename T, bool PAD_TAIL, bool HAS_L1>
224+__aicore__ inline void ApplyProximalAdagrad<T, PAD_TAIL, HAS_L1>::CopyInTile(
225+ int64_t gmOffset, int64_t currentNum)
226+{
227+ LocalTensor<T> varLocal = varInQue_.template AllocTensor<T>();
228+ LocalTensor<T> accumLocal = accumInQue_.template AllocTensor<T>();
229+ LocalTensor<T> gradLocal = gradInQue_.template AllocTensor<T>();
230+ 
231+ DataCopyExtParams copyParams;
232+ copyParams.blockCount = 1;
233+ copyParams.blockLen = static_cast<uint32_t>(currentNum * sizeof(T));
234+ copyParams.srcStride = 0;
235+ copyParams.dstStride = 0;
236+ 
237+ // ISSUE-001 fix: defensively right-pad the tail lanes ([currentNum, alignedNum))
238+ // with safe values so Compute's Rsqrt/Div on padding lanes does not produce
239+ // Inf/NaN. We compute padding count up to the next 8-element (32B) block.
240+ // The padding lanes are never written back to GM (DataCopyPad honours the
241+ // exact byte length), so the value choice only matters for in-UB compute
242+ // stability:
243+ // - var / grad: pad with 0.0f -> grad=0 keeps accum unchanged, prox=var.
244+ // - accum: pad with 1.0f -> Rsqrt(accum + grad^2) = Rsqrt(1) = 1,
245+ // avoiding the Rsqrt(0)=+Inf -> NaN cascade.
246+ constexpr int64_t kAlignBlock = 32 / sizeof(T);
247+ int64_t alignedNum =
248+ ((currentNum + kAlignBlock - 1) / kAlignBlock) * kAlignBlock;
249+ uint8_t rightPadCount = static_cast<uint8_t>(alignedNum - currentNum);
250+ DataCopyPadExtParams<T> padZeroParams{true, 0, rightPadCount, static_cast<T>(0)};
251+ DataCopyPadExtParams<T> padOneParams{true, 0, rightPadCount, static_cast<T>(1.0f)};
252+ 
253+ DataCopyPad(varLocal, varGm_[gmOffset], copyParams, padZeroParams);
254+ DataCopyPad(accumLocal, accumGm_[gmOffset], copyParams, padOneParams);
255+ DataCopyPad(gradLocal, gradGm_[gmOffset], copyParams, padZeroParams);
256+ 
257+ varInQue_.EnQue(varLocal);
258+ accumInQue_.EnQue(accumLocal);
259+ gradInQue_.EnQue(gradLocal);
260+}
261+ 
262+// =============================================================================
263+// CopyOutTile: write back var / accum to their inplace GM slots. DataCopyPad
264+// honours the exact byte length, so non-aligned tails do not over-write
265+// neighbouring cores' data.
266+// =============================================================================
267+template <typename T, bool PAD_TAIL, bool HAS_L1>
268+__aicore__ inline void ApplyProximalAdagrad<T, PAD_TAIL, HAS_L1>::CopyOutTile(
269+ int64_t gmOffset, int64_t currentNum)
270+{
271+ LocalTensor<T> varOutLocal = varOutQue_.template DeQue<T>();
272+ LocalTensor<T> accumOutLocal = accumOutQue_.template DeQue<T>();
273+ 
274+ DataCopyExtParams copyParams;
275+ copyParams.blockCount = 1;
276+ copyParams.blockLen = static_cast<uint32_t>(currentNum * sizeof(T));
277+ copyParams.srcStride = 0;
278+ copyParams.dstStride = 0;
279+ 
280+ DataCopyPad(varOutGm_[gmOffset], varOutLocal, copyParams);
281+ DataCopyPad(accumOutGm_[gmOffset], accumOutLocal, copyParams);
282+ 
283+ varOutQue_.FreeTensor(varOutLocal);
284+ accumOutQue_.FreeTensor(accumOutLocal);
285+}
286+ 
287+// =============================================================================
288+// Compute: main per-tile computation (float32 only in iteration 1/2).
289+//
290+// accum' = accum + grad*grad
291+// eta = lr * rsqrt(accum')
292+// prox = var - eta * grad
293+//
294+// HAS_L1 = true (TilingKey 10001 / 10002 - general path):
295+// thresh = max(|prox| - eta*l1, 0)
296+// signed = prox >= 0 ? thresh : -thresh
297+// var' = signed / (1 + eta*l2)
298+// Includes a runtime fast-path: if l1Scalar_ == 0 the kernel skips the
299+// sign / soft-threshold steps and behaves like the HAS_L1=false branch.
300+//
301+// HAS_L1 = false (TilingKey 10003 - dedicated fast path):
302+// var' = prox / (1 + eta*l2)
303+// =============================================================================
304+template <typename T, bool PAD_TAIL, bool HAS_L1>
305+__aicore__ inline void ApplyProximalAdagrad<T, PAD_TAIL, HAS_L1>::Compute(
306+ int64_t currentNum)
307+{
308+ LocalTensor<T> varLocal = varInQue_.template DeQue<T>();
309+ LocalTensor<T> accumLocal = accumInQue_.template DeQue<T>();
310+ LocalTensor<T> gradLocal = gradInQue_.template DeQue<T>();
311+ 
312+ LocalTensor<T> varOutLocal = varOutQue_.template AllocTensor<T>();
313+ LocalTensor<T> accumOutLocal = accumOutQue_.template AllocTensor<T>();
314+ 
315+ LocalTensor<float> etaTmp = etaBuf_.template Get<float>();
316+ LocalTensor<float> proxTmp = proxBuf_.template Get<float>();
317+ LocalTensor<float> threshTmp = threshBuf_.template Get<float>();
318+ LocalTensor<float> scaleTmp = scaleBuf_.template Get<float>();
319+ 
320+ // Align work count to 32B / sizeof(float) = 8 elements. This guarantees
321+ // Compare / Select 256-byte alignment (ubFactor is a 2048-elem multiple,
322+ // and we round non-aligned tails up to the next 8-elem block; the extra
323+ // padding lanes are computed but never written back to GM thanks to
324+ // DataCopyPad honouring the exact byte length).
325+ constexpr int64_t kAlignBlock = 32 / sizeof(float);
326+ int64_t alignedNum =
327+ ((currentNum + kAlignBlock - 1) / kAlignBlock) * kAlignBlock;
328+ int32_t n = static_cast<int32_t>(alignedNum);
329+ 
330+ // Compile-time specialisation: iteration 1/2 only supports float32. If
331+ // new dtypes are added later they should Cast at the beginning and Cast
332+ // back at the end, leaving the inner block below unchanged.
333+ if constexpr (std::is_same_v<T, float>) {
334+ // ----- Common: S1 / S2 / S3 -----
335+ // S1: accum' = accum + grad*grad -> write into accumOutLocal.
336+ Mul(accumOutLocal, gradLocal, gradLocal, n);
337+ Add(accumOutLocal, accumLocal, accumOutLocal, n);
338+ 
339+ // S2: eta = lr * rsqrt(accum')
340+ // NOTE: Caller contract guarantees accum >= 0 (typically > 0). If
341+ // accum + grad^2 == 0, Rsqrt returns +Inf and downstream eta*grad
342+ // becomes Inf*0 = NaN. This matches the semantics of PyTorch's
343+ // ApplyProximalAdagrad and TensorFlow's ApplyProximalAdagrad, where
344+ // the safe-input contract is the caller's responsibility (see
345+ // README.md / DESIGN.md). Defensive padding in CopyInTile ensures
346+ // the [currentNum, alignedNum) tail lanes use accum=1, not 0, so
347+ // padding lanes never trigger this path.
348+ Rsqrt(etaTmp, accumOutLocal, n);
349+ Muls(etaTmp, etaTmp, lrScalar_, n);
350+ 
351+ // S3: prox = var - eta * grad
352+ Mul(proxTmp, etaTmp, gradLocal, n);
353+ Sub(proxTmp, varLocal, proxTmp, n);
354+ 
355+ // ----- S4 / S5 branch -----
356+ if constexpr (HAS_L1) {
357+ // Runtime fast-path: if l1 == 0 the sign + soft-threshold steps
358+ // collapse to identity, so we behave like the HAS_L1=false branch.
359+ if (l1Scalar_ == 0.0f) {
360+ // scale = 1 + eta*l2; var = prox / scale
361+ Muls(scaleTmp, etaTmp, l2Scalar_, n);
362+ Adds(scaleTmp, scaleTmp, 1.0f, n);
363+ Div(varOutLocal, proxTmp, scaleTmp, n);
364+ } else {
365+ // S4a: thresh = max(|prox| - eta*l1, 0)
366+ Abs(threshTmp, proxTmp, n);
367+ Muls(scaleTmp, etaTmp, l1Scalar_, n); // scaleTmp reused as eta*l1
368+ Sub(threshTmp, threshTmp, scaleTmp, n);
369+ Maxs(threshTmp, threshTmp, 0.0f, n);
370+ 
371+ // S4b: signed_thresh = prox >= 0 ? thresh : -thresh
372+ Muls(scaleTmp, threshTmp, -1.0f, n); // scaleTmp = -thresh
373+ // SUG-001 NOTE: We temporarily reinterpret varOutLocal's storage
374+ // as a uint8_t mask buffer. This relies on an implicit ordering
375+ // invariant: varOutLocal must NOT be read between this point and
376+ // the final Div() below that overwrites it. The mask is consumed
377+ // by the immediately following Select() call and never read again
378+ // in this Compute(). If future maintenance inserts any read of
379+ // varOutLocal before the final Div, this aliasing must be
380+ // replaced by a dedicated mask buffer (e.g. add a maskBuf_ in
381+ // Init or carve out the tail of threshBuf_).
382+ LocalTensor<uint8_t> maskTensor =
383+ varOutLocal.template ReinterpretCast<uint8_t>();
384+ CompareScalar(maskTensor, proxTmp, 0.0f, CMPMODE::GE, n);
385+ Select(proxTmp, maskTensor, threshTmp, scaleTmp,
386+ SELMODE::VSEL_TENSOR_TENSOR_MODE, n);
387+ 
388+ // S5: scale = 1 + eta*l2
389+ Muls(scaleTmp, etaTmp, l2Scalar_, n);
390+ Adds(scaleTmp, scaleTmp, 1.0f, n);
391+ 
392+ // var' = signed_thresh / scale
393+ Div(varOutLocal, proxTmp, scaleTmp, n);
394+ }
395+ } else {
396+ // HAS_L1 == false (TilingKey 10003): dedicated simplified path.
397+ // var = prox / (1 + eta*l2)
398+ Muls(scaleTmp, etaTmp, l2Scalar_, n);
399+ Adds(scaleTmp, scaleTmp, 1.0f, n);
400+ Div(varOutLocal, proxTmp, scaleTmp, n);
401+ }
402+ }
403+ 
404+ varOutQue_.template EnQue<T>(varOutLocal);
405+ accumOutQue_.template EnQue<T>(accumOutLocal);
406+ 
407+ varInQue_.FreeTensor(varLocal);
408+ accumInQue_.FreeTensor(accumLocal);
409+ gradInQue_.FreeTensor(gradLocal);
410+}
411+ 
412+// =============================================================================
413+// Process: main loop over UB-sized chunks.
414+// =============================================================================
415+template <typename T, bool PAD_TAIL, bool HAS_L1>
416+__aicore__ inline void ApplyProximalAdagrad<T, PAD_TAIL, HAS_L1>::Process()
417+{
418+ if (blockLen_ <= 0) {
419+ return;
420+ }
421+ int64_t loopCount = (blockLen_ + ubFactor_ - 1) / ubFactor_;
422+ for (int64_t i = 0; i < loopCount; i++) {
423+ int64_t gmOffset = i * ubFactor_;
424+ int64_t currentNum = (i == (loopCount - 1))
425+ ? (blockLen_ - gmOffset)
426+ : ubFactor_;
427+ CopyInTile(gmOffset, currentNum);
428+ Compute(currentNum);
429+ CopyOutTile(gmOffset, currentNum);
430+ }
431+}
432+ 
433+} // namespace NsApplyProximalAdagrad
434+ 
435+#endif // APPLY_PROXIMAL_ADAGRAD_H
Aexperimental/optim/apply_proximal_adagrad/op_kernel/apply_proximal_adagrad_tiling_data.h+37-0
@@ -0,0 +1,37 @@
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_proximal_adagrad_tiling_data.h
18+ * \brief ApplyProximalAdagrad TilingData structure (arch35).
19+ *
20+ * Standard C++ struct form (per REQUIREMENTS s8.4 "禁用 BEGIN_TILING_DATA_DEF").
21+ */
22+ 
23+#ifndef _APPLY_PROXIMAL_ADAGRAD_TILING_DATA_H_
24+#define _APPLY_PROXIMAL_ADAGRAD_TILING_DATA_H_
25+ 
26+#include <cstdint>
27+ 
28+struct ApplyProximalAdagradTilingData {
29+ // Total number of elements in var/accum/grad (all share the same shape).
30+ int64_t totalElements = 0;
31+ // Number of elements per-core (aligned up to ubBlockSize to keep DMA safe).
32+ int64_t blockFactor = 0;
33+ // Number of elements processed per UB iteration.
34+ int64_t ubFactor = 0;
35+};
36+ 
37+#endif // _APPLY_PROXIMAL_ADAGRAD_TILING_DATA_H_
Aexperimental/optim/apply_proximal_adagrad/op_kernel/apply_proximal_adagrad_tiling_key.h+60-0
@@ -0,0 +1,60 @@
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_proximal_adagrad_tiling_key.h
18+ * \brief ApplyProximalAdagrad TilingKey template-argument declaration (arch35).
19+ *
20+ * Iteration-2 introduces three TilingKey routes (DESIGN s5.5):
21+ * 10001 - aligned tail block + L1 >= 0 general path
22+ * 10002 - non-aligned tail block + L1 >= 0 general path (DataCopyPad write-back)
23+ * 10003 - any tail block + L1 == 0 fast path (skip sign/soft-threshold)
24+ *
25+ * They are encoded via two boolean template parameters (PAD_TAIL, HAS_L1):
26+ * (PAD_TAIL=0, HAS_L1=1) -> 10001
27+ * (PAD_TAIL=1, HAS_L1=1) -> 10002
28+ * (PAD_TAIL=0|1, HAS_L1=0) -> 10003
29+ *
30+ * - PAD_TAIL is derived in host tiling from (totalElements % ubBlockSize != 0).
31+ * - HAS_L1 is normally driven from host = 1 (lr/l1/l2 sit in Device GM and we
32+ * cannot synchronously inspect them at tiling time). The kernel additionally
33+ * performs a cheap runtime fast-path check when l1Scalar == 0. We still
34+ * register the HAS_L1 = 0 binary so UT (and a future L0-API hint) can drive
35+ * the dedicated TilingKey 10003.
36+ *
37+ * Uses ASCENDC_TPL_ARGS_DECL template-argument mechanism.
38+ * TILING_KEY_IS macro is forbidden (see REQUIREMENTS s8.4).
39+ */
40+ 
41+#ifndef __APPLY_PROXIMAL_ADAGRAD_TILING_KEY_H__
42+#define __APPLY_PROXIMAL_ADAGRAD_TILING_KEY_H__
43+ 
44+#include "ascendc/host_api/tiling/template_argument.h"
45+ 
46+ASCENDC_TPL_ARGS_DECL(ApplyProximalAdagrad,
47+ ASCENDC_TPL_DATATYPE_DECL(D_T_VAR, C_DT_FLOAT, ASCENDC_TPL_INPUT(0)),
48+ ASCENDC_TPL_UINT_DECL(PAD_TAIL, 8, ASCENDC_TPL_UI_LIST, 0, 1),
49+ ASCENDC_TPL_UINT_DECL(HAS_L1, 8, ASCENDC_TPL_UI_LIST, 0, 1)
50+);
51+ 
52+ASCENDC_TPL_SEL(
53+ ASCENDC_TPL_ARGS_SEL(
54+ ASCENDC_TPL_DATATYPE_SEL(D_T_VAR, C_DT_FLOAT),
55+ ASCENDC_TPL_UINT_SEL(PAD_TAIL, ASCENDC_TPL_UI_LIST, 0, 1),
56+ ASCENDC_TPL_UINT_SEL(HAS_L1, ASCENDC_TPL_UI_LIST, 0, 1)
57+ ),
58+);
59+ 
60+#endif // __APPLY_PROXIMAL_ADAGRAD_TILING_KEY_H__