已合并
【特性】新增fusedSgd #6276
majiajian_hw创建于 6月18日
【特性】新增fusedSgd #6276
已合并
共 32 个文件变更+4281-0
| @@ -2023,6 +2023,14 @@ loss-c@ops-nn: | |||
| 2023 | - ops/ops-nn/optim/apply_adam_w_v2/op_kernel/apply_adam_w_v2_fp.h | 2023 | - ops/ops-nn/optim/apply_adam_w_v2/op_kernel/apply_adam_w_v2_fp.h |
| 2024 | - ops/ops-nn/optim/apply_adam_w_v2/op_kernel/apply_adam_w_v2_mix_dtype.h | 2024 | - ops/ops-nn/optim/apply_adam_w_v2/op_kernel/apply_adam_w_v2_mix_dtype.h |
| 2025 | - ops/ops-nn/optim/apply_adam_w_v2/op_kernel/apply_adam_w_v2.cpp | 2025 | - ops/ops-nn/optim/apply_adam_w_v2/op_kernel/apply_adam_w_v2.cpp |
| 2026 | - ops/ops-nn/optim/fused_sgd/docs/ | ||
| 2027 | - ops/ops-nn/optim/fused_sgd/README.md | ||
| 2028 | - ops/ops-nn/optim/fused_sgd/op_host/ | ||
| 2029 | - ops/ops-nn/optim/fused_sgd/op_kernel/fused_sgd_base.h | ||
| 2030 | - ops/ops-nn/optim/fused_sgd/op_kernel/fused_sgd_f16_bf16.h | ||
| 2031 | - ops/ops-nn/optim/fused_sgd/op_kernel/fused_sgd_f32.h | ||
| 2032 | - ops/ops-nn/optim/fused_sgd/op_kernel/fused_sgd_tiling_data.h | ||
| 2033 | - ops/ops-nn/optim/fused_sgd/op_kernel/fused_sgd.cpp | ||
| 2026 | - ops/ops-nn/optim/apply_fused_ema_adam/ | 2034 | - ops/ops-nn/optim/apply_fused_ema_adam/ |
| 2027 | opensource_style: null | 2035 | opensource_style: null |
| 2028 | kernel_style: null | 2036 | kernel_style: null |
| @@ -2089,6 +2097,8 @@ loss-c@ops-nn: | |||
| 2089 | - ops/ops-nn/optim/apply_adam_w_quant/tests/ | 2097 | - ops/ops-nn/optim/apply_adam_w_quant/tests/ |
| 2090 | - ops/ops-nn/optim/apply_adam_w_v2/examples/ | 2098 | - ops/ops-nn/optim/apply_adam_w_v2/examples/ |
| 2091 | - ops/ops-nn/optim/apply_adam_w_v2/tests/ | 2099 | - ops/ops-nn/optim/apply_adam_w_v2/tests/ |
| 2100 | - ops/ops-nn/optim/fused_sgd/examples/ | ||
| 2101 | - ops/ops-nn/optim/fused_sgd/tests/ | ||
| 2092 | - ops/ops-nn/optim/apply_fused_ema_adam/examples/ | 2102 | - ops/ops-nn/optim/apply_fused_ema_adam/examples/ |
| 2093 | - ops/ops-nn/optim/apply_fused_ema_adam/tests/ | 2103 | - ops/ops-nn/optim/apply_fused_ema_adam/tests/ |
| 2094 | non_delivery: null | 2104 | non_delivery: null |
| @@ -3289,6 +3289,16 @@ | |||
| 3289 | <td>AI Core</td> | 3289 | <td>AI Core</td> |
| 3290 | <td>结合Adagrad自适应学习率与FOBOS（Forward-Backward Splitting）Proximal近端算法的优化器，基于梯度平方累加器自适应调整学习率，并通过软阈值（L1正则化）与缩放（L2正则化）对模型参数进行原地更新。</td> | 3290 | <td>结合Adagrad自适应学习率与FOBOS（Forward-Backward Splitting）Proximal近端算法的优化器，基于梯度平方累加器自适应调整学习率，并通过软阈值（L1正则化）与缩放（L2正则化）对模型参数进行原地更新。</td> |
| 3291 | </tr> | 3291 | </tr> |
| 3292 | <tr> | ||
| 3293 | <td>optim</td> | ||
| 3294 | <td><a href="../../optim/fused_sgd/README.md">fused_sgd</a></td> | ||
 | |||
| 3295 | <td>✓</td> | ||
| 3296 | <td>✓</td> | ||
| 3297 | <td>✓</td> | ||
| 3298 | <td>✗</td> | ||
| 3299 | <td>AI Core</td> | ||
| 3300 | <td>实现FusedSgd融合优化器功能。</td> | ||
| 3301 | </tr> | ||
| 3292 | <tr> | 3302 | <tr> |
| 3293 | <td>pooling</td> | 3303 | <td>pooling</td> |
| 3294 | <td><a href="../../pooling/adaptive_avg_pool3d/README.md">adaptive_avg_pool3d</a></td> | 3304 | <td><a href="../../pooling/adaptive_avg_pool3d/README.md">adaptive_avg_pool3d</a></td> |
| @@ -0,0 +1,11 @@ | |||
| 1 | # ---------------------------------------------------------------------------- | ||
| 2 | # Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | # This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | # CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | # Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | # See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | # ---------------------------------------------------------------------------- | ||
| 10 | message(STATUS "=== Debug: start fused_sgd.CMakeLists.txt ") | ||
| 11 | add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE fused_sgd ACLNNTYPE aclnn_exclude) | ||
| @@ -0,0 +1,167 @@ | |||
| 1 | # FusedSgd | ||
| 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 | - **算子功能**：实现fusedSgd算子。将传统 SGD 更新过程中原本分散的多个细粒度操作（如梯度缩放、权重衰减、动量更新、参数赋值等）融合为单个NPU Kernel执行。 | ||
| 17 | - **计算公式**： | ||
| 18 | |||
| 19 | $$ | ||
| 20 | \begin{aligned} | ||
| 21 | &\tilde{g}_t = \begin{cases} | ||
| 22 | g_t / s & s \neq \text{None} \\ | ||
| 23 | g_t & \text{otherwise} | ||
| 24 | \end{cases} \\ | ||
| 25 | |||
| 26 | &\hat{g}_t = \begin{cases} | ||
| 27 | -\tilde{g}_t & \text{maximize} \\ | ||
| 28 | \tilde{g}_t & \text{otherwise} | ||
| 29 | \end{cases} \\ | ||
| 30 | |||
| 31 | &\bar{g}_t = \hat{g}_t + weightDecay \cdot \theta_t \\ | ||
| 32 | |||
| 33 | &v_{t+1} = \begin{cases} | ||
| 34 | \bar{g}_t & \text{first step} \\ | ||
| 35 | \mu v_t + (1-dampening)\bar{g}_t & \text{otherwise} | ||
| 36 | \end{cases} \\ | ||
| 37 | |||
| 38 | &g_t^{\text{final}} = \begin{cases} | ||
| 39 | \bar{g}_t + \mu v_{t+1} & \text{nesterov} \\ | ||
| 40 | v_{t+1} & v_t \neq \text{None} \\ | ||
| 41 | \bar{g}_t & \text{otherwise} | ||
| 42 | \end{cases} \\ | ||
| 43 | |||
| 44 | &\theta_{t+1} = \theta_t - lr \cdot g_t^{\text{final}} \\ | ||
| 45 | &g_{t+1} = \tilde{g}_t | ||
| 46 | |||
| 47 | \end{aligned} | ||
| 48 | $$ | ||
| 49 | |||
| 50 | ## 参数说明 | ||
| 51 | |||
| 52 | <table style="undefined;table-layout: fixed; width: 1080px"><colgroup> | ||
| 53 | <col style="width: 155px"> | ||
| 54 | <col style="width: 162px"> | ||
| 55 | <col style="width: 380px"> | ||
| 56 | <col style="width: 276px"> | ||
| 57 | <col style="width: 107px"> | ||
| 58 | </colgroup> | ||
| 59 | <thead> | ||
| 60 | <tr> | ||
| 61 | <th>参数名</th> | ||
| 62 | <th>输入/输出/属性</th> | ||
| 63 | <th>描述</th> | ||
| 64 | <th>数据类型</th> | ||
| 65 | <th>数据格式</th> | ||
| 66 | </tr></thead> | ||
| 67 | <tbody> | ||
| 68 | <tr> | ||
| 69 | <td>params</td> | ||
| 70 | <td>输入/输出</td> | ||
| 71 | <td>待更新参数，对应公式中的θ。</td> | ||
| 72 | <td>BFLOAT16、FLOAT16、FLOAT</td> | ||
| 73 | <td>ND</td> | ||
| 74 | </tr> | ||
| 75 | <tr> | ||
| 76 | <td>grads</td> | ||
| 77 | <td>输入/输出</td> | ||
| 78 | <td>待更新参数对应的梯度，对应公式中的g。</td> | ||
| 79 | <td>BFLOAT16、FLOAT16、FLOAT</td> | ||
| 80 | <td>ND</td> | ||
| 81 | </tr> | ||
| 82 | <tr> | ||
| 83 | <td>x</td> | ||
| 84 | <td>输入</td> | ||
| 85 | <td>待更新参数对应的动量，对应公式中的v。</td> | ||
| 86 | <td>BFLOAT16、FLOAT16、FLOAT</td> | ||
| 87 | <td>ND</td> | ||
| 88 | </tr> | ||
| 89 | <tr> | ||
| 90 | <td>gradScale</td> | ||
| 91 | <td>输入</td> | ||
| 92 | <td>梯度缩放大小，对应公式中的s。</td> | ||
| 93 | <td>FLOAT</td> | ||
| 94 | <td>ND</td> | ||
| 95 | </tr> | ||
| 96 | <tr> | ||
| 97 | <td>y</td> | ||
| 98 | <td>输出</td> | ||
| 99 | <td>待更新参数对应的动量，对应公式中的v。</td> | ||
| 100 | <td>BFLOAT16、FLOAT16、FLOAT</td> | ||
| 101 | <td>ND</td> | ||
| 102 | </tr> | ||
| 103 | <tr> | ||
| 104 | <td>weightDecay</td> | ||
| 105 | <td>属性</td> | ||
| 106 | <td>权重衰减值，对应公式中的weightDecay，默认为0。</td> | ||
| 107 | <td>FLOAT</td> | ||
| 108 | <td>-</td> | ||
| 109 | </tr> | ||
| 110 | <tr> | ||
| 111 | <td>momentum</td> | ||
| 112 | <td>属性</td> | ||
| 113 | <td>动量值，对应公式中的μ，默认为0。</td> | ||
| 114 | <td>FLOAT</td> | ||
| 115 | <td>-</td> | ||
| 116 | </tr> | ||
| 117 | <tr> | ||
| 118 | <td>lr</td> | ||
| 119 | <td>属性</td> | ||
| 120 | <td>学习率，对应公式中的lr，默认为1e-3。</td> | ||
| 121 | <td>FLOAT</td> | ||
| 122 | <td>-</td> | ||
| 123 | </tr> | ||
| 124 | <tr> | ||
| 125 | <td>dampening</td> | ||
| 126 | <td>属性</td> | ||
| 127 | <td>动量阻尼系数，对应公式中的dampening，默认为0。</td> | ||
| 128 | <td>FLOAT</td> | ||
| 129 | <td>-</td> | ||
| 130 | </tr> | ||
| 131 | <tr> | ||
| 132 | <td>nesterov</td> | ||
| 133 | <td>属性</td> | ||
| 134 | <td>是否启用Nesterov动量，对应公式中的nesterov，默认为False。</td> | ||
| 135 | <td>BOOL</td> | ||
| 136 | <td>-</td> | ||
| 137 | </tr> | ||
| 138 | <tr> | ||
| 139 | <td>maximize</td> | ||
| 140 | <td>属性</td> | ||
| 141 | <td>是否为最大化目标函数，对应公式中的maximize，默认为False。</td> | ||
| 142 | <td>BOOL</td> | ||
| 143 | <td>-</td> | ||
| 144 | </tr> | ||
| 145 | <tr> | ||
| 146 | <td>isFirstStep</td> | ||
| 147 | <td>属性</td> | ||
| 148 | <td>是否第一步更新，对应公式中的FirstStep，默认为True。</td> | ||
| 149 | <td>BOOL</td> | ||
| 150 | <td>-</td> | ||
| 151 | </tr> | ||
| 152 | </tbody></table> | ||
| 153 | |||
| 154 | ## 约束说明 | ||
| 155 | |||
| 156 | - params、grads、x、gradScale的数据类型在支持的范围之内。 | ||
| 157 | - params、grads与x及其中各个tensor具有相同的数据类型。 | ||
| 158 | - params、grads与x中tensor的shape维度小于等于8，gradScale的shape为[1]。 | ||
| 159 | - lr、momentum、weightDecay、dampening的值大于等于0。 | ||
| 160 | - params、grads与x（x不为空时）中相同索引tensor的shape相同。 | ||
| 161 | - params、grads与x（x不为空时）中不能有空指针。 | ||
| 162 | |||
| 163 | ## 调用说明 | ||
| 164 | |||
| 165 | | 调用方式 | 样例代码 | 说明 | | ||
| 166 | | --------- | ------------------------------------------------------------ | ------------------------------------------------------------ | | ||
| 167 | | aclnn接口 | [test_aclnn_fused_sgd](./examples/test_aclnn_fused_sgd.cpp) | 通过[aclnnFusedSgd](docs/aclnnFusedSgd.md)接口方式调用FusedSgd算子。 | | ||
| @@ -0,0 +1,555 @@ | |||
| 1 | # aclnnFusedSgd | ||
| 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 | - 算子功能：实现fusedSgd算子。将传统 SGD 更新过程中原本分散的多个细粒度操作（如梯度缩放、权重衰减、动量更新、参数赋值等）融合为单个NPU Kernel执行。 | ||
| 17 | - 计算公式： | ||
| 18 | |||
| 19 | $$ | ||
| 20 | \begin{aligned} | ||
| 21 | &\tilde{g}_t = \begin{cases} | ||
| 22 | g_t / s & s \neq \text{None} \\ | ||
| 23 | g_t & \text{otherwise} | ||
| 24 | \end{cases} \\ | ||
| 25 | |||
| 26 | &\hat{g}_t = \begin{cases} | ||
| 27 | -\tilde{g}_t & \text{maximize} \\ | ||
| 28 | \tilde{g}_t & \text{otherwise} | ||
| 29 | \end{cases} \\ | ||
| 30 | |||
| 31 | &\bar{g}_t = \hat{g}_t + weightDecay \cdot \theta_t \\ | ||
| 32 | |||
| 33 | &v_{t+1} = \begin{cases} | ||
| 34 | \bar{g}_t & \text{first step} \\ | ||
| 35 | \mu v_t + (1-dampening)\bar{g}_t & \text{otherwise} | ||
| 36 | \end{cases} \\ | ||
| 37 | |||
| 38 | &g_t^{\text{final}} = \begin{cases} | ||
| 39 | \bar{g}_t + \mu v_{t+1} & \text{nesterov} \\ | ||
| 40 | v_{t+1} & v_t \neq \text{None} \\ | ||
| 41 | \bar{g}_t & \text{otherwise} | ||
| 42 | \end{cases} \\ | ||
| 43 | |||
| 44 | &\theta_{t+1} = \theta_t - lr \cdot g_t^{\text{final}} \\ | ||
| 45 | &g_{t+1} = \tilde{g}_t | ||
| 46 | |||
| 47 | \end{aligned} | ||
| 48 | $$ | ||
| 49 | |||
| 50 | ## 函数原型 | ||
| 51 | |||
| 52 | 算子执行接口为[两段式接口](../../../docs/zh/context/两段式接口.md)，必须先调用“aclnnFusedSgdGetWorkspaceSize”接口获取入参并根据计算流程计算所需workspace大小，再调用“aclnnFusedSgd”接口执行计算。 | ||
| 53 | |||
| 54 | ```cpp | ||
| 55 | aclnnStatus aclnnFusedSgdGetWorkspaceSize( | ||
| 56 | const aclTensorList *paramsRef, | ||
| 57 | const aclTensorList *gradsRef, | ||
| 58 | const aclTensorList *momentumBufferListOptionalRef, | ||
| 59 | const aclTensor *gradScaleOptional, | ||
| 60 | float weightDecay, | ||
| 61 | float momentum, | ||
| 62 | float lr, | ||
| 63 | float dampening, | ||
| 64 | bool nesterov, | ||
| 65 | bool maximize, | ||
| 66 | bool isFirstStep, | ||
| 67 | uint64_t *workspaceSize, | ||
| 68 | aclOpExecutor **executor); | ||
| 69 | ``` | ||
| 70 | |||
| 71 | ```cpp | ||
| 72 | aclnnStatus aclnnFusedSgd( | ||
| 73 | void *workspace, | ||
| 74 | uint64_t workspace_size, | ||
| 75 | aclOpExecutor *executor, | ||
| 76 | const aclrtStream stream) | ||
| 77 | ``` | ||
| 78 | |||
| 79 | ## aclnnFusedSgdGetWorkspaceSize | ||
| 80 | |||
| 81 | - **参数说明** | ||
| 82 | |||
| 83 | <div style="overflow-x: auto;"> | ||
| 84 | <table style="undefined;table-layout: fixed; width: 1497px"><colgroup> | ||
| 85 | <col style="width: 150px"> | ||
| 86 | <col style="width: 120px"> | ||
| 87 | <col style="width: 300px"> | ||
| 88 | <col style="width: 330px"> | ||
| 89 | <col style="width: 212px"> | ||
| 90 | <col style="width: 100px"> | ||
| 91 | <col style="width: 140px"> | ||
| 92 | <col style="width: 145px"> | ||
| 93 | </colgroup> | ||
| 94 | <thead> | ||
| 95 | <tr> | ||
| 96 | <th>参数名</th> | ||
| 97 | <th>输入/输出</th> | ||
| 98 | <th>描述</th> | ||
| 99 | <th>使用说明</th> | ||
| 100 | <th>数据类型</th> | ||
| 101 | <th>数据格式</th> | ||
| 102 | <th>维度(shape)</th> | ||
| 103 | <th>非连续Tensor</th> | ||
| 104 | </tr></thead> | ||
| 105 | <tbody> | ||
| 106 | <tr> | ||
| 107 | <td>paramsRef</td> | ||
| 108 | <td>输入/输出</td> | ||
| 109 | <td>更新参数。</td> | ||
| 110 | <td>不支持空Tensor。</td> | ||
| 111 | <td>BFLOAT16、FLOAT16、FLOAT32</td> | ||
| 112 | <td>ND</td> | ||
| 113 | <td>小于等于8维</td> | ||
| 114 | <td>√</td> | ||
| 115 | </tr> | ||
| 116 | <tr> | ||
| 117 | <td>gradsRef</td> | ||
| 118 | <td>输入/输出</td> | ||
| 119 | <td>更新梯度</td> | ||
| 120 | <td>不支持空Tensor。</td> | ||
| 121 | <td>BFLOAT16、FLOAT16、FLOAT32</td> | ||
| 122 | <td>ND</td> | ||
| 123 | <td>小于等于8维</td> | ||
| 124 | <td>√</td> | ||
| 125 | </tr> | ||
| 126 | <tr> | ||
| 127 | <td>momentumBufferListOptionalRef</td> | ||
| 128 | <td>输入/输出</td> | ||
| 129 | <td>更新动量</td> | ||
| 130 | <td>支持空Tensor。</td> | ||
| 131 | <td>BFLOAT16、FLOAT16、FLOAT32</td> | ||
| 132 | <td>ND</td> | ||
| 133 | <td>小于等于8维</td> | ||
| 134 | <td>√</td> | ||
| 135 | </tr> | ||
| 136 | <tr> | ||
| 137 | <td>gradScaleOptional</td> | ||
| 138 | <td>输入</td> | ||
| 139 | <td>梯度缩放值</td> | ||
| 140 | <td>支持空Tensor。</td> | ||
| 141 | <td>BFLOAT16、FLOAT16、FLOAT32</td> | ||
| 142 | <td>ND</td> | ||
| 143 | <td>shape等于[1]</td> | ||
| 144 | <td>√</td> | ||
| 145 | </tr> | ||
| 146 | <tr> | ||
| 147 | <td>weightDecay</td> | ||
| 148 | <td>属性</td> | ||
| 149 | <td>权重衰减值。</td> | ||
| 150 | <td>-</td> | ||
| 151 | <td>FLOAT</td> | ||
| 152 | <td>-</td> | ||
| 153 | <td>-</td> | ||
| 154 | <td>-</td> | ||
| 155 | </tr> | ||
| 156 | <tr> | ||
| 157 | <td>momentum</td> | ||
| 158 | <td>属性</td> | ||
| 159 | <td>动量值。</td> | ||
| 160 | <td>-</td> | ||
| 161 | <td>FLOAT</td> | ||
| 162 | <td>-</td> | ||
| 163 | <td>-</td> | ||
| 164 | <td>-</td> | ||
| 165 | </tr> | ||
| 166 | <tr> | ||
| 167 | <td>lr</td> | ||
| 168 | <td>属性</td> | ||
| 169 | <td>学习率。</td> | ||
| 170 | <td>-</td> | ||
| 171 | <td>FLOAT</td> | ||
| 172 | <td>-</td> | ||
| 173 | <td>-</td> | ||
| 174 | <td>-</td> | ||
| 175 | </tr> | ||
| 176 | <tr> | ||
| 177 | <td>dampening</td> | ||
| 178 | <td>属性</td> | ||
| 179 | <td>动量的阻尼系数。</td> | ||
| 180 | <td>-</td> | ||
| 181 | <td>FLOAT</td> | ||
| 182 | <td>-</td> | ||
| 183 | <td>-</td> | ||
| 184 | <td>-</td> | ||
| 185 | </tr> | ||
| 186 | <tr> | ||
| 187 | <td>nesterov</td> | ||
| 188 | <td>属性</td> | ||
| 189 | <td>是否启用 Nesterov 动量。</td> | ||
| 190 | <td>-</td> | ||
| 191 | <td>BOOL</td> | ||
| 192 | <td>-</td> | ||
| 193 | <td>-</td> | ||
| 194 | <td>-</td> | ||
| 195 | </tr> | ||
| 196 | <tr> | ||
| 197 | <td>maximize</td> | ||
| 198 | <td>属性</td> | ||
| 199 | <td>是否为最大化目标函数。</td> | ||
| 200 | <td>-</td> | ||
| 201 | <td>BOOL</td> | ||
| 202 | <td>-</td> | ||
| 203 | <td>-</td> | ||
| 204 | <td>-</td> | ||
| 205 | </tr> | ||
| 206 | <tr> | ||
| 207 | <td>isFirstStep</td> | ||
| 208 | <td>属性</td> | ||
| 209 | <td>是否第一步更新。</td> | ||
| 210 | <td>-</td> | ||
| 211 | <td>BOOL</td> | ||
| 212 | <td>-</td> | ||
| 213 | <td>-</td> | ||
| 214 | <td>-</td> | ||
| 215 | </tr> | ||
| 216 | <tr> | ||
| 217 | <td>workspaceSize</td> | ||
| 218 | <td>输出</td> | ||
| 219 | <td>返回用户需要在Device侧申请的workspace大小。</td> | ||
| 220 | <td>-</td> | ||
| 221 | <td>-</td> | ||
| 222 | <td>-</td> | ||
| 223 | <td>1</td> | ||
| 224 | <td>-</td> | ||
| 225 | </tr> | ||
| 226 | <tr> | ||
| 227 | <td>executor</td> | ||
| 228 | <td>输出</td> | ||
| 229 | <td>返回op执行器，包含了算子计算流程。</td> | ||
| 230 | <td>-</td> | ||
| 231 | <td>-</td> | ||
| 232 | <td>-</td> | ||
| 233 | <td>1</td> | ||
| 234 | <td>-</td> | ||
| 235 | </tr> | ||
| 236 | </tbody></table> | ||
| 237 | </div> | ||
| 238 | - **返回值** | ||
| 239 | |||
| 240 | 返回aclnnStatus状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。 | ||
| 241 | |||
| 242 | 第一段接口完成入参校验，出现以下场景时报错： | ||
| 243 | |||
| 244 | <div style="overflow-x: auto;"> | ||
| 245 | <table style="undefined;table-layout: fixed; width: 1000px"> <colgroup> | ||
| 246 | <col style="width: 250px"> | ||
| 247 | <col style="width: 100px"> | ||
| 248 | <col style="width: 650px"> | ||
| 249 | </colgroup> | ||
| 250 | <thead> | ||
| 251 | <tr> | ||
| 252 | <th>返回值</th> | ||
| 253 | <th>错误码</th> | ||
| 254 | <th>描述</th> | ||
| 255 | </tr></thead> | ||
| 256 | <tbody> | ||
| 257 | <tr> | ||
| 258 | <td>ACLNN_ERR_PARAM_NULLPTR</td> | ||
| 259 | <td>161001</td> | ||
| 260 | <td>传入的paramsRef、gradsRef、momentumBufferListOptionalRef、grad_scale是空指针。</td> | ||
| 261 | </tr> | ||
| 262 | <tr> | ||
| 263 | <td>ACLNN_ERR_PARAM_INVALID</td> | ||
| 264 | <td>161002</td> | ||
| 265 | <td><ul><li>paramsRef、gradsRef、momentumBufferListOptionalRef、grad_scale的数据类型不在支持的范围之内。</li> | ||
| 266 | <li>paramsRef、gradsRef与momentumBufferListOptionalRef及其中各个tensor的数据类型不同。</li> | ||
| 267 | <li>paramsRef、gradsRef与momentumBufferListOptionalRef中tensor的shape维度小于等于8，grad_scale的shape不为[1]。</li> | ||
| 268 | <li>paramsRef、gradsRef与momentumBufferListOptionalRef（momentumBufferListOptionalRef不为空时）中相同索引tensor的shape不同。</li> | ||
| 269 | <li>paramsRef、gradsRef与momentumBufferListOptionalRef（momentumBufferListOptionalRef不为空时）中有空指针。</li> | ||
| 270 | <li>lr、momentum、weightDecay、dampening的值小于0。</li> | ||
| 271 | <li>gradScale等于0。</li> | ||
| 272 | </ul></td> | ||
| 273 | </tr> | ||
| 274 | </tbody> | ||
| 275 | </table> | ||
| 276 | </div> | ||
| 277 | |||
| 278 | ## aclnnFusedSgd | ||
| 279 | |||
| 280 | * **参数说明**： | ||
| 281 | |||
| 282 | <div style="overflow-x: auto;"> | ||
| 283 | <table style="undefined;table-layout: fixed; width: 900px"><colgroup> | ||
| 284 | <col style="width: 150px"> | ||
| 285 | <col style="width: 100px"> | ||
| 286 | <col style="width: 650px"> | ||
| 287 | </colgroup> | ||
| 288 | <thead> | ||
| 289 | <tr> | ||
| 290 | <th>参数名</th> | ||
| 291 | <th>输入/输出</th> | ||
| 292 | <th>描述</th> | ||
| 293 | </tr></thead> | ||
| 294 | <tbody> | ||
| 295 | <tr> | ||
| 296 | <td>workspace</td> | ||
| 297 | <td>输入</td> | ||
| 298 | <td>在Device侧申请的workspace内存地址。</td> | ||
| 299 | </tr> | ||
| 300 | <tr> | ||
| 301 | <td>workspaceSize</td> | ||
| 302 | <td>输入</td> | ||
| 303 | <td>在Device侧申请的workspace大小，由第一段接口aclnnFusedSgd获取。</td> | ||
| 304 | </tr> | ||
| 305 | <tr> | ||
| 306 | <td>executor</td> | ||
| 307 | <td>输入</td> | ||
| 308 | <td>op执行器，包含了算子计算流程。</td> | ||
| 309 | </tr> | ||
| 310 | <tr> | ||
| 311 | <td>stream</td> | ||
| 312 | <td>输入</td> | ||
| 313 | <td>指定执行任务的Stream。</td> | ||
| 314 | </tr> | ||
| 315 | </tbody></table> | ||
| 316 | </div> | ||
| 317 | |||
| 318 | * **返回值**： | ||
| 319 | 返回aclnnStatus状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn返回码.md)。 | ||
| 320 | |||
| 321 | ## 约束说明 | ||
| 322 | |||
| 323 | 1. paramsRef、gradsRef、momentumBufferListOptionalRef、grad_scale的数据类型在支持的范围之内。 | ||
| 324 | 2. paramsRef、gradsRef与momentumBufferListOptionalRef及其中各个tensor具有相同的数据类型。 | ||
| 325 | 3. paramsRef、gradsRef与momentumBufferListOptionalRef中tensor的shape维度小于等于8，grad_scale的shape为[1]。 | ||
| 326 | 4. paramsRef、gradsRef与momentumBufferListOptionalRef（momentumBufferListOptionalRef不为空时）中相同索引tensor的shape相同。 | ||
| 327 | 5. paramsRef、gradsRef与momentumBufferListOptionalRef（momentumBufferListOptionalRef不为空时）中不能有空指针。 | ||
| 328 | 6. lr、momentum、weightDecay、dampening的值大于等于0。 | ||
| 329 | 7. gradScale不等于0。 | ||
| 330 | |||
| 331 | ## 调用示例 | ||
| 332 | |||
| 333 | 示例代码如下，仅供参考，具体编译和执行过程请参考[编译与运行样例](../../../docs/zh/context/编译与运行样例.md)。 | ||
| 334 | |||
| 335 | ```Cpp | ||
| 336 | #include "acl/acl.h" | ||
| 337 | #include "aclnnop/aclnn_fused_sgd.h" | ||
| 338 | #include <iostream> | ||
| 339 | #include <vector> | ||
| 340 | |||
| 341 | #define CHECK_RET(cond, return_expr) \ | ||
| 342 | do { \ | ||
| 343 | if (!(cond)) { \ | ||
| 344 | return_expr; \ | ||
| 345 | } \ | ||
| 346 | } while (0) | ||
| 347 | |||
| 348 | #define LOG_PRINT(message, ...) \ | ||
| 349 | do { \ | ||
| 350 | printf(message, ##__VA_ARGS__); \ | ||
| 351 | } while (0) | ||
| 352 | |||
| 353 | int64_t GetShapeSize(const std::vector<int64_t> &shape) { | ||
| 354 | int64_t shapeSize = 1; | ||
| 355 | for (auto i : shape) { | ||
| 356 | shapeSize *= i; | ||
| 357 | } | ||
| 358 | return shapeSize; | ||
| 359 | } | ||
| 360 | |||
| 361 | void PrintOutResult(std::vector<int64_t> &shape, void **deviceAddr) { | ||
| 362 | auto size = GetShapeSize(shape); | ||
| 363 | std::vector<float> resultData(size, 0); | ||
| 364 | auto ret = aclrtMemcpy( | ||
| 365 | resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, | ||
| 366 | size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 367 | CHECK_RET( | ||
| 368 | ret == ACL_SUCCESS, | ||
| 369 | LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); | ||
| 370 | return ); | ||
| 371 | for (int64_t i = 0; i < size; i++) { | ||
| 372 | LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]); | ||
| 373 | } | ||
| 374 | } | ||
| 375 | |||
| 376 | int Init(int32_t deviceId, aclrtStream *stream) { | ||
| 377 | auto ret = aclInit(nullptr); | ||
| 378 | CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); | ||
| 379 | return ret); | ||
| 380 | ret = aclrtSetDevice(deviceId); | ||
| 381 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 382 | LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); | ||
| 383 | return ret); | ||
| 384 | ret = aclrtCreateStream(stream); | ||
| 385 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 386 | LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); | ||
| 387 | return ret); | ||
| 388 | return 0; | ||
| 389 | } | ||
| 390 | |||
| 391 | template <typename T> | ||
| 392 | int CreateAclTensor(const std::vector<T> &hostData, | ||
| 393 | const std::vector<int64_t> &shape, void **deviceAddr, | ||
| 394 | aclDataType dataType, aclTensor **tensor) { | ||
| 395 | auto size = GetShapeSize(shape) * sizeof(T); | ||
| 396 | auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 397 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 398 | LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); | ||
| 399 | return ret); | ||
| 400 | ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, | ||
| 401 | ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 402 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 403 | LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); | ||
| 404 | return ret); | ||
| 405 | |||
| 406 | std::vector<int64_t> strides(shape.size(), 1); | ||
| 407 | for (int64_t i = shape.size() - 2; i >= 0; i--) { | ||
| 408 | strides[i] = shape[i + 1] * strides[i + 1]; | ||
| 409 | } | ||
| 410 | |||
| 411 | *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, | ||
| 412 | strides.data(), 0, aclFormat::ACL_FORMAT_ND, | ||
| 413 | shape.data(), shape.size(), *deviceAddr); | ||
| 414 | return 0; | ||
| 415 | } | ||
| 416 | |||
| 417 | int main() { | ||
| 418 | int32_t deviceId = 0; | ||
| 419 | aclrtStream stream; | ||
| 420 | auto ret = Init(deviceId, &stream); | ||
| 421 | CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); | ||
| 422 | return ret); | ||
| 423 | |||
| 424 | std::vector<float> paramsRefHostData1 = {1, 2, 3, 4, 5, 6, 7, 8}; | ||
| 425 | std::vector<float> gradsRefHostData1 = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}; | ||
| 426 | std::vector<float> momentumHostData1 = {0, 0, 0, 0, 0, 0, 0, 0}; | ||
| 427 | std::vector<float> paramsRefHostData2 = {9, 10, 11, 12}; | ||
| 428 | std::vector<float> gradsRefHostData2 = {0.9, 1.0, 1.1, 1.2}; | ||
| 429 | std::vector<float> momentumHostData2 = {0, 0, 0, 0}; | ||
| 430 | std::vector<float> gradScaleOptionalHostData = {1.0}; | ||
| 431 | std::vector<int64_t> inputShape1 = {2, 2, 2}; | ||
| 432 | std::vector<int64_t> inputShape2 = {2, 2}; | ||
| 433 | std::vector<int64_t> scalarShape = {1}; | ||
| 434 | |||
| 435 | void *paramsRef1DeviceAddr = nullptr; | ||
| 436 | void *gradsRef1DeviceAddr = nullptr; | ||
| 437 | void *momentum1DeviceAddr = nullptr; | ||
| 438 | void *paramsRef2DeviceAddr = nullptr; | ||
| 439 | void *gradsRef2DeviceAddr = nullptr; | ||
| 440 | void *momentum2DeviceAddr = nullptr; | ||
| 441 | void *gradScaleOptionalDeviceAddr = nullptr; | ||
| 442 | |||
| 443 | aclTensor *paramsRef1 = nullptr; | ||
| 444 | aclTensor *gradsRef1 = nullptr; | ||
| 445 | aclTensor *momentum1 = nullptr; | ||
| 446 | aclTensor *paramsRef2 = nullptr; | ||
| 447 | aclTensor *gradsRef2 = nullptr; | ||
| 448 | aclTensor *momentum2 = nullptr; | ||
| 449 | aclTensor *gradScaleOptional = nullptr; | ||
| 450 | |||
| 451 | ret = CreateAclTensor(paramsRefHostData1, inputShape1, ¶msRef1DeviceAddr, aclDataType::ACL_FLOAT, ¶msRef1); | ||
| 452 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 453 | ret = CreateAclTensor(gradsRefHostData1, inputShape1, &gradsRef1DeviceAddr, aclDataType::ACL_FLOAT, &gradsRef1); | ||
| 454 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 455 | ret = CreateAclTensor(momentumHostData1, inputShape1, &momentum1DeviceAddr, aclDataType::ACL_FLOAT, &momentum1); | ||
| 456 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 457 | ret = CreateAclTensor(paramsRefHostData2, inputShape2, ¶msRef2DeviceAddr, aclDataType::ACL_FLOAT, ¶msRef2); | ||
| 458 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 459 | ret = CreateAclTensor(gradsRefHostData2, inputShape2, &gradsRef2DeviceAddr, aclDataType::ACL_FLOAT, &gradsRef2); | ||
| 460 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 461 | ret = CreateAclTensor(momentumHostData2, inputShape2, &momentum2DeviceAddr, aclDataType::ACL_FLOAT, &momentum2); | ||
| 462 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 463 | ret = CreateAclTensor(gradScaleOptionalHostData, scalarShape, &gradScaleOptionalDeviceAddr, aclDataType::ACL_FLOAT, &gradScaleOptional); | ||
| 464 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 465 | |||
| 466 | std::vector<aclTensor*> paramsRefListData = {paramsRef1, paramsRef2}; | ||
| 467 | std::vector<aclTensor*> gradsRefListData = {gradsRef1, gradsRef2}; | ||
| 468 | std::vector<aclTensor*> momentumListData = {momentum1, momentum2}; | ||
| 469 | aclTensorList* paramsRefList = aclCreateTensorList(paramsRefListData.data(), paramsRefListData.size()); | ||
| 470 | aclTensorList* gradsRefList = aclCreateTensorList(gradsRefListData.data(), gradsRefListData.size()); | ||
| 471 | aclTensorList* momentumList = aclCreateTensorList(momentumListData.data(), momentumListData.size()); | ||
| 472 | |||
| 473 | float weightDecay = 0.01f; | ||
| 474 | float momentumVal = 0.9f; | ||
| 475 | float lr = 0.001f; | ||
| 476 | float dampening = 0.0f; | ||
| 477 | bool nesterov = false; | ||
| 478 | bool maximize = false; | ||
| 479 | bool isFirstStep = true; | ||
| 480 | |||
| 481 | uint64_t workspaceSize = 0; | ||
| 482 | aclOpExecutor *executor; | ||
| 483 | |||
| 484 | ret = aclnnFusedSgdGetWorkspaceSize(paramsRefList, gradsRefList, momentumList, gradScaleOptional, | ||
| 485 | weightDecay, momentumVal, lr, dampening, | ||
| 486 | nesterov, maximize, isFirstStep, | ||
| 487 | &workspaceSize, &executor); | ||
| 488 | CHECK_RET( | ||
| 489 | ret == ACL_SUCCESS, | ||
| 490 | LOG_PRINT("aclnnFusedSgdGetWorkspaceSize failed. ERROR: %d\n", ret); | ||
| 491 | return ret); | ||
| 492 | |||
| 493 | void *workspaceAddr = nullptr; | ||
| 494 | if (workspaceSize > 0) { | ||
| 495 | ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 496 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 497 | LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); | ||
| 498 | return ret); | ||
| 499 | } | ||
| 500 | |||
| 501 | ret = aclnnFusedSgd(workspaceAddr, workspaceSize, executor, stream); | ||
| 502 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 503 | LOG_PRINT("aclnnFusedSgd failed. ERROR: %d\n", ret); | ||
| 504 | return ret); | ||
| 505 | |||
| 506 | ret = aclrtSynchronizeStream(stream); | ||
| 507 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 508 | LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); | ||
| 509 | return ret); | ||
| 510 | |||
| 511 | LOG_PRINT("====== Tensor 1 paramsRef results ======\n"); | ||
| 512 | PrintOutResult(inputShape1, ¶msRef1DeviceAddr); | ||
| 513 | LOG_PRINT("====== Tensor 1 gradsRef results ======\n"); | ||
| 514 | PrintOutResult(inputShape1, &gradsRef1DeviceAddr); | ||
| 515 | LOG_PRINT("------ Momentum buffer 1 ------\n"); | ||
| 516 | PrintOutResult(inputShape1, &momentum1DeviceAddr); | ||
| 517 | LOG_PRINT("====== Tensor 2 paramsRef results ======\n"); | ||
| 518 | PrintOutResult(inputShape2, ¶msRef2DeviceAddr); | ||
| 519 | LOG_PRINT("====== Tensor 2 gradsRef results ======\n"); | ||
| 520 | PrintOutResult(inputShape2, &gradsRef2DeviceAddr); | ||
| 521 | LOG_PRINT("------ Momentum buffer 2 ------\n"); | ||
| 522 | PrintOutResult(inputShape2, &momentum2DeviceAddr); | ||
| 523 | |||
| 524 | aclDestroyTensorList(paramsRefList); | ||
| 525 | aclDestroyTensorList(gradsRefList); | ||
| 526 | aclDestroyTensorList(momentumList); | ||
| 527 | aclDestroyTensor(gradScaleOptional); | ||
| 528 | |||
| 529 | aclrtFree(paramsRef1DeviceAddr); | ||
| 530 | aclrtFree(gradsRef1DeviceAddr); | ||
| 531 | aclrtFree(momentum1DeviceAddr); | ||
| 532 | aclrtFree(paramsRef2DeviceAddr); | ||
| 533 | aclrtFree(gradsRef2DeviceAddr); | ||
| 534 | aclrtFree(momentum2DeviceAddr); | ||
| 535 | aclrtFree(gradScaleOptionalDeviceAddr); | ||
| 536 | if (workspaceSize > 0) { | ||
| 537 | aclrtFree(workspaceAddr); | ||
| 538 | } | ||
| 539 | |||
| 540 | ret = aclrtDestroyStream(stream); | ||
| 541 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 542 | LOG_PRINT("destroy stream failed. ERROR: %d\n", ret); | ||
| 543 | return ret); | ||
| 544 | ret = aclrtResetDevice(deviceId); | ||
| 545 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 546 | LOG_PRINT("reset device failed. ERROR: %d\n", ret); | ||
| 547 | return ret); | ||
| 548 | ret = aclFinalize(); | ||
| 549 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 550 | LOG_PRINT("finalize acl failed. ERROR: %d\n", ret); | ||
| 551 | return ret); | ||
| 552 | return 0; | ||
| 553 | } | ||
| 554 | |||
| 555 | ``` | ||
| @@ -0,0 +1,228 @@ | |||
| 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 | |||
| 13 | |||
| 14 | |||
| 15 | |||
| 16 | |||
| 17 | do { \ | ||
| 18 | if (!(cond)) { \ | ||
| 19 | return_expr; \ | ||
| 20 | } \ | ||
| 21 | } while (0) | ||
| 22 | |||
| 23 | |||
| 24 | do { \ | ||
| 25 | printf(message, ##__VA_ARGS__); \ | ||
| 26 | } while (0) | ||
| 27 | |||
| 28 | int64_t GetShapeSize(const std::vector<int64_t> &shape) { | ||
| 29 | int64_t shapeSize = 1; | ||
| 30 | for (auto i : shape) { | ||
| 31 | shapeSize *= i; | ||
| 32 | } | ||
| 33 | return shapeSize; | ||
| 34 | } | ||
| 35 | |||
| 36 | void PrintOutResult(std::vector<int64_t> &shape, void **deviceAddr) { | ||
| 37 | auto size = GetShapeSize(shape); | ||
| 38 | std::vector<float> resultData(size, 0); | ||
| 39 | auto ret = aclrtMemcpy( | ||
| 40 | resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr, | ||
| 41 | size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST); | ||
| 42 | CHECK_RET( | ||
| 43 | ret == ACL_SUCCESS, | ||
| 44 | LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); | ||
| 45 | return ); | ||
| 46 | for (int64_t i = 0; i < size; i++) { | ||
| 47 | LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]); | ||
| 48 | } | ||
| 49 | } | ||
| 50 | |||
| 51 | int Init(int32_t deviceId, aclrtStream *stream) { | ||
| 52 | auto ret = aclInit(nullptr); | ||
| 53 | CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); | ||
| 54 | return ret); | ||
| 55 | ret = aclrtSetDevice(deviceId); | ||
| 56 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 57 | LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); | ||
| 58 | return ret); | ||
| 59 | ret = aclrtCreateStream(stream); | ||
| 60 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 61 | LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); | ||
| 62 | return ret); | ||
| 63 | return 0; | ||
| 64 | } | ||
| 65 | |||
| 66 | template <typename T> | ||
| 67 | int CreateAclTensor(const std::vector<T> &hostData, | ||
| 68 | const std::vector<int64_t> &shape, void **deviceAddr, | ||
| 69 | aclDataType dataType, aclTensor **tensor) { | ||
| 70 | auto size = GetShapeSize(shape) * sizeof(T); | ||
| 71 | auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 72 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 73 | LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); | ||
| 74 | return ret); | ||
| 75 | ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, | ||
| 76 | ACL_MEMCPY_HOST_TO_DEVICE); | ||
| 77 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 78 | LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); | ||
| 79 | return ret); | ||
| 80 | |||
| 81 | std::vector<int64_t> strides(shape.size(), 1); | ||
| 82 | for (int64_t i = shape.size() - 2; i >= 0; i--) { | ||
| 83 | strides[i] = shape[i + 1] * strides[i + 1]; | ||
| 84 | } | ||
| 85 | |||
| 86 | *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, | ||
| 87 | strides.data(), 0, aclFormat::ACL_FORMAT_ND, | ||
| 88 | shape.data(), shape.size(), *deviceAddr); | ||
| 89 | return 0; | ||
| 90 | } | ||
| 91 | |||
| 92 | int main() { | ||
| 93 | int32_t deviceId = 0; | ||
| 94 | aclrtStream stream; | ||
| 95 | auto ret = Init(deviceId, &stream); | ||
| 96 | CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); | ||
| 97 | return ret); | ||
| 98 | |||
| 99 | std::vector<float> paramsRefHostData1 = {1, 2, 3, 4, 5, 6, 7, 8}; | ||
| 100 | std::vector<float> gradsRefHostData1 = {0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8}; | ||
| 101 | std::vector<float> momentumHostData1 = {0, 0, 0, 0, 0, 0, 0, 0}; | ||
| 102 | std::vector<float> paramsRefHostData2 = {9, 10, 11, 12}; | ||
| 103 | std::vector<float> gradsRefHostData2 = {0.9, 1.0, 1.1, 1.2}; | ||
| 104 | std::vector<float> momentumHostData2 = {0, 0, 0, 0}; | ||
| 105 | std::vector<float> gradScaleOptionalHostData = {1.0}; | ||
| 106 | std::vector<int64_t> inputShape1 = {2, 2, 2}; | ||
| 107 | std::vector<int64_t> inputShape2 = {2, 2}; | ||
| 108 | std::vector<int64_t> scalarShape = {1}; | ||
| 109 | |||
| 110 | void *paramsRef1DeviceAddr = nullptr; | ||
| 111 | void *gradsRef1DeviceAddr = nullptr; | ||
| 112 | void *momentum1DeviceAddr = nullptr; | ||
| 113 | void *paramsRef2DeviceAddr = nullptr; | ||
| 114 | void *gradsRef2DeviceAddr = nullptr; | ||
| 115 | void *momentum2DeviceAddr = nullptr; | ||
| 116 | void *gradScaleOptionalDeviceAddr = nullptr; | ||
| 117 | |||
| 118 | aclTensor *paramsRef1 = nullptr; | ||
| 119 | aclTensor *gradsRef1 = nullptr; | ||
| 120 | aclTensor *momentum1 = nullptr; | ||
| 121 | aclTensor *paramsRef2 = nullptr; | ||
| 122 | aclTensor *gradsRef2 = nullptr; | ||
| 123 | aclTensor *momentum2 = nullptr; | ||
| 124 | aclTensor *gradScaleOptional = nullptr; | ||
| 125 | |||
| 126 | ret = CreateAclTensor(paramsRefHostData1, inputShape1, ¶msRef1DeviceAddr, aclDataType::ACL_FLOAT, ¶msRef1); | ||
| 127 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 128 | ret = CreateAclTensor(gradsRefHostData1, inputShape1, &gradsRef1DeviceAddr, aclDataType::ACL_FLOAT, &gradsRef1); | ||
| 129 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 130 | ret = CreateAclTensor(momentumHostData1, inputShape1, &momentum1DeviceAddr, aclDataType::ACL_FLOAT, &momentum1); | ||
| 131 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 132 | ret = CreateAclTensor(paramsRefHostData2, inputShape2, ¶msRef2DeviceAddr, aclDataType::ACL_FLOAT, ¶msRef2); | ||
| 133 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 134 | ret = CreateAclTensor(gradsRefHostData2, inputShape2, &gradsRef2DeviceAddr, aclDataType::ACL_FLOAT, &gradsRef2); | ||
| 135 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 136 | ret = CreateAclTensor(momentumHostData2, inputShape2, &momentum2DeviceAddr, aclDataType::ACL_FLOAT, &momentum2); | ||
| 137 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 138 | ret = CreateAclTensor(gradScaleOptionalHostData, scalarShape, &gradScaleOptionalDeviceAddr, aclDataType::ACL_FLOAT, &gradScaleOptional); | ||
| 139 | CHECK_RET(ret == ACL_SUCCESS, return ret); | ||
| 140 | |||
| 141 | std::vector<aclTensor*> paramsRefListData = {paramsRef1, paramsRef2}; | ||
| 142 | std::vector<aclTensor*> gradsRefListData = {gradsRef1, gradsRef2}; | ||
| 143 | std::vector<aclTensor*> momentumListData = {momentum1, momentum2}; | ||
| 144 | aclTensorList* paramsRefList = aclCreateTensorList(paramsRefListData.data(), paramsRefListData.size()); | ||
| 145 | aclTensorList* gradsRefList = aclCreateTensorList(gradsRefListData.data(), gradsRefListData.size()); | ||
| 146 | aclTensorList* momentumList = aclCreateTensorList(momentumListData.data(), momentumListData.size()); | ||
| 147 | |||
| 148 | float weightDecay = 0.01f; | ||
| 149 | float momentumVal = 0.9f; | ||
| 150 | float lr = 0.001f; | ||
| 151 | float dampening = 0.0f; | ||
| 152 | bool nesterov = false; | ||
| 153 | bool maximize = false; | ||
| 154 | bool isFirstStep = true; | ||
| 155 | |||
| 156 | uint64_t workspaceSize = 0; | ||
| 157 | aclOpExecutor *executor; | ||
| 158 | |||
| 159 | ret = aclnnFusedSgdGetWorkspaceSize(paramsRefList, gradsRefList, momentumList, gradScaleOptional, | ||
| 160 | weightDecay, momentumVal, lr, dampening, | ||
| 161 | nesterov, maximize, isFirstStep, | ||
| 162 | &workspaceSize, &executor); | ||
| 163 | CHECK_RET( | ||
| 164 | ret == ACL_SUCCESS, | ||
| 165 | LOG_PRINT("aclnnFusedSgdGetWorkspaceSize failed. ERROR: %d\n", ret); | ||
| 166 | return ret); | ||
| 167 | |||
| 168 | void *workspaceAddr = nullptr; | ||
| 169 | if (workspaceSize > 0) { | ||
| 170 | ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST); | ||
| 171 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 172 | LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); | ||
| 173 | return ret); | ||
| 174 | } | ||
| 175 | |||
| 176 | ret = aclnnFusedSgd(workspaceAddr, workspaceSize, executor, stream); | ||
| 177 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 178 | LOG_PRINT("aclnnFusedSgd failed. ERROR: %d\n", ret); | ||
| 179 | return ret); | ||
| 180 | |||
| 181 | ret = aclrtSynchronizeStream(stream); | ||
| 182 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 183 | LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); | ||
| 184 | return ret); | ||
| 185 | |||
| 186 | LOG_PRINT("====== Tensor 1 paramsRef results ======\n"); | ||
| 187 | PrintOutResult(inputShape1, ¶msRef1DeviceAddr); | ||
| 188 | LOG_PRINT("====== Tensor 1 gradsRef results ======\n"); | ||
| 189 | PrintOutResult(inputShape1, &gradsRef1DeviceAddr); | ||
| 190 | LOG_PRINT("------ Momentum buffer 1 ------\n"); | ||
| 191 | PrintOutResult(inputShape1, &momentum1DeviceAddr); | ||
| 192 | LOG_PRINT("====== Tensor 2 paramsRef results ======\n"); | ||
| 193 | PrintOutResult(inputShape2, ¶msRef2DeviceAddr); | ||
| 194 | LOG_PRINT("====== Tensor 2 gradsRef results ======\n"); | ||
| 195 | PrintOutResult(inputShape2, &gradsRef2DeviceAddr); | ||
| 196 | LOG_PRINT("------ Momentum buffer 2 ------\n"); | ||
| 197 | PrintOutResult(inputShape2, &momentum2DeviceAddr); | ||
| 198 | |||
| 199 | aclDestroyTensorList(paramsRefList); | ||
| 200 | aclDestroyTensorList(gradsRefList); | ||
| 201 | aclDestroyTensorList(momentumList); | ||
| 202 | aclDestroyTensor(gradScaleOptional); | ||
| 203 | |||
| 204 | aclrtFree(paramsRef1DeviceAddr); | ||
| 205 | aclrtFree(gradsRef1DeviceAddr); | ||
| 206 | aclrtFree(momentum1DeviceAddr); | ||
| 207 | aclrtFree(paramsRef2DeviceAddr); | ||
| 208 | aclrtFree(gradsRef2DeviceAddr); | ||
| 209 | aclrtFree(momentum2DeviceAddr); | ||
| 210 | aclrtFree(gradScaleOptionalDeviceAddr); | ||
| 211 | if (workspaceSize > 0) { | ||
| 212 | aclrtFree(workspaceAddr); | ||
| 213 | } | ||
| 214 | |||
| 215 | ret = aclrtDestroyStream(stream); | ||
| 216 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 217 | LOG_PRINT("destroy stream failed. ERROR: %d\n", ret); | ||
| 218 | return ret); | ||
| 219 | ret = aclrtResetDevice(deviceId); | ||
| 220 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 221 | LOG_PRINT("reset device failed. ERROR: %d\n", ret); | ||
| 222 | return ret); | ||
| 223 | ret = aclFinalize(); | ||
| 224 | CHECK_RET(ret == ACL_SUCCESS, | ||
| 225 | LOG_PRINT("finalize acl failed. ERROR: %d\n", ret); | ||
| 226 | return ret); | ||
| 227 | return 0; | ||
| 228 | } | ||
| @@ -0,0 +1,354 @@ | |||
| 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 | |||
| 13 | |||
| 14 | |||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | |||
| 23 | |||
| 24 | |||
| 25 | |||
| 26 | |||
| 27 | |||
| 28 | using namespace op; | ||
| 29 | |||
| 30 | extern "C" { | ||
| 31 | |||
| 32 | |||
| 33 | static const std::initializer_list<op::DataType> ASCEND910_DTYPE_SUPPORT_LIST = { | ||
| 34 | op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16}; | ||
| 35 | |||
| 36 | static const std::initializer_list<op::DataType> ASCEND910B_DTYPE_SUPPORT_LIST = { | ||
| 37 | op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16}; | ||
| 38 | |||
| 39 | static const std::initializer_list<op::DataType> ASCEND950_DTYPE_SUPPORT_LIST = { | ||
| 40 | op::DataType::DT_FLOAT, op::DataType::DT_FLOAT16, op::DataType::DT_BF16}; | ||
| 41 | |||
| 42 | static inline const std::initializer_list<op::DataType>& GetDtypeSupportListFromSocVersion() | ||
| 43 | { | ||
| 44 | auto curArch = GetCurrentPlatformInfo().GetCurNpuArch(); | ||
| 45 | switch (curArch) { | ||
| 46 | case NpuArch::DAV_2201: { | ||
| 47 | return ASCEND910B_DTYPE_SUPPORT_LIST; | ||
| 48 | } | ||
| 49 | case NpuArch::DAV_3510: { | ||
| 50 | return ASCEND950_DTYPE_SUPPORT_LIST; | ||
| 51 | } | ||
| 52 | case NpuArch::DAV_1001: { | ||
| 53 | return ASCEND910_DTYPE_SUPPORT_LIST; | ||
| 54 | } | ||
| 55 | default: { | ||
| 56 | return ASCEND910_DTYPE_SUPPORT_LIST; | ||
| 57 | } | ||
| 58 | } | ||
| 59 | } | ||
| 60 | |||
| 61 | static bool CheckNotNull( | ||
| 62 | const aclTensorList *paramsRef, | ||
| 63 | const aclTensorList *gradsRef, | ||
| 64 | const aclTensorList *momentumBufferListOptionalRef) | ||
| 65 | { | ||
| 66 | OP_CHECK_NULL(paramsRef, return false); | ||
| 67 | for (uint64_t i = 0; i < paramsRef->Size(); i++) { | ||
| 68 | OP_CHECK_NULL((*paramsRef)[i], return false); | ||
| 69 | } | ||
| 70 | OP_CHECK_NULL(gradsRef, return false); | ||
| 71 | for (uint64_t i = 0; i < gradsRef->Size(); i++) { | ||
| 72 | OP_CHECK_NULL((*gradsRef)[i], return false); | ||
| 73 | } | ||
| 74 | if(momentumBufferListOptionalRef != nullptr) { | ||
| 75 | for (uint64_t i = 0; i < momentumBufferListOptionalRef->Size(); i++) { | ||
| 76 | OP_CHECK_NULL((*momentumBufferListOptionalRef)[i], return false); | ||
| 77 | } | ||
| 78 | } | ||
| 79 | return true; | ||
| 80 | } | ||
| 81 | |||
| 82 | static bool CheckTensorListCount( | ||
| 83 | const aclTensorList *paramsRef, | ||
| 84 | const aclTensorList *gradsRef, | ||
| 85 | const aclTensorList *momentumBufferListOptionalRef) | ||
| 86 | { | ||
| 87 | auto tensorCount = paramsRef->Size(); | ||
| 88 | if (tensorCount == 0) { | ||
| 89 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "paramsRef tensor list is empty."); | ||
| 90 | return false; | ||
| 91 | } | ||
| 92 | if (gradsRef->Size() != tensorCount) { | ||
| 93 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "gradsRef tensor count does not match paramsRef."); | ||
| 94 | return false; | ||
| 95 | } | ||
| 96 | if (momentumBufferListOptionalRef != nullptr && momentumBufferListOptionalRef->Size() != tensorCount) { | ||
| 97 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "momentumBufferListOptionalRef tensor count does not match paramsRef."); | ||
| 98 | return false; | ||
| 99 | } | ||
| 100 | return true; | ||
| 101 | } | ||
| 102 | |||
| 103 | static bool CheckDtype( | ||
| 104 | const aclTensorList *paramsRef, | ||
| 105 | const aclTensorList *gradsRef, | ||
| 106 | const aclTensorList *momentumBufferListOptionalRef) | ||
| 107 | { | ||
| 108 | const std::initializer_list<op::DataType> dtypeSupportList = GetDtypeSupportListFromSocVersion(); | ||
| 109 | auto paramsTensor = (*paramsRef)[0]; | ||
| 110 | |||
| 111 | OP_CHECK_DTYPE_NOT_SUPPORT(paramsTensor, dtypeSupportList, return false); | ||
| 112 | op::DataType inputType = paramsTensor->GetDataType(); | ||
| 113 | for (uint64_t i = 0; i < paramsRef->Size(); i++) { | ||
| 114 | if ((*paramsRef)[i]->GetDataType() != inputType || (*gradsRef)[i]->GetDataType() != inputType) { | ||
| 115 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "expects all input tensors with the same dtype."); | ||
| 116 | return false; | ||
| 117 | } | ||
| 118 | } | ||
| 119 | if(momentumBufferListOptionalRef != nullptr) { | ||
| 120 | auto momentumTensor = (*momentumBufferListOptionalRef)[0]; | ||
| 121 | OP_CHECK_DTYPE_NOT_SUPPORT(momentumTensor, dtypeSupportList, return false); | ||
| 122 | OP_CHECK_DTYPE_NOT_SAME(paramsTensor, momentumTensor, return false); | ||
| 123 | for (uint64_t i = 0; i < momentumBufferListOptionalRef->Size(); i++) { | ||
| 124 | if ((*momentumBufferListOptionalRef)[i]->GetDataType() != inputType) { | ||
| 125 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "expects all input tensors with the same dtype."); | ||
| 126 | return false; | ||
| 127 | } | ||
| 128 | } | ||
| 129 | } | ||
| 130 | return true; | ||
| 131 | } | ||
| 132 | |||
| 133 | static void CheckOptionalTensorListEmpty(const aclTensorList *&tensorList) | ||
| 134 | { | ||
| 135 | if (tensorList == nullptr) { | ||
| 136 | OP_LOGI("momentumBufferListOptionalRef is nullptr"); | ||
| 137 | return; | ||
| 138 | } | ||
| 139 | if (tensorList->Size() == 0) { | ||
| 140 | OP_LOGI("momentumBufferListOptionalRef is set nullptr because len(momentumBufferListOptionalRef) is 0."); | ||
| 141 | tensorList = nullptr; | ||
| 142 | } | ||
| 143 | } | ||
| 144 | |||
| 145 | static void CheckIsFirstStep(bool isFirstStep) | ||
| 146 | { | ||
| 147 | if (isFirstStep) { | ||
| 148 | OP_LOGW("isFirstStep argument has no effect when momentumBufferListOptionalRef is empty"); | ||
| 149 | } | ||
| 150 | } | ||
🔵 Low Priority
该函数在 变更：第169-174行新增的 建议：将   | |||
| 151 | |||
| 152 | static bool CheckAttr( | ||
| 153 | const aclTensorList *momentumBufferListOptionalRef, | ||
| 154 | float weightDecay, float momentum, float lr, | ||
| 155 | float dampening, bool nesterov) | ||
| 156 | { | ||
| 157 | if (weightDecay < 0) { | ||
| 158 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "weightDecay[%f] shoule be greater or equal than 0", weightDecay); | ||
| 159 | return false; | ||
| 160 | } | ||
| 161 | if (momentum < 0) { | ||
| 162 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "momentum[%f] shoule be greater or equal than 0", momentum); | ||
| 163 | return false; | ||
| 164 | } | ||
| 165 | if (lr < 0) { | ||
| 166 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "lr[%f] shoule be greater or equal than 0", lr); | ||
| 167 | return false; | ||
| 168 | } | ||
| 169 | const float EPS = 1e-6f; | ||
| 170 | if (nesterov && (momentum <= 0 || std::abs(dampening) >= EPS)) { | ||
| 171 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "nesterov[%d] momentum requires a momentum[%f] and zero dampening[%f].", static_cast<int32_t>(nesterov), momentum, dampening); | ||
| 172 | return false; | ||
| 173 | } | ||
| 174 | if ((momentumBufferListOptionalRef == nullptr && std::abs(momentum) >= EPS)) { | ||
| 175 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "momentum[%f] is invalid. momentum shoube be 0 when momentumBufferListOptionalRef is nullptr.", momentum); | ||
| 176 | return false; | ||
| 177 | } | ||
| 178 | if ((momentumBufferListOptionalRef != nullptr && (momentum < 0.0 || std::abs(momentum) < EPS))) { | ||
| 179 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "momentum[%f] is invalid. momentum shoube be greater than 0 when momentumBufferListOptionalRef is not nullptr.", momentum); | ||
| 180 | return false; | ||
🔵 Low Priority 在
此外第199行句子中 "requires a momentum" 语义冗余。这些拼写错误会影响日志可读性，增加排查问题的难度。 变更：第182、186、190、199、203行新增的错误日志字符串。 建议：将所有 "shoule" 修正为 "should"，将 "shoube" 修正为 "should be"。第199行建议改为 majiajian_hw 6月22日 评论： 6月22日 评论： | |||
| 181 | } | ||
| 182 | return true; | ||
| 183 | } | ||
| 184 | |||
| 185 | static bool CheckShape( | ||
| 186 | const aclTensorList *paramsRef, | ||
| 187 | const aclTensorList *gradsRef, | ||
| 188 | const aclTensorList *momentumBufferListOptionalRef, | ||
| 189 | const aclTensor *gradScaleOptional) | ||
| 190 | { | ||
| 191 | for (uint64_t i = 0; i < paramsRef->Size(); i++) { | ||
| 192 | op::Shape expectShape = (*paramsRef)[i]->GetViewShape(); | ||
| 193 | if((*gradsRef)[i]->GetViewShape() != expectShape || | ||
| 194 | (momentumBufferListOptionalRef != nullptr && | ||
| 195 | (*momentumBufferListOptionalRef)[i]->GetViewShape() != expectShape)) { | ||
| 196 | OP_LOGE(ACLNN_ERR_PARAM_INVALID, "expects all input tensors with the same shape."); | ||
| 197 | return false; | ||
| 198 | } | ||
| 199 | } | ||
| 200 | if (gradScaleOptional != nullptr && !(gradScaleOptional->IsEmpty())) { | ||
| 201 | op::Shape shape = gradScaleOptional->GetViewShape(); | ||
| 202 | if (shape.GetDimNum() != 1 || shape.GetDim(0) != 1) { | ||
| 203 | return false; | ||
| 204 | } | ||
| 205 | } | ||
| 206 | return true; | ||
| 207 | } | ||
| 208 | |||
| 209 | static aclnnStatus CheckParams( | ||
| 210 | const aclTensorList *paramsRef, | ||
| 211 | const aclTensorList *gradsRef, | ||
| 212 | const aclTensorList *momentumBufferListOptionalRef, | ||
| 213 | const aclTensor *gradScaleOptional, | ||
| 214 | float weightDecay, float momentum, float lr, float dampening, bool nesterov) | ||
| 215 | { | ||
| 216 | CHECK_RET(CheckAttr(momentumBufferListOptionalRef, weightDecay, momentum, lr, dampening, nesterov), ACLNN_ERR_PARAM_INVALID); | ||
| 217 | CHECK_RET(CheckNotNull(paramsRef, gradsRef, momentumBufferListOptionalRef), ACLNN_ERR_PARAM_NULLPTR); | ||
| 218 | CHECK_RET(CheckTensorListCount(paramsRef, gradsRef, momentumBufferListOptionalRef), ACLNN_ERR_PARAM_INVALID); | ||
| 219 | CHECK_RET(CheckDtype(paramsRef, gradsRef, momentumBufferListOptionalRef), ACLNN_ERR_PARAM_INVALID); | ||
| 220 | CHECK_RET(CheckShape(paramsRef, gradsRef, momentumBufferListOptionalRef, gradScaleOptional), ACLNN_ERR_PARAM_INVALID); | ||
| 221 | return ACLNN_SUCCESS; | ||
| 222 | } | ||
| 223 | |||
| 224 | const aclTensor* FlattenDims(const aclTensor* tensor, aclOpExecutor* executor) | ||
| 225 | { | ||
| 226 | op::Shape shapeTensor = tensor->GetViewShape(); | ||
| 227 | int64_t dimNum = shapeTensor.GetDimNum(); | ||
| 228 | |||
| 229 | op::Shape newShape; | ||
| 230 | int64_t catdimSize = 1; | ||
| 231 | for (int64_t i = 0; i < dimNum; i++) { | ||
| 232 | catdimSize *= shapeTensor.GetDim(i); | ||
| 233 | } | ||
| 234 | newShape.AppendDim(catdimSize); | ||
| 235 | auto reshapeTensor = executor->CreateView(tensor, tensor->GetViewShape(), tensor->GetViewOffset()); | ||
| 236 | reshapeTensor->SetViewShape(newShape); | ||
| 237 | reshapeTensor->SetOriginalShape(newShape); | ||
| 238 | reshapeTensor->SetStorageShape(newShape); | ||
| 239 | return reshapeTensor; | ||
| 240 | } | ||
| 241 | |||
| 242 | static const aclTensorList* MakeContiguousTensorList( | ||
| 243 | const aclTensorList* tensorList, aclOpExecutor* executor) | ||
| 244 | { | ||
| 245 | op::FVector<const aclTensor*> contiguousTensors; | ||
| 246 | for (uint64_t i = 0; i < tensorList->Size(); i++) { | ||
| 247 | if ((*tensorList)[i]->IsEmpty()) { | ||
| 248 | continue; | ||
| 249 | } | ||
| 250 | auto contiguous = l0op::Contiguous((*tensorList)[i], executor); | ||
| 251 | CHECK_RET(contiguous != nullptr, nullptr); | ||
| 252 | contiguous = FlattenDims(contiguous, executor); | ||
| 253 | contiguousTensors.emplace_back(contiguous); | ||
| 254 | } | ||
| 255 | return executor->AllocTensorList(contiguousTensors.data(), contiguousTensors.size()); | ||
| 256 | } | ||
| 257 | |||
| 258 | static void ViewCopyTensorList( | ||
| 259 | const aclTensorList* src, const aclTensorList* dst, aclOpExecutor* executor) | ||
| 260 | { | ||
| 261 | uint64_t cnt = 0; | ||
| 262 | for (uint64_t i = 0; i < dst->Size(); i++) { | ||
| 263 | if ((*dst)[i]->IsEmpty()) { | ||
| 264 | continue; | ||
| 265 | } | ||
| 266 | l0op::ViewCopy((*src)[cnt], (*dst)[i], executor); | ||
| 267 | cnt += 1; | ||
| 268 | } | ||
| 269 | } | ||
| 270 | |||
| 271 | aclnnStatus aclnnFusedSgdGetWorkspaceSize( | ||
| 272 | const aclTensorList *paramsRef, | ||
| 273 | const aclTensorList *gradsRef, | ||
| 274 | const aclTensorList *momentumBufferListOptionalRef, | ||
| 275 | const aclTensor *gradScaleOptional, | ||
| 276 | float weightDecay, | ||
| 277 | float momentum, | ||
| 278 | float lr, | ||
| 279 | float dampening, | ||
| 280 | bool nesterov, | ||
| 281 | bool maximize, | ||
| 282 | bool isFirstStep, | ||
| 283 | uint64_t *workspaceSize, | ||
| 284 | aclOpExecutor **executor) | ||
| 285 | { | ||
| 286 | L2_DFX_PHASE_1(aclnnFusedSgd, | ||
| 287 | DFX_IN(paramsRef, gradsRef, momentumBufferListOptionalRef, gradScaleOptional, | ||
| 288 | weightDecay, momentum, lr, dampening, nesterov, maximize, isFirstStep), | ||
| 289 | DFX_OUT(paramsRef, gradsRef, momentumBufferListOptionalRef)); | ||
| 290 | |||
| 291 | auto uniqueExecutor = CREATE_EXECUTOR(); | ||
| 292 | CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR); | ||
| 293 | |||
| 294 | if (paramsRef->Size() == 0) { | ||
| 295 | uniqueExecutor.ReleaseTo(executor); | ||
| 296 | return ACLNN_SUCCESS; | ||
| 297 | } | ||
| 298 | |||
| 299 | CheckOptionalTensorListEmpty(momentumBufferListOptionalRef); | ||
| 300 | |||
| 301 | CheckIsFirstStep(isFirstStep); | ||
| 302 | |||
| 303 | auto ret = CheckParams(paramsRef, gradsRef, momentumBufferListOptionalRef, gradScaleOptional, weightDecay, momentum, lr, dampening, nesterov); | ||
| 304 | CHECK_RET(ret == ACLNN_SUCCESS, ret); | ||
| 305 | |||
| 306 | if (gradScaleOptional != nullptr && !(gradScaleOptional->IsEmpty())) { | ||
| 307 | gradScaleOptional = l0op::Cast(gradScaleOptional, DataType::DT_FLOAT, uniqueExecutor.get()); | ||
| 308 | } | ||
| 309 | |||
| 310 | auto paramsContiguous = MakeContiguousTensorList(paramsRef, uniqueExecutor.get()); | ||
| 311 | CHECK_RET(paramsContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); | ||
| 312 | |||
| 313 | auto gradsContiguous = MakeContiguousTensorList(gradsRef, uniqueExecutor.get()); | ||
| 314 | CHECK_RET(gradsContiguous != nullptr, ACLNN_ERR_INNER_NULLPTR); | ||
| 315 | |||
| 316 | const aclTensorList* momentumContiguous = nullptr; | ||
| 317 | if(momentumBufferListOptionalRef != nullptr) { | ||
| 318 | momentumContiguous = MakeContiguousTensorList(momentumBufferListOptionalRef, uniqueExecutor.get()); | ||
| 319 | } | ||
| 320 | |||
| 321 | auto [paramsOut, gradsOut, momentumOut] = l0op::FusedSgd( | ||
| 322 | paramsContiguous, gradsContiguous, momentumContiguous, | ||
| 323 | gradScaleOptional, | ||
| 324 | weightDecay, momentum, lr, dampening, nesterov, maximize, isFirstStep, | ||
| 325 | uniqueExecutor.get()); | ||
| 326 | CHECK_RET(paramsOut != nullptr, ACLNN_ERR_INNER_NULLPTR); | ||
| 327 | CHECK_RET(gradsOut != nullptr, ACLNN_ERR_INNER_NULLPTR); | ||
| 328 | CHECK_RET(momentumContiguous == nullptr || momentumOut != nullptr, ACLNN_ERR_INNER_NULLPTR); | ||
| 329 | |||
| 330 | ViewCopyTensorList(paramsOut, paramsRef, uniqueExecutor.get()); | ||
| 331 | ViewCopyTensorList(gradsOut, gradsRef, uniqueExecutor.get()); | ||
| 332 | if(momentumBufferListOptionalRef != nullptr) { | ||
| 333 | ViewCopyTensorList(momentumOut, momentumBufferListOptionalRef, uniqueExecutor.get()); | ||
| 334 | } | ||
| 335 | |||
| 336 | *workspaceSize = uniqueExecutor->GetWorkspaceSize(); | ||
| 337 | uniqueExecutor.ReleaseTo(executor); | ||
| 338 | |||
| 339 | return ACLNN_SUCCESS; | ||
| 340 | } | ||
| 341 | |||
| 342 | aclnnStatus aclnnFusedSgd( | ||
| 343 | void *workspace, | ||
| 344 | uint64_t workspaceSize, | ||
| 345 | aclOpExecutor *executor, | ||
| 346 | aclrtStream stream) | ||
| 347 | { | ||
| 348 | L2_DFX_PHASE_2(aclnnFusedSgd); | ||
| 349 | return CommonOpExecutorRun(workspace, workspaceSize, executor, stream); | ||
| 350 | } | ||
| 351 | |||
| 352 | |||
| 353 | } | ||
| 354 | |||
| @@ -0,0 +1,79 @@ | |||
| 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 | |||
| 13 | |||
| 14 | |||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | extern "C" { | ||
| 19 | |||
| 20 | |||
| 21 | /** | ||
| 22 | * @brief aclnnFusedSgd的第一段接口，根据具体的计算流程，计算workspace大小。 | ||
| 23 | * @domain aclnn_ops_train | ||
| 24 | * | ||
| 25 | * 算子功能：融合SGD优化器，支持momentum、weight_decay、nesterov、maximize， | ||
| 26 | * 将参数更新、动量更新、梯度缩放等操作融合为单个kernel。 | ||
| 27 | * | ||
| 28 | * @param [in] paramsRef: device侧的aclTensorList，需要更新的参数列表。 | ||
| 29 | * 数据类型支持FLOAT、FLOAT16、BFLOAT16。数据格式支持ND。 | ||
| 30 | * @param [in] gradsRef: device侧的aclTensorList，梯度列表。数据类型、shape需要与params一致。 | ||
| 31 | * @param [in] momentumBufferListOptionalRef: device侧的aclTensorList，动量缓冲列表。数据类型、shape需要与params一致。 | ||
| 32 | * @param [in] gradScaleOptional: device侧的aclTensor（可选），梯度缩放因子。数据类型支持FLOAT。 | ||
| 33 | * @param [in] weightDecay: 权重衰减系数，数据类型FLOAT。 | ||
| 34 | * @param [in] momentum: 动量因子，数据类型FLOAT。 | ||
| 35 | * @param [in] lr: 学习率，数据类型FLOAT。 | ||
| 36 | * @param [in] dampening: 动量阻尼系数，数据类型FLOAT。 | ||
| 37 | * @param [in] nesterov: 是否启用Nesterov动量，数据类型BOOL。 | ||
| 38 | * @param [in] maximize: 是否最大化目标函数，数据类型BOOL。 | ||
| 39 | * @param [in] isFirstStep: 是否为第一个优化步，数据类型BOOL。 | ||
| 40 | * @param [out] workspaceSize: 返回用户在device侧申请的workspace大小。 | ||
| 41 | * @param [out] executor: 返回op执行器。 | ||
| 42 | * @return aclnnStatus: 返回状态码。 | ||
| 43 | */ | ||
| 44 | ACLNN_API aclnnStatus aclnnFusedSgdGetWorkspaceSize( | ||
| 45 | const aclTensorList *paramsRef, | ||
| 46 | const aclTensorList *gradsRef, | ||
| 47 | const aclTensorList *momentumBufferListOptionalRef, | ||
| 48 | const aclTensor *gradScaleOptional, | ||
| 49 | float weightDecay, | ||
| 50 | float momentum, | ||
| 51 | float lr, | ||
| 52 | float dampening, | ||
| 53 | bool nesterov, | ||
| 54 | bool maximize, | ||
| 55 | bool isFirstStep, | ||
| 56 | uint64_t *workspaceSize, | ||
| 57 | aclOpExecutor **executor); | ||
| 58 | |||
| 59 | /** | ||
| 60 | * @brief aclnnFusedSgd的第二段接口，用于执行计算。 | ||
| 61 | * | ||
| 62 | * 算子功能：执行融合SGD优化器。 | ||
| 63 | * @param [in] workspace: 在device侧申请的workspace内存起址。 | ||
| 64 | * @param [in] workspaceSize: workspace大小，由aclnnFusedSgdGetWorkspaceSize获取。 | ||
| 65 | * @param [in] executor: op执行器。 | ||
| 66 | * @param [in] stream: acl stream流。 | ||
| 67 | * @return aclnnStatus: 返回状态码。 | ||
| 68 | */ | ||
| 69 | ACLNN_API aclnnStatus aclnnFusedSgd( | ||
| 70 | void *workspace, | ||
| 71 | uint64_t workspaceSize, | ||
| 72 | aclOpExecutor *executor, | ||
| 73 | aclrtStream stream); | ||
| 74 | |||
| 75 | |||
| 76 | } | ||
| 77 | |||
| 78 | |||
| 79 | |||
| @@ -0,0 +1,93 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd.cpp | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | |||
| 23 | |||
| 24 | |||
| 25 | |||
| 26 | using namespace op; | ||
| 27 | |||
| 28 | namespace l0op { | ||
| 29 | OP_TYPE_REGISTER(FusedSgd); | ||
| 30 | |||
| 31 | std::tuple<const aclTensorList*, const aclTensorList*, const aclTensorList*> FusedSgd( | ||
| 32 | const aclTensorList* paramsRef, | ||
| 33 | const aclTensorList* gradsRef, | ||
| 34 | const aclTensorList* momentumBufferListOptionalRef, | ||
| 35 | const aclTensor* gradScaleOptional, | ||
| 36 | float weightDecay, | ||
| 37 | float momentum, | ||
| 38 | float lr, | ||
| 39 | float dampening, | ||
| 40 | bool nesterov, | ||
| 41 | bool maximize, | ||
| 42 | bool isFirstStep, | ||
| 43 | aclOpExecutor *executor) | ||
| 44 | { | ||
| 45 | L0_DFX(FusedSgd, paramsRef, gradsRef, momentumBufferListOptionalRef, gradScaleOptional, | ||
| 46 | weightDecay, momentum, lr, dampening, nesterov, maximize, isFirstStep); | ||
| 47 | |||
| 48 | const aclTensorList* momentumBufferListOptionalRefOut = nullptr; | ||
| 49 | if (momentumBufferListOptionalRef == nullptr) { | ||
| 50 | const op::Shape momentumBufferListOptionalRefOutShape = {1}; | ||
| 51 | const aclTensor* tmpTensor = executor->AllocTensor(momentumBufferListOptionalRefOutShape, (*paramsRef)[0]->GetDataType(), (*paramsRef)[0]->GetStorageFormat()); | ||
| 52 | op::FVector<const aclTensor*> tensorListA; | ||
| 53 | tensorListA.emplace_back(tmpTensor); | ||
| 54 | momentumBufferListOptionalRefOut = executor->AllocTensorList(tensorListA.data(), tensorListA.size()); | ||
| 55 | if(momentumBufferListOptionalRefOut == nullptr) { | ||
| 56 | return std::tuple<const aclTensorList*, const aclTensorList*, const aclTensorList*>( | ||
| 57 | nullptr, nullptr, nullptr); | ||
| 58 | } | ||
| 59 | } else { | ||
| 60 | momentumBufferListOptionalRefOut = momentumBufferListOptionalRef; | ||
| 61 | } | ||
| 62 | |||
| 63 | const aclTensorList* momentumBufferListOptionalRefInput = momentumBufferListOptionalRef; | ||
| 64 | if (momentumBufferListOptionalRef == nullptr) { | ||
| 65 | const op::Shape momentumBufferListOptionalRefInputShape = {0}; | ||
| 66 | const aclTensor* tmpTensor = executor->AllocTensor(momentumBufferListOptionalRefInputShape, (*paramsRef)[0]->GetDataType(), (*paramsRef)[0]->GetStorageFormat()); | ||
| 67 | op::FVector<const aclTensor*> tensorListB; | ||
| 68 | tensorListB.emplace_back(tmpTensor); | ||
| 69 | momentumBufferListOptionalRefInput = executor->AllocTensorList(tensorListB.data(), tensorListB.size()); | ||
| 70 | if(momentumBufferListOptionalRefInput == nullptr) { | ||
| 71 | return std::tuple<const aclTensorList*, const aclTensorList*, const aclTensorList*>( | ||
| 72 | nullptr, nullptr, nullptr); | ||
| 73 | } | ||
| 74 | } | ||
| 75 | |||
| 76 | auto retAicore = ADD_TO_LAUNCHER_LIST_AICORE(FusedSgd, | ||
| 77 | OP_INPUT(paramsRef, gradsRef, momentumBufferListOptionalRefInput, gradScaleOptional), | ||
| 78 | OP_OUTPUT(paramsRef, gradsRef, momentumBufferListOptionalRefOut), | ||
| 79 | OP_ATTR(weightDecay, momentum, lr, dampening, nesterov, maximize, isFirstStep)); | ||
| 80 | if (retAicore != ACLNN_SUCCESS) { | ||
| 81 | OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "FusedSgd ADD_TO_LAUNCHER_LIST_AICORE failed."); | ||
| 82 | return std::tuple<const aclTensorList*, const aclTensorList*, const aclTensorList*>(nullptr, nullptr, nullptr); | ||
| 83 | } | ||
| 84 | if (momentumBufferListOptionalRef == nullptr) { | ||
| 85 | return std::tuple<const aclTensorList*, const aclTensorList*, const aclTensorList*>( | ||
| 86 | paramsRef, gradsRef, nullptr); | ||
| 87 | } else { | ||
| 88 | return std::tuple<const aclTensorList*, const aclTensorList*, const aclTensorList*>( | ||
| 89 | paramsRef, gradsRef, momentumBufferListOptionalRefOut); | ||
| 90 | } | ||
| 91 | } | ||
| 92 | |||
| 93 | } // namespace l0op | ||
| @@ -0,0 +1,36 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd.h | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | namespace l0op { | ||
| 21 | std::tuple<const aclTensorList*, const aclTensorList*, const aclTensorList*> FusedSgd( | ||
| 22 | const aclTensorList* paramsRef, | ||
| 23 | const aclTensorList* gradsRef, | ||
| 24 | const aclTensorList* momentumBufferListOptionalRef, | ||
| 25 | const aclTensor* gradScaleOptional, | ||
| 26 | float weightDecay, | ||
| 27 | float momentum, | ||
| 28 | float lr, | ||
| 29 | float dampening, | ||
| 30 | bool nesterov, | ||
| 31 | bool maximize, | ||
| 32 | bool isFirstStep, | ||
| 33 | aclOpExecutor *executor); | ||
| 34 | } | ||
| 35 | |||
| 36 | |||
| @@ -0,0 +1,12 @@ | |||
| 1 | # ---------------------------------------------------------------------------- | ||
| 2 | # Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | # This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | # CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | # Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | # See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | # ---------------------------------------------------------------------------- | ||
| 10 | |||
| 11 | message(STATUS "=== Debug: start ops.optim.fused_sgd.graph_plugin.CMakeLists.txt ") | ||
| 12 | add_graph_plugin_sources() | ||
| @@ -0,0 +1,341 @@ | |||
| 1 | { | ||
| 2 | "op_type": "FusedSgd", | ||
| 3 | "op_list": [ | ||
| 4 | { | ||
| 5 | "bin_filename": "FusedSgd_0", | ||
| 6 | "inputs": [ | ||
| 7 | [ | ||
| 8 | { | ||
| 9 | "name": "params", | ||
| 10 | "index": 0, | ||
| 11 | "dtype": "bfloat16", | ||
| 12 | "format": "ND", | ||
| 13 | "paramType": "dynamic", | ||
| 14 | "shape": [-2] | ||
| 15 | } | ||
| 16 | ], | ||
| 17 | [ | ||
| 18 | { | ||
| 19 | "name": "grads", | ||
| 20 | "index": 1, | ||
| 21 | "dtype": "bfloat16", | ||
| 22 | "format": "ND", | ||
| 23 | "paramType": "dynamic", | ||
| 24 | "shape": [-2] | ||
| 25 | } | ||
| 26 | ], | ||
| 27 | [ | ||
| 28 | { | ||
| 29 | "name": "momentum_buffer_list", | ||
| 30 | "index": 2, | ||
| 31 | "dtype": "bfloat16", | ||
| 32 | "format": "ND", | ||
| 33 | "paramType": "dynamic", | ||
| 34 | "shape": [-2] | ||
| 35 | } | ||
| 36 | ], | ||
| 37 | { | ||
| 38 | "name": "grad_scale", | ||
| 39 | "index": 3, | ||
| 40 | "dtype": "float32", | ||
| 41 | "format": "ND", | ||
| 42 | "paramType": "optional", | ||
| 43 | "shape": [-2] | ||
| 44 | } | ||
| 45 | ], | ||
| 46 | "outputs": [ | ||
| 47 | [ | ||
| 48 | { | ||
| 49 | "name": "params", | ||
| 50 | "index": 0, | ||
| 51 | "dtype": "bfloat16", | ||
| 52 | "format": "ND", | ||
| 53 | "paramType": "dynamic", | ||
| 54 | "shape": [-2] | ||
| 55 | } | ||
| 56 | ], | ||
| 57 | [ | ||
| 58 | { | ||
| 59 | "name": "grads", | ||
| 60 | "index": 1, | ||
| 61 | "dtype": "bfloat16", | ||
| 62 | "format": "ND", | ||
| 63 | "paramType": "dynamic", | ||
| 64 | "shape": [-2] | ||
| 65 | } | ||
| 66 | ], | ||
| 67 | [ | ||
| 68 | { | ||
| 69 | "name": "momentum_buffer_list_out", | ||
| 70 | "index": 2, | ||
| 71 | "dtype": "bfloat16", | ||
| 72 | "format": "ND", | ||
| 73 | "paramType": "dynamic", | ||
| 74 | "shape": [-2] | ||
| 75 | } | ||
| 76 | ] | ||
| 77 | ], | ||
| 78 | "attrs": [ | ||
| 79 | { | ||
| 80 | "name": "weight_decay", | ||
| 81 | "dtype": "float", | ||
| 82 | "value": 0.0 | ||
| 83 | }, | ||
| 84 | { | ||
| 85 | "name": "momentum", | ||
| 86 | "dtype": "float", | ||
| 87 | "value": 0.0 | ||
| 88 | }, | ||
| 89 | { | ||
| 90 | "name": "lr", | ||
| 91 | "dtype": "float", | ||
| 92 | "value": 0.01 | ||
| 93 | }, | ||
| 94 | { | ||
| 95 | "name": "dampening", | ||
| 96 | "dtype": "float", | ||
| 97 | "value": 0.0 | ||
| 98 | }, | ||
| 99 | { | ||
| 100 | "name": "nesterov", | ||
| 101 | "dtype": "bool", | ||
| 102 | "value": false | ||
| 103 | }, | ||
| 104 | { | ||
| 105 | "name": "maximize", | ||
| 106 | "dtype": "bool", | ||
| 107 | "value": false | ||
| 108 | }, | ||
| 109 | { | ||
| 110 | "name": "is_first_step", | ||
| 111 | "dtype": "bool", | ||
| 112 | "value": true | ||
| 113 | } | ||
| 114 | ] | ||
| 115 | }, | ||
| 116 | { | ||
| 117 | "bin_filename": "FusedSgd_1", | ||
| 118 | "inputs": [ | ||
| 119 | [ | ||
| 120 | { | ||
| 121 | "name": "params", | ||
| 122 | "index": 0, | ||
| 123 | "dtype": "float16", | ||
| 124 | "format": "ND", | ||
| 125 | "paramType": "dynamic", | ||
| 126 | "shape": [-2] | ||
| 127 | } | ||
| 128 | ], | ||
| 129 | [ | ||
| 130 | { | ||
| 131 | "name": "grads", | ||
| 132 | "index": 1, | ||
| 133 | "dtype": "float16", | ||
| 134 | "format": "ND", | ||
| 135 | "paramType": "dynamic", | ||
| 136 | "shape": [-2] | ||
| 137 | } | ||
| 138 | ], | ||
| 139 | [ | ||
| 140 | { | ||
| 141 | "name": "momentum_buffer_list", | ||
| 142 | "index": 2, | ||
| 143 | "dtype": "float16", | ||
| 144 | "format": "ND", | ||
| 145 | "paramType": "dynamic", | ||
| 146 | "shape": [-2] | ||
| 147 | } | ||
| 148 | ], | ||
| 149 | { | ||
| 150 | "name": "grad_scale", | ||
| 151 | "index": 3, | ||
| 152 | "dtype": "float32", | ||
| 153 | "format": "ND", | ||
| 154 | "paramType": "optional", | ||
| 155 | "shape": [-2] | ||
| 156 | } | ||
| 157 | ], | ||
| 158 | "outputs": [ | ||
| 159 | [ | ||
| 160 | { | ||
| 161 | "name": "params", | ||
| 162 | "index": 0, | ||
| 163 | "dtype": "float16", | ||
| 164 | "format": "ND", | ||
| 165 | "paramType": "dynamic", | ||
| 166 | "shape": [-2] | ||
| 167 | } | ||
| 168 | ], | ||
| 169 | [ | ||
| 170 | { | ||
| 171 | "name": "grads", | ||
| 172 | "index": 1, | ||
| 173 | "dtype": "float16", | ||
| 174 | "format": "ND", | ||
| 175 | "paramType": "dynamic", | ||
| 176 | "shape": [-2] | ||
| 177 | } | ||
| 178 | ], | ||
| 179 | [ | ||
| 180 | { | ||
| 181 | "name": "momentum_buffer_list_out", | ||
| 182 | "index": 2, | ||
| 183 | "dtype": "float16", | ||
| 184 | "format": "ND", | ||
| 185 | "paramType": "dynamic", | ||
| 186 | "shape": [-2] | ||
| 187 | } | ||
| 188 | ] | ||
| 189 | ], | ||
| 190 | "attrs": [ | ||
| 191 | { | ||
| 192 | "name": "weight_decay", | ||
| 193 | "dtype": "float", | ||
| 194 | "value": 0.0 | ||
| 195 | }, | ||
| 196 | { | ||
| 197 | "name": "momentum", | ||
| 198 | "dtype": "float", | ||
| 199 | "value": 0.0 | ||
| 200 | }, | ||
| 201 | { | ||
| 202 | "name": "lr", | ||
| 203 | "dtype": "float", | ||
| 204 | "value": 0.01 | ||
| 205 | }, | ||
| 206 | { | ||
| 207 | "name": "dampening", | ||
| 208 | "dtype": "float", | ||
| 209 | "value": 0.0 | ||
| 210 | }, | ||
| 211 | { | ||
| 212 | "name": "nesterov", | ||
| 213 | "dtype": "bool", | ||
| 214 | "value": false | ||
| 215 | }, | ||
| 216 | { | ||
| 217 | "name": "maximize", | ||
| 218 | "dtype": "bool", | ||
| 219 | "value": false | ||
| 220 | }, | ||
| 221 | { | ||
| 222 | "name": "is_first_step", | ||
| 223 | "dtype": "bool", | ||
| 224 | "value": true | ||
| 225 | } | ||
| 226 | ] | ||
| 227 | }, | ||
| 228 | { | ||
| 229 | "bin_filename": "FusedSgd_2", | ||
| 230 | "inputs": [ | ||
| 231 | [ | ||
| 232 | { | ||
| 233 | "name": "params", | ||
| 234 | "index": 0, | ||
| 235 | "dtype": "float32", | ||
| 236 | "format": "ND", | ||
| 237 | "paramType": "dynamic", | ||
| 238 | "shape": [-2] | ||
| 239 | } | ||
| 240 | ], | ||
| 241 | [ | ||
| 242 | { | ||
| 243 | "name": "grads", | ||
| 244 | "index": 1, | ||
| 245 | "dtype": "float32", | ||
| 246 | "format": "ND", | ||
| 247 | "paramType": "dynamic", | ||
| 248 | "shape": [-2] | ||
| 249 | } | ||
| 250 | ], | ||
| 251 | [ | ||
| 252 | { | ||
| 253 | "name": "momentum_buffer_list", | ||
| 254 | "index": 2, | ||
| 255 | "dtype": "float32", | ||
| 256 | "format": "ND", | ||
| 257 | "paramType": "dynamic", | ||
| 258 | "shape": [-2] | ||
| 259 | } | ||
| 260 | ], | ||
| 261 | { | ||
| 262 | "name": "grad_scale", | ||
| 263 | "index": 3, | ||
| 264 | "dtype": "float32", | ||
| 265 | "format": "ND", | ||
| 266 | "paramType": "optional", | ||
| 267 | "shape": [-2] | ||
| 268 | } | ||
| 269 | ], | ||
| 270 | "outputs": [ | ||
| 271 | [ | ||
| 272 | { | ||
| 273 | "name": "params", | ||
| 274 | "index": 0, | ||
| 275 | "dtype": "float32", | ||
| 276 | "format": "ND", | ||
| 277 | "paramType": "dynamic", | ||
| 278 | "shape": [-2] | ||
| 279 | } | ||
| 280 | ], | ||
| 281 | [ | ||
| 282 | { | ||
| 283 | "name": "grads", | ||
| 284 | "index": 1, | ||
| 285 | "dtype": "float32", | ||
| 286 | "format": "ND", | ||
| 287 | "paramType": "dynamic", | ||
| 288 | "shape": [-2] | ||
| 289 | } | ||
| 290 | ], | ||
| 291 | [ | ||
| 292 | { | ||
| 293 | "name": "momentum_buffer_list_out", | ||
| 294 | "index": 2, | ||
| 295 | "dtype": "float32", | ||
| 296 | "format": "ND", | ||
| 297 | "paramType": "dynamic", | ||
| 298 | "shape": [-2] | ||
| 299 | } | ||
| 300 | ] | ||
| 301 | ], | ||
| 302 | "attrs": [ | ||
| 303 | { | ||
| 304 | "name": "weight_decay", | ||
| 305 | "dtype": "float", | ||
| 306 | "value": 0.0 | ||
| 307 | }, | ||
| 308 | { | ||
| 309 | "name": "momentum", | ||
| 310 | "dtype": "float", | ||
| 311 | "value": 0.0 | ||
| 312 | }, | ||
| 313 | { | ||
| 314 | "name": "lr", | ||
| 315 | "dtype": "float", | ||
| 316 | "value": 0.01 | ||
| 317 | }, | ||
| 318 | { | ||
| 319 | "name": "dampening", | ||
| 320 | "dtype": "float", | ||
| 321 | "value": 0.0 | ||
| 322 | }, | ||
| 323 | { | ||
| 324 | "name": "nesterov", | ||
| 325 | "dtype": "bool", | ||
| 326 | "value": false | ||
| 327 | }, | ||
| 328 | { | ||
| 329 | "name": "maximize", | ||
| 330 | "dtype": "bool", | ||
| 331 | "value": false | ||
| 332 | }, | ||
| 333 | { | ||
| 334 | "name": "is_first_step", | ||
| 335 | "dtype": "bool", | ||
| 336 | "value": true | ||
| 337 | } | ||
| 338 | ] | ||
| 339 | } | ||
| 340 | ] | ||
| 341 | } | ||
| @@ -0,0 +1,341 @@ | |||
| 1 | { | ||
| 2 | "op_type": "FusedSgd", | ||
| 3 | "op_list": [ | ||
| 4 | { | ||
| 5 | "bin_filename": "FusedSgd_0", | ||
| 6 | "inputs": [ | ||
| 7 | [ | ||
| 8 | { | ||
| 9 | "name": "params", | ||
| 10 | "index": 0, | ||
| 11 | "dtype": "bfloat16", | ||
| 12 | "format": "ND", | ||
| 13 | "paramType": "dynamic", | ||
| 14 | "shape": [-2] | ||
| 15 | } | ||
| 16 | ], | ||
| 17 | [ | ||
| 18 | { | ||
| 19 | "name": "grads", | ||
| 20 | "index": 1, | ||
| 21 | "dtype": "bfloat16", | ||
| 22 | "format": "ND", | ||
| 23 | "paramType": "dynamic", | ||
| 24 | "shape": [-2] | ||
| 25 | } | ||
| 26 | ], | ||
| 27 | [ | ||
| 28 | { | ||
| 29 | "name": "momentum_buffer_list", | ||
| 30 | "index": 2, | ||
| 31 | "dtype": "bfloat16", | ||
| 32 | "format": "ND", | ||
| 33 | "paramType": "dynamic", | ||
| 34 | "shape": [-2] | ||
| 35 | } | ||
| 36 | ], | ||
| 37 | { | ||
| 38 | "name": "grad_scale", | ||
| 39 | "index": 3, | ||
| 40 | "dtype": "float32", | ||
| 41 | "format": "ND", | ||
| 42 | "paramType": "optional", | ||
| 43 | "shape": [-2] | ||
| 44 | } | ||
| 45 | ], | ||
| 46 | "outputs": [ | ||
| 47 | [ | ||
| 48 | { | ||
| 49 | "name": "params", | ||
| 50 | "index": 0, | ||
| 51 | "dtype": "bfloat16", | ||
| 52 | "format": "ND", | ||
| 53 | "paramType": "dynamic", | ||
| 54 | "shape": [-2] | ||
| 55 | } | ||
| 56 | ], | ||
| 57 | [ | ||
| 58 | { | ||
| 59 | "name": "grads", | ||
| 60 | "index": 1, | ||
| 61 | "dtype": "bfloat16", | ||
| 62 | "format": "ND", | ||
| 63 | "paramType": "dynamic", | ||
| 64 | "shape": [-2] | ||
| 65 | } | ||
| 66 | ], | ||
| 67 | [ | ||
| 68 | { | ||
| 69 | "name": "momentum_buffer_list_out", | ||
| 70 | "index": 2, | ||
| 71 | "dtype": "bfloat16", | ||
| 72 | "format": "ND", | ||
| 73 | "paramType": "dynamic", | ||
| 74 | "shape": [-2] | ||
| 75 | } | ||
| 76 | ] | ||
| 77 | ], | ||
| 78 | "attrs": [ | ||
| 79 | { | ||
| 80 | "name": "weight_decay", | ||
| 81 | "dtype": "float", | ||
| 82 | "value": 0.0 | ||
| 83 | }, | ||
| 84 | { | ||
| 85 | "name": "momentum", | ||
| 86 | "dtype": "float", | ||
| 87 | "value": 0.0 | ||
| 88 | }, | ||
| 89 | { | ||
| 90 | "name": "lr", | ||
| 91 | "dtype": "float", | ||
| 92 | "value": 0.01 | ||
| 93 | }, | ||
| 94 | { | ||
| 95 | "name": "dampening", | ||
| 96 | "dtype": "float", | ||
| 97 | "value": 0.0 | ||
| 98 | }, | ||
| 99 | { | ||
| 100 | "name": "nesterov", | ||
| 101 | "dtype": "bool", | ||
| 102 | "value": false | ||
| 103 | }, | ||
| 104 | { | ||
| 105 | "name": "maximize", | ||
| 106 | "dtype": "bool", | ||
| 107 | "value": false | ||
| 108 | }, | ||
| 109 | { | ||
| 110 | "name": "is_first_step", | ||
| 111 | "dtype": "bool", | ||
| 112 | "value": true | ||
| 113 | } | ||
| 114 | ] | ||
| 115 | }, | ||
| 116 | { | ||
| 117 | "bin_filename": "FusedSgd_1", | ||
| 118 | "inputs": [ | ||
| 119 | [ | ||
| 120 | { | ||
| 121 | "name": "params", | ||
| 122 | "index": 0, | ||
| 123 | "dtype": "float16", | ||
| 124 | "format": "ND", | ||
| 125 | "paramType": "dynamic", | ||
| 126 | "shape": [-2] | ||
| 127 | } | ||
| 128 | ], | ||
| 129 | [ | ||
| 130 | { | ||
| 131 | "name": "grads", | ||
| 132 | "index": 1, | ||
| 133 | "dtype": "float16", | ||
| 134 | "format": "ND", | ||
| 135 | "paramType": "dynamic", | ||
| 136 | "shape": [-2] | ||
| 137 | } | ||
| 138 | ], | ||
| 139 | [ | ||
| 140 | { | ||
| 141 | "name": "momentum_buffer_list", | ||
| 142 | "index": 2, | ||
| 143 | "dtype": "float16", | ||
| 144 | "format": "ND", | ||
| 145 | "paramType": "dynamic", | ||
| 146 | "shape": [-2] | ||
| 147 | } | ||
| 148 | ], | ||
| 149 | { | ||
| 150 | "name": "grad_scale", | ||
| 151 | "index": 3, | ||
| 152 | "dtype": "float32", | ||
| 153 | "format": "ND", | ||
| 154 | "paramType": "optional", | ||
| 155 | "shape": [-2] | ||
| 156 | } | ||
| 157 | ], | ||
| 158 | "outputs": [ | ||
| 159 | [ | ||
| 160 | { | ||
| 161 | "name": "params", | ||
| 162 | "index": 0, | ||
| 163 | "dtype": "float16", | ||
| 164 | "format": "ND", | ||
| 165 | "paramType": "dynamic", | ||
| 166 | "shape": [-2] | ||
| 167 | } | ||
| 168 | ], | ||
| 169 | [ | ||
| 170 | { | ||
| 171 | "name": "grads", | ||
| 172 | "index": 1, | ||
| 173 | "dtype": "float16", | ||
| 174 | "format": "ND", | ||
| 175 | "paramType": "dynamic", | ||
| 176 | "shape": [-2] | ||
| 177 | } | ||
| 178 | ], | ||
| 179 | [ | ||
| 180 | { | ||
| 181 | "name": "momentum_buffer_list_out", | ||
| 182 | "index": 2, | ||
| 183 | "dtype": "float16", | ||
| 184 | "format": "ND", | ||
| 185 | "paramType": "dynamic", | ||
| 186 | "shape": [-2] | ||
| 187 | } | ||
| 188 | ] | ||
| 189 | ], | ||
| 190 | "attrs": [ | ||
| 191 | { | ||
| 192 | "name": "weight_decay", | ||
| 193 | "dtype": "float", | ||
| 194 | "value": 0.0 | ||
| 195 | }, | ||
| 196 | { | ||
| 197 | "name": "momentum", | ||
| 198 | "dtype": "float", | ||
| 199 | "value": 0.0 | ||
| 200 | }, | ||
| 201 | { | ||
| 202 | "name": "lr", | ||
| 203 | "dtype": "float", | ||
| 204 | "value": 0.01 | ||
| 205 | }, | ||
| 206 | { | ||
| 207 | "name": "dampening", | ||
| 208 | "dtype": "float", | ||
| 209 | "value": 0.0 | ||
| 210 | }, | ||
| 211 | { | ||
| 212 | "name": "nesterov", | ||
| 213 | "dtype": "bool", | ||
| 214 | "value": false | ||
| 215 | }, | ||
| 216 | { | ||
| 217 | "name": "maximize", | ||
| 218 | "dtype": "bool", | ||
| 219 | "value": false | ||
| 220 | }, | ||
| 221 | { | ||
| 222 | "name": "is_first_step", | ||
| 223 | "dtype": "bool", | ||
| 224 | "value": true | ||
| 225 | } | ||
| 226 | ] | ||
| 227 | }, | ||
| 228 | { | ||
| 229 | "bin_filename": "FusedSgd_2", | ||
| 230 | "inputs": [ | ||
| 231 | [ | ||
| 232 | { | ||
| 233 | "name": "params", | ||
| 234 | "index": 0, | ||
| 235 | "dtype": "float32", | ||
| 236 | "format": "ND", | ||
| 237 | "paramType": "dynamic", | ||
| 238 | "shape": [-2] | ||
| 239 | } | ||
| 240 | ], | ||
| 241 | [ | ||
| 242 | { | ||
| 243 | "name": "grads", | ||
| 244 | "index": 1, | ||
| 245 | "dtype": "float32", | ||
| 246 | "format": "ND", | ||
| 247 | "paramType": "dynamic", | ||
| 248 | "shape": [-2] | ||
| 249 | } | ||
| 250 | ], | ||
| 251 | [ | ||
| 252 | { | ||
| 253 | "name": "momentum_buffer_list", | ||
| 254 | "index": 2, | ||
| 255 | "dtype": "float32", | ||
| 256 | "format": "ND", | ||
| 257 | "paramType": "dynamic", | ||
| 258 | "shape": [-2] | ||
| 259 | } | ||
| 260 | ], | ||
| 261 | { | ||
| 262 | "name": "grad_scale", | ||
| 263 | "index": 3, | ||
| 264 | "dtype": "float32", | ||
| 265 | "format": "ND", | ||
| 266 | "paramType": "optional", | ||
| 267 | "shape": [-2] | ||
| 268 | } | ||
| 269 | ], | ||
| 270 | "outputs": [ | ||
| 271 | [ | ||
| 272 | { | ||
| 273 | "name": "params", | ||
| 274 | "index": 0, | ||
| 275 | "dtype": "float32", | ||
| 276 | "format": "ND", | ||
| 277 | "paramType": "dynamic", | ||
| 278 | "shape": [-2] | ||
| 279 | } | ||
| 280 | ], | ||
| 281 | [ | ||
| 282 | { | ||
| 283 | "name": "grads", | ||
| 284 | "index": 1, | ||
| 285 | "dtype": "float32", | ||
| 286 | "format": "ND", | ||
| 287 | "paramType": "dynamic", | ||
| 288 | "shape": [-2] | ||
| 289 | } | ||
| 290 | ], | ||
| 291 | [ | ||
| 292 | { | ||
| 293 | "name": "momentum_buffer_list_out", | ||
| 294 | "index": 2, | ||
| 295 | "dtype": "float32", | ||
| 296 | "format": "ND", | ||
| 297 | "paramType": "dynamic", | ||
| 298 | "shape": [-2] | ||
| 299 | } | ||
| 300 | ] | ||
| 301 | ], | ||
| 302 | "attrs": [ | ||
| 303 | { | ||
| 304 | "name": "weight_decay", | ||
| 305 | "dtype": "float", | ||
| 306 | "value": 0.0 | ||
| 307 | }, | ||
| 308 | { | ||
| 309 | "name": "momentum", | ||
| 310 | "dtype": "float", | ||
| 311 | "value": 0.0 | ||
| 312 | }, | ||
| 313 | { | ||
| 314 | "name": "lr", | ||
| 315 | "dtype": "float", | ||
| 316 | "value": 0.01 | ||
| 317 | }, | ||
| 318 | { | ||
| 319 | "name": "dampening", | ||
| 320 | "dtype": "float", | ||
| 321 | "value": 0.0 | ||
| 322 | }, | ||
| 323 | { | ||
| 324 | "name": "nesterov", | ||
| 325 | "dtype": "bool", | ||
| 326 | "value": false | ||
| 327 | }, | ||
| 328 | { | ||
| 329 | "name": "maximize", | ||
| 330 | "dtype": "bool", | ||
| 331 | "value": false | ||
| 332 | }, | ||
| 333 | { | ||
| 334 | "name": "is_first_step", | ||
| 335 | "dtype": "bool", | ||
| 336 | "value": true | ||
| 337 | } | ||
| 338 | ] | ||
| 339 | } | ||
| 340 | ] | ||
| 341 | } | ||
| @@ -0,0 +1,80 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd_def.cpp | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | namespace ops { | ||
| 18 | static constexpr float LR_DEFAULT = 0.001; | ||
| 19 | static constexpr float WEIGHT_DECAY_DEFAULT = 0.0; | ||
| 20 | static constexpr float MOMENTUM_DEFAULT = 0.0; | ||
| 21 | static constexpr float DAMPENING_DEFAULT = 0.0; | ||
| 22 | |||
| 23 | class FusedSgd : public OpDef | ||
| 24 | { | ||
| 25 | public: | ||
| 26 | explicit FusedSgd(const char* name) : OpDef(name) | ||
| 27 | { | ||
| 28 | this->Input("params") | ||
| 29 | .ParamType(DYNAMIC) | ||
| 30 | .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) | ||
| 31 | .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 32 | .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 33 | .AutoContiguous(); | ||
| 34 | this->Input("grads") | ||
| 35 | .ParamType(DYNAMIC) | ||
| 36 | .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) | ||
| 37 | .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 38 | .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 39 | .AutoContiguous(); | ||
| 40 | this->Input("x") | ||
| 41 | .ParamType(DYNAMIC) | ||
| 42 | .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) | ||
| 43 | .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 44 | .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 45 | .AutoContiguous(); | ||
| 46 | this->Input("grad_scale") | ||
| 47 | .ParamType(OPTIONAL) | ||
| 48 | .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT}) | ||
| 49 | .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 50 | .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 51 | .AutoContiguous(); | ||
| 52 | this->Output("params") | ||
| 53 | .ParamType(DYNAMIC) | ||
| 54 | .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) | ||
| 55 | .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 56 | .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 57 | this->Output("grads") | ||
| 58 | .ParamType(DYNAMIC) | ||
| 59 | .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) | ||
| 60 | .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 61 | .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 62 | this->Output("y") | ||
| 63 | .ParamType(DYNAMIC) | ||
| 64 | .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16}) | ||
| 65 | .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}) | ||
| 66 | .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND}); | ||
| 67 | this->Attr("weight_decay").AttrType(OPTIONAL).Float(WEIGHT_DECAY_DEFAULT); | ||
| 68 | this->Attr("momentum").AttrType(OPTIONAL).Float(MOMENTUM_DEFAULT); | ||
| 69 | this->Attr("lr").AttrType(OPTIONAL).Float(LR_DEFAULT); | ||
| 70 | this->Attr("dampening").AttrType(OPTIONAL).Float(DAMPENING_DEFAULT); | ||
| 71 | this->Attr("nesterov").AttrType(OPTIONAL).Bool(false); | ||
| 72 | this->Attr("maximize").AttrType(OPTIONAL).Bool(false); | ||
| 73 | this->Attr("is_first_step").AttrType(OPTIONAL).Bool(true); | ||
| 74 | this->AICore().AddConfig("ascend910b"); | ||
| 75 | this->AICore().AddConfig("ascend910_93"); | ||
| 76 | } | ||
| 77 | }; | ||
| 78 | |||
| 79 | OP_ADD(FusedSgd); | ||
| 80 | } // namespace ops | ||
| @@ -0,0 +1,132 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd_infershape.cpp | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | using namespace ge; | ||
| 19 | static constexpr size_t INPUT_PARAMS_INDEX = 0; | ||
| 20 | static constexpr size_t INPUT_GRADS_INDEX = 1; | ||
| 21 | static constexpr size_t INPUT_MOMENTUM_BUFFER_INDEX = 2; | ||
| 22 | static constexpr size_t OUTPUT_PARAMS_INDEX = 0; | ||
| 23 | static constexpr size_t OUTPUT_GRADS_INDEX = 1; | ||
| 24 | static constexpr size_t OUTPUT_MOMENTUM_BUFFER_INDEX = 2; | ||
| 25 | |||
| 26 | namespace ops { | ||
| 27 | static ge::graphStatus InferShapeForFusedSgd(gert::InferShapeContext* context) | ||
| 28 | { | ||
| 29 | OP_LOGD(context, "Begin to do InferShapeForFusedSgd."); | ||
| 30 | |||
| 31 | auto computeNodeInfo = context->GetComputeNodeInfo(); | ||
| 32 | OP_CHECK_NULL_WITH_CONTEXT(context, computeNodeInfo); | ||
| 33 | |||
| 34 | auto paramsInstanceInfo = computeNodeInfo->GetInputInstanceInfo(INPUT_PARAMS_INDEX); | ||
| 35 | OP_CHECK_NULL_WITH_CONTEXT(context, paramsInstanceInfo); | ||
| 36 | auto inputNum = paramsInstanceInfo->GetInstanceNum(); | ||
| 37 | if (inputNum == 0) { | ||
| 38 | OP_LOGE(context, "input num must be greater than 0"); | ||
| 39 | return ge::GRAPH_FAILED; | ||
| 40 | } | ||
| 41 | |||
| 42 | // 判断tensorlist是否为空 | ||
| 43 | auto momentumBufferListInput = context->GetDynamicInputShape(INPUT_MOMENTUM_BUFFER_INDEX, 0); | ||
| 44 | if (momentumBufferListInput != nullptr) { | ||
| 45 | uint32_t momentumBufferListDims = momentumBufferListInput->GetDimNum(); | ||
| 46 | bool flag = true; | ||
| 47 | if (momentumBufferListDims < 1) { | ||
| 48 | flag = false; | ||
| 49 | } | ||
| 50 | if (flag) { | ||
| 51 | for(uint32_t i = 0; i < momentumBufferListDims; i++) { | ||
| 52 | int64_t dimValue = momentumBufferListInput->GetDim(i); | ||
| 53 | if (dimValue == 0) { | ||
| 54 | flag = false; | ||
| 55 | break; | ||
| 56 | } | ||
| 57 | } | ||
| 58 | } | ||
| 59 | if (!flag) { | ||
| 60 | momentumBufferListInput = nullptr; | ||
| 61 | } | ||
| 62 | } | ||
| 63 | |||
| 64 | auto paramsOutInstanceInfo = context->GetIrOutputInstanceInfo(OUTPUT_PARAMS_INDEX); | ||
| 65 | auto gradsOutInstanceInfo = context->GetIrOutputInstanceInfo(OUTPUT_GRADS_INDEX); | ||
| 66 | auto momentumOutInstanceInfo = context->GetIrOutputInstanceInfo(OUTPUT_MOMENTUM_BUFFER_INDEX); | ||
| 67 | |||
| 68 | for (uint32_t i = 0; i < inputNum; i++) { | ||
| 69 | const gert::Shape* paramsShape = context->GetDynamicInputShape(INPUT_PARAMS_INDEX, i); | ||
| 70 | OP_CHECK_NULL_WITH_CONTEXT(context, paramsShape); | ||
| 71 | const gert::Shape* gradsShape = context->GetDynamicInputShape(INPUT_GRADS_INDEX, i); | ||
| 72 | OP_CHECK_NULL_WITH_CONTEXT(context, gradsShape); | ||
| 73 | const gert::Shape* momentumShape = nullptr; | ||
| 74 | if (momentumBufferListInput != nullptr) { | ||
| 75 | momentumShape = context->GetDynamicInputShape(INPUT_MOMENTUM_BUFFER_INDEX, i); | ||
| 76 | OP_CHECK_NULL_WITH_CONTEXT(context, momentumShape); | ||
| 77 | } | ||
| 78 | |||
| 79 | if (*paramsShape != *gradsShape || (momentumShape != nullptr && *paramsShape != *momentumShape)) { | ||
| 80 | OP_LOGE(context, "params, grads and momentum_buffer_list should have the same shape"); | ||
| 81 | return ge::GRAPH_FAILED; | ||
| 82 | } | ||
| 83 | |||
| 84 | gert::Shape* paramsRefShape = context->GetOutputShape(paramsOutInstanceInfo->GetInstanceStart() + i); | ||
| 85 | OP_CHECK_NULL_WITH_CONTEXT(context, paramsRefShape); | ||
| 86 | gert::Shape* gradsRefShape = context->GetOutputShape(gradsOutInstanceInfo->GetInstanceStart() + i); | ||
| 87 | OP_CHECK_NULL_WITH_CONTEXT(context, gradsRefShape); | ||
| 88 | gert::Shape* momentumRefShape = nullptr; | ||
| 89 | if (momentumBufferListInput != nullptr) { | ||
| 90 | momentumRefShape = context->GetOutputShape(momentumOutInstanceInfo->GetInstanceStart() + i); | ||
| 91 | OP_CHECK_NULL_WITH_CONTEXT(context, momentumRefShape); | ||
| 92 | } | ||
| 93 | |||
| 94 | *paramsRefShape = *paramsShape; | ||
| 95 | *gradsRefShape = *gradsShape; | ||
| 96 | if (momentumBufferListInput != nullptr) { | ||
| 97 | *momentumRefShape = *momentumShape; | ||
| 98 | } | ||
| 99 | } | ||
| 100 | |||
| 101 | OP_LOGD(context, "End to do InferShapeForFusedSgd."); | ||
| 102 | return ge::GRAPH_SUCCESS; | ||
| 103 | } | ||
| 104 | |||
| 105 | static graphStatus InferDataTypeForFusedSgd(gert::InferDataTypeContext* context) | ||
| 106 | { | ||
| 107 | auto computeNodeInfo = context->GetComputeNodeInfo(); | ||
| 108 | OP_CHECK_NULL_WITH_CONTEXT(context, computeNodeInfo); | ||
| 109 | |||
| 110 | auto paramsInstanceInfo = computeNodeInfo->GetInputInstanceInfo(INPUT_PARAMS_INDEX); | ||
| 111 | OP_CHECK_NULL_WITH_CONTEXT(context, paramsInstanceInfo); | ||
| 112 | auto inputNum = paramsInstanceInfo->GetInstanceNum(); | ||
| 113 | |||
| 114 | auto paramsOutInstanceInfo = context->GetIrOutputInstanceInfo(OUTPUT_PARAMS_INDEX); | ||
| 115 | auto gradsOutInstanceInfo = context->GetIrOutputInstanceInfo(OUTPUT_GRADS_INDEX); | ||
| 116 | auto momentumOutInstanceInfo = context->GetIrOutputInstanceInfo(OUTPUT_MOMENTUM_BUFFER_INDEX); | ||
| 117 | |||
| 118 | for (uint32_t i = 0; i < inputNum; i++) { | ||
| 119 | context->SetOutputDataType(paramsOutInstanceInfo->GetInstanceStart() + i, | ||
| 120 | context->GetDynamicInputDataType(INPUT_PARAMS_INDEX, i)); | ||
| 121 | context->SetOutputDataType(gradsOutInstanceInfo->GetInstanceStart() + i, | ||
| 122 | context->GetDynamicInputDataType(INPUT_GRADS_INDEX, i)); | ||
| 123 | context->SetOutputDataType(momentumOutInstanceInfo->GetInstanceStart() + i, | ||
| 124 | context->GetDynamicInputDataType(INPUT_MOMENTUM_BUFFER_INDEX, i)); | ||
🟡 Medium Priority 在 当 momentum buffer 为空时（tensor 全部无元素）， 同一文件第27行的 变更：第123-124行新增，对 momentum buffer 的无条件访问。 建议：在 majiajian_hw 6月22日 评论： 6月22日 评论： | |||
| 125 | } | ||
| 126 | return GRAPH_SUCCESS; | ||
| 127 | } | ||
| 128 | |||
| 129 | IMPL_OP_INFERSHAPE(FusedSgd) | ||
| 130 | .InferShape(InferShapeForFusedSgd) | ||
| 131 | .InferDataType(InferDataTypeForFusedSgd); | ||
| 132 | } // namespace ops | ||
| @@ -0,0 +1,342 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd_tiling.cpp | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | using namespace std; | ||
| 23 | namespace optiling { | ||
| 24 | constexpr uint32_t INPUT_PARAMS_IDX = 0; | ||
| 25 | constexpr uint32_t INPUT_GRADS_IDX = 1; | ||
| 26 | constexpr uint32_t INPUT_MOMENTUM_BUFFER_IDX = 2; | ||
| 27 | constexpr uint32_t INPUT_GRAD_SCALE_IDX = 3; | ||
| 28 | constexpr uint32_t ATTR_WEIGHT_DECAY_IDX = 0; | ||
| 29 | constexpr uint32_t ATTR_MOMENTUM_IDX = 1; | ||
| 30 | constexpr uint32_t ATTR_LR_IDX = 2; | ||
| 31 | constexpr uint32_t ATTR_DAMPENING_IDX = 3; | ||
| 32 | constexpr uint32_t ATTR_NESTEROV_IDX = 4; | ||
| 33 | constexpr uint32_t ATTR_MAXIMIZE_IDX = 5; | ||
| 34 | constexpr uint32_t ATTR_IS_FIRST_STEP_IDX = 6; | ||
| 35 | constexpr uint32_t ONE_BLK_NUM = 16; | ||
| 36 | constexpr uint32_t ONE_BLK_NUM_FP32 = 8; | ||
| 37 | constexpr uint32_t BYTE_ONE_BLK = 32; | ||
| 38 | constexpr uint32_t TBUFFER_NUM = 3; | ||
| 39 | constexpr uint32_t BUFFER_NUM = 2; | ||
| 40 | constexpr uint32_t FP16_BF16_DTYPE_SIZE = 2; | ||
| 41 | constexpr uint32_t FP32_DTYPE_SIZE = 4; | ||
| 42 | |||
| 43 | std::string FusedSgdTiling::TilingDataToString() const | ||
| 44 | { | ||
| 45 | return "weightDecay = " + std::to_string(weightDecay_) + \ | ||
| 46 | ", momentum = " + std::to_string(momentum_) + \ | ||
| 47 | ", lr = " + std::to_string(lr_) + \ | ||
| 48 | ", dampening = " + std::to_string(dampening_) + \ | ||
| 49 | ", nesterov = " + std::to_string(nesterov_) + \ | ||
| 50 | ", maximize = " + std::to_string(maximize_) + \ | ||
| 51 | ", isFirstStep = " + std::to_string(isFirstStep_) + \ | ||
| 52 | ", useGradScale = " + std::to_string(useGradScale_) + \ | ||
| 53 | ", useMomentum = " + std::to_string(useMomentum_) + \ | ||
| 54 | ", tensorNum = " + std::to_string(tensorNum_) + \ | ||
| 55 | ", tensorsPerCore = " + std::to_string(tensorsPerCore_) + \ | ||
| 56 | ", usedCoreNum = " + std::to_string(usedCoreNum_) + \ | ||
| 57 | ", coreCalcMax = " + std::to_string(coreCalcMax_); | ||
| 58 | } | ||
| 59 | |||
| 60 | // 获取硬件信息 | ||
| 61 | ge::graphStatus FusedSgdTiling::GetPlatformInfo() | ||
| 62 | { | ||
| 63 | auto platformInfo = context_->GetPlatformInfo(); | ||
| 64 | OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo); | ||
| 65 | auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo); | ||
| 66 | coreNum_ = static_cast<uint32_t>(ascendcPlatform.GetCoreNumAiv()); | ||
| 67 | OP_CHECK_IF(coreNum_ == 0, OP_LOGE(context_, "coreNum is 0"), return ge::GRAPH_FAILED); | ||
| 68 | ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize_); | ||
| 69 | OP_CHECK_IF(ubSize_ == 0, OP_LOGE(context_, "ubSize is 0"), return ge::GRAPH_FAILED); | ||
| 70 | sysWorkspaceSize_ = ascendcPlatform.GetLibApiWorkSpaceSize(); | ||
| 71 | return ge::GRAPH_SUCCESS; | ||
| 72 | } | ||
| 73 | |||
| 74 | ge::graphStatus FusedSgdTiling::GetAttrInfo() | ||
| 75 | { | ||
| 76 | // 获取属性 | ||
| 77 | auto* attrs = context_->GetAttrs(); | ||
| 78 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrs); | ||
| 79 | const float* attrWeightDecay = attrs->GetAttrPointer<float>(ATTR_WEIGHT_DECAY_IDX); | ||
| 80 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrWeightDecay); | ||
| 81 | weightDecay_ = static_cast<float>(*attrWeightDecay); | ||
| 82 | const float* attrMomentum = attrs->GetAttrPointer<float>(ATTR_MOMENTUM_IDX); | ||
| 83 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrMomentum); | ||
| 84 | momentum_ = static_cast<float>(*attrMomentum); | ||
| 85 | const float* attrLr = attrs->GetAttrPointer<float>(ATTR_LR_IDX); | ||
| 86 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrLr); | ||
| 87 | lr_ = static_cast<float>(*attrLr); | ||
| 88 | const float* attrDampening = attrs->GetAttrPointer<float>(ATTR_DAMPENING_IDX); | ||
| 89 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrDampening); | ||
| 90 | dampening_ = static_cast<float>(*attrDampening); | ||
| 91 | const bool* attrNesterov = attrs->GetAttrPointer<bool>(ATTR_NESTEROV_IDX); | ||
| 92 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrNesterov); | ||
| 93 | nesterov_ = static_cast<uint32_t>(*attrNesterov ? 1 : 0); | ||
| 94 | const bool* attrMaximize = attrs->GetAttrPointer<bool>(ATTR_MAXIMIZE_IDX); | ||
| 95 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrMaximize); | ||
| 96 | maximize_ = static_cast<uint32_t>(*attrMaximize ? 1 : 0); | ||
| 97 | const bool* attrIsFirstStep = attrs->GetAttrPointer<bool>(ATTR_IS_FIRST_STEP_IDX); | ||
| 98 | OP_CHECK_NULL_WITH_CONTEXT(context_, attrIsFirstStep); | ||
| 99 | isFirstStep_ = static_cast<uint32_t>(*attrIsFirstStep ? 1 : 0); | ||
| 100 | |||
| 101 | return ge::GRAPH_SUCCESS; | ||
| 102 | } | ||
| 103 | |||
| 104 | void FusedSgdTiling::CheckOptionalInputs() | ||
| 105 | { | ||
| 106 | // 判断gradScale为空 | ||
| 107 | auto shapeInput = context_->GetOptionalInputTensor(INPUT_GRAD_SCALE_IDX); | ||
| 108 | if (shapeInput != nullptr) { | ||
| 109 | const gert::Shape& inputShapeGradScale = shapeInput->GetStorageShape(); | ||
| 110 | uint32_t gradScaleDims = inputShapeGradScale.GetDimNum(); | ||
| 111 | bool flag = true; | ||
| 112 | if (gradScaleDims < 1) { | ||
| 113 | flag = false; | ||
| 114 | } | ||
| 115 | if (flag) { | ||
| 116 | for(uint32_t i = 0; i < gradScaleDims; i++) { | ||
| 117 | int64_t dimValue = inputShapeGradScale.GetDim(i); | ||
| 118 | if (dimValue == 0) { | ||
| 119 | flag = false; | ||
| 120 | break; | ||
| 121 | } | ||
| 122 | } | ||
| 123 | } | ||
| 124 | if (flag) { | ||
| 125 | useGradScale_ = 1; | ||
| 126 | } else { | ||
| 127 | useGradScale_ = 0; | ||
| 128 | } | ||
| 129 | } else { | ||
| 130 | useGradScale_ = 0; | ||
| 131 | } | ||
| 132 | |||
| 133 | // 判断tensorlist为空 | ||
| 134 | auto momentumBufferListInput = context_->GetDynamicInputShape(INPUT_MOMENTUM_BUFFER_IDX, 0); | ||
| 135 | if (momentumBufferListInput != nullptr) { | ||
| 136 | const gert::Shape& inputShapeMomentumBufferList = momentumBufferListInput->GetStorageShape(); | ||
| 137 | uint32_t momentumBufferListDims = inputShapeMomentumBufferList.GetDimNum(); | ||
| 138 | bool flag = true; | ||
| 139 | if (momentumBufferListDims < 1) { | ||
| 140 | flag = false; | ||
| 141 | } | ||
| 142 | if (flag) { | ||
| 143 | for(uint32_t i = 0; i < momentumBufferListDims; i++) { | ||
| 144 | int64_t dimValue = inputShapeMomentumBufferList.GetDim(i); | ||
| 145 | if (dimValue == 0) { | ||
| 146 | flag = false; | ||
| 147 | break; | ||
| 148 | } | ||
| 149 | } | ||
| 150 | } | ||
| 151 | if (flag) { | ||
| 152 | useMomentum_ = 1; | ||
| 153 | } else { | ||
| 154 | useMomentum_ = 0; | ||
| 155 | } | ||
| 156 | } else { | ||
| 157 | useMomentum_ = 0; | ||
| 158 | } | ||
| 159 | } | ||
| 160 | |||
| 161 | static ge::graphStatus CheckInputDtype(gert::TilingContext* context, uint32_t useMomentum_) | ||
| 162 | { | ||
| 163 | auto dtypeInput = context->GetDynamicInputDesc(INPUT_PARAMS_IDX, 0); | ||
| 164 | OP_CHECK_NULL_WITH_CONTEXT(context, dtypeInput); | ||
| 165 | auto paramsDtype = dtypeInput->GetDataType(); | ||
| 166 | |||
| 167 | dtypeInput = context->GetDynamicInputDesc(INPUT_GRADS_IDX, 0); | ||
| 168 | OP_CHECK_NULL_WITH_CONTEXT(context, dtypeInput); | ||
| 169 | auto gradsDtype = dtypeInput->GetDataType(); | ||
| 170 | |||
| 171 | bool isDiffDtype = | ||
| 172 | (paramsDtype != gradsDtype); | ||
| 173 | ge::DataType momentumDtype; | ||
| 174 | if(useMomentum_) { | ||
| 175 | dtypeInput = context->GetDynamicInputDesc(INPUT_MOMENTUM_BUFFER_IDX, 0); | ||
| 176 | OP_CHECK_NULL_WITH_CONTEXT(context, dtypeInput); | ||
| 177 | momentumDtype = dtypeInput->GetDataType(); | ||
| 178 | isDiffDtype = isDiffDtype || (paramsDtype != momentumDtype); | ||
| 179 | } | ||
| 180 | |||
| 181 | if (isDiffDtype) { | ||
| 182 | std::string dtypeMsg = Ops::Base::ToString(paramsDtype) + ", " + | ||
| 183 | Ops::Base::ToString(gradsDtype); | ||
| 184 | if(!useMomentum_) { | ||
| 185 | OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON( | ||
| 186 | context->GetNodeName(), "params, grads", dtypeMsg.c_str(), | ||
| 187 | "params, grads should have the same dtype"); | ||
| 188 | } else { | ||
| 189 | dtypeMsg = dtypeMsg + " and " + Ops::Base::ToString(momentumDtype); | ||
| 190 | OP_LOGE_FOR_INVALID_DTYPES_WITH_REASON( | ||
| 191 | context->GetNodeName(), "params, grads and momentum_buffer_list", dtypeMsg.c_str(), | ||
| 192 | "params, grads and momentum_buffer_list should have the same dtype"); | ||
| 193 | } | ||
| 194 | return ge::GRAPH_FAILED; | ||
| 195 | } | ||
| 196 | |||
| 197 | bool isInvalidType = (paramsDtype != ge::DT_FLOAT) && (paramsDtype != ge::DT_BF16) && (paramsDtype != ge::DT_FLOAT16); | ||
| 198 | if (isInvalidType) { | ||
| 199 | OP_LOGE_FOR_INVALID_DTYPE(context->GetNodeName(), "params/grads/momentum_buffer_list", | ||
| 200 | Ops::Base::ToString(paramsDtype).c_str(), "float16, bfloat16 and float"); | ||
| 201 | return ge::GRAPH_FAILED; | ||
| 202 | } | ||
| 203 | |||
| 204 | return ge::GRAPH_SUCCESS; | ||
| 205 | } | ||
| 206 | |||
| 207 | ge::graphStatus FusedSgdTiling::GetInputTensorInfo() | ||
| 208 | { | ||
| 209 | auto computeNodeInfo = context_->GetComputeNodeInfo(); | ||
| 210 | OP_CHECK_NULL_WITH_CONTEXT(context_, computeNodeInfo); | ||
| 211 | |||
| 212 | auto anchorInstanceInfo = computeNodeInfo->GetInputInstanceInfo(INPUT_PARAMS_IDX); | ||
| 213 | OP_CHECK_NULL_WITH_CONTEXT(context_, anchorInstanceInfo); | ||
| 214 | tensorNum_ = static_cast<uint64_t>(anchorInstanceInfo->GetInstanceNum()); | ||
| 215 | if (tensorNum_ == 0) { | ||
| 216 | OP_LOGE(context_, "tensor num can not be 0"); | ||
| 217 | return ge::GRAPH_FAILED; | ||
| 218 | } | ||
| 219 | |||
| 220 | // 检查可选输入是否为空 | ||
| 221 | CheckOptionalInputs(); | ||
| 222 | |||
| 223 | for (uint64_t i = 0; i < tensorNum_; i++) { | ||
| 224 | auto paramsShapePtr = context_->GetDynamicInputShape(INPUT_PARAMS_IDX, i); | ||
| 225 | OP_CHECK_NULL_WITH_CONTEXT(context_, paramsShapePtr); | ||
| 226 | auto gradsShapePtr = context_->GetDynamicInputShape(INPUT_GRADS_IDX, i); | ||
| 227 | OP_CHECK_NULL_WITH_CONTEXT(context_, gradsShapePtr); | ||
| 228 | |||
| 229 | gert::Shape paramsShape = paramsShapePtr->GetStorageShape(); | ||
| 230 | gert::Shape gradsShape = gradsShapePtr->GetStorageShape(); | ||
| 231 | bool isDiffSize = paramsShape != gradsShape; | ||
| 232 | |||
| 233 | gert::Shape momentumShape; | ||
| 234 | if(useMomentum_) { | ||
| 235 | auto momentumShapePtr = context_->GetDynamicInputShape(INPUT_MOMENTUM_BUFFER_IDX, i); | ||
| 236 | OP_CHECK_NULL_WITH_CONTEXT(context_, momentumShapePtr); | ||
| 237 | momentumShape = momentumShapePtr->GetStorageShape(); | ||
| 238 | isDiffSize = isDiffSize || paramsShape != momentumShape; | ||
| 239 | } | ||
| 240 | if (isDiffSize) { | ||
| 241 | std::string shapesMsg = Ops::Base::ToString(paramsShape) + ", " + | ||
| 242 | Ops::Base::ToString(gradsShape); | ||
| 243 | if(!useMomentum_) { | ||
| 244 | OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( | ||
| 245 | context_->GetNodeName(), "params, grads", shapesMsg.c_str(), | ||
| 246 | "params, grads should have the same shape"); | ||
| 247 | } else { | ||
| 248 | shapesMsg = shapesMsg + " and " + Ops::Base::ToString(momentumShape); | ||
| 249 | OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON( | ||
| 250 | context_->GetNodeName(), "params, grads and momentum_buffer_list", shapesMsg.c_str(), | ||
| 251 | "params, grads and momentum_buffer_list should have the same shape"); | ||
| 252 | } | ||
| 253 | return ge::GRAPH_FAILED; | ||
| 254 | } | ||
| 255 | } | ||
| 256 | |||
| 257 | return CheckInputDtype(context_, useMomentum_); | ||
| 258 | } | ||
| 259 | |||
| 260 | ge::graphStatus FusedSgdTiling::CalculateOutputInfo() | ||
| 261 | { | ||
| 262 | usedCoreNum_ = tensorNum_ < static_cast<uint64_t>(coreNum_) ? tensorNum_ : static_cast<uint64_t>(coreNum_); | ||
| 263 | tensorsPerCore_ = static_cast<uint32_t>((tensorNum_ + usedCoreNum_ - 1) / usedCoreNum_); | ||
| 264 | |||
| 265 | dtypeSize_ = context_->GetDynamicInputDesc(INPUT_PARAMS_IDX, 0)->GetDataType() == ge::DT_FLOAT ? FP32_DTYPE_SIZE : FP16_BF16_DTYPE_SIZE; | ||
| 266 | uint64_t tBuffersize = BUFFER_NUM * BYTE_ONE_BLK; | ||
| 267 | uint64_t bufferSize = ubSize_ - tBuffersize; | ||
| 268 | // 计算处理一个元素所需的ub大小 | ||
| 269 | uint64_t coreOnesize; | ||
| 270 | if (dtypeSize_ == FP32_DTYPE_SIZE) { | ||
| 271 | coreOnesize = FP32_DTYPE_SIZE * 3 * 2 * BUFFER_NUM; | ||
| 272 | } else { | ||
| 273 | coreOnesize = ((dtypeSize_ + FP32_DTYPE_SIZE) * 3 + FP32_DTYPE_SIZE * 3) * BUFFER_NUM; | ||
| 274 | } | ||
| 275 | uint64_t alignSize = dtypeSize_ == FP32_DTYPE_SIZE ? ONE_BLK_NUM_FP32 : ONE_BLK_NUM; | ||
| 276 | OP_LOGI(context_, "bufferSize = %lu", bufferSize); | ||
| 277 | OP_LOGI(context_, "coreOnesize = %lu", coreOnesize); | ||
| 278 | OP_LOGI(context_, "alignSize = %lu", alignSize); | ||
| 279 | // 计算ub一次最多能处理的数据量 | ||
| 280 | coreCalcMax_ = bufferSize / coreOnesize / alignSize * alignSize; | ||
| 281 | OP_LOGI(context_, "coreCalcMax_ = %lu", coreCalcMax_); | ||
| 282 | |||
| 283 | return ge::GRAPH_SUCCESS; | ||
| 284 | } | ||
| 285 | |||
| 286 | void FusedSgdTiling::SetTilingData(FusedSgdTilingData* tilingData) { | ||
| 287 | tilingData->weightDecay = weightDecay_; | ||
| 288 | tilingData->momentum = momentum_; | ||
| 289 | tilingData->lr = lr_; | ||
| 290 | tilingData->dampening = dampening_; | ||
| 291 | tilingData->nesterov = nesterov_; | ||
| 292 | tilingData->maximize = maximize_; | ||
| 293 | tilingData->isFirstStep = isFirstStep_; | ||
| 294 | tilingData->useGradScale = useGradScale_; | ||
| 295 | tilingData->useMomentum = useMomentum_; | ||
| 296 | tilingData->tensorNum = tensorNum_; | ||
| 297 | tilingData->tensorsPerCore = tensorsPerCore_; | ||
| 298 | tilingData->usedCoreNum = usedCoreNum_; | ||
| 299 | tilingData->coreCalcMax = coreCalcMax_; | ||
| 300 | |||
| 301 | size_t* workspaceSize = context_->GetWorkspaceSizes(1); | ||
| 302 | *workspaceSize = sysWorkspaceSize_; | ||
| 303 | context_->SetTilingKey(0); | ||
| 304 | context_->SetBlockDim(usedCoreNum_); | ||
| 305 | } | ||
| 306 | |||
| 307 | ge::graphStatus Tiling4FusedSgd(gert::TilingContext* context) | ||
| 308 | { | ||
| 309 | OP_LOGD(context, "Tiling4FusedSgd"); | ||
| 310 | FusedSgdTiling tiling(context); | ||
| 311 | OP_CHECK_IF( | ||
| 312 | tiling.GetPlatformInfo() != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetPlatformInfo error"), | ||
| 313 | return ge::GRAPH_FAILED); | ||
| 314 | OP_CHECK_IF( | ||
| 315 | tiling.GetAttrInfo() != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetAttrInfo error"), | ||
| 316 | return ge::GRAPH_FAILED); | ||
| 317 | OP_CHECK_IF( | ||
| 318 | tiling.GetInputTensorInfo() != ge::GRAPH_SUCCESS, OP_LOGE(context, "GetInputTensorInfo error"), | ||
| 319 | return ge::GRAPH_FAILED); | ||
| 320 | OP_CHECK_IF( | ||
| 321 | tiling.CalculateOutputInfo() != ge::GRAPH_SUCCESS, OP_LOGE(context, "CalculateOutputInfo error"), | ||
| 322 | return ge::GRAPH_FAILED); | ||
| 323 | |||
| 324 | FusedSgdTilingData* tilingData = context->GetTilingData<FusedSgdTilingData>(); | ||
| 325 | OP_CHECK_NULL_WITH_CONTEXT(context, tilingData); | ||
| 326 | OP_CHECK_IF( | ||
| 327 | memset_s(tilingData, sizeof(FusedSgdTilingData), 0, sizeof(FusedSgdTilingData)) != EOK, | ||
| 328 | OP_LOGE(context, "set tiling data error"), return ge::GRAPH_FAILED); | ||
| 329 | tiling.SetTilingData(tilingData); | ||
| 330 | OP_LOGD(context, "tiling data: %s", tiling.TilingDataToString().c_str()); | ||
| 331 | return ge::GRAPH_SUCCESS; | ||
| 332 | } | ||
| 333 | |||
| 334 | ge::graphStatus TilingPrepare4FusedSgd([[maybe_unused]] gert::TilingParseContext* context) | ||
| 335 | { | ||
| 336 | return ge::GRAPH_SUCCESS; | ||
| 337 | } | ||
| 338 | |||
| 339 | IMPL_OP_OPTILING(FusedSgd) | ||
| 340 | .Tiling(Tiling4FusedSgd) | ||
| 341 | .TilingParse<FusedSgdCompileInfo>(TilingPrepare4FusedSgd); | ||
| 342 | } // namespace optiling | ||
| @@ -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 | * \file fused_sgd_tiling.h | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | |||
| 23 | namespace optiling { | ||
| 24 | |||
| 25 | struct FusedSgdCompileInfo { | ||
| 26 | }; | ||
| 27 | |||
| 28 | class FusedSgdTiling { | ||
| 29 | public: | ||
| 30 | explicit FusedSgdTiling(gert::TilingContext* context) : context_(context) {}; | ||
| 31 | ge::graphStatus GetPlatformInfo(); | ||
| 32 | ge::graphStatus GetAttrInfo(); | ||
| 33 | ge::graphStatus GetInputTensorInfo(); | ||
| 34 | ge::graphStatus CalculateOutputInfo(); | ||
| 35 | void CheckOptionalInputs(); | ||
| 36 | void SetTilingData(FusedSgdTilingData* tilingData); | ||
| 37 | std::string TilingDataToString() const; | ||
| 38 | |||
| 39 | private: | ||
| 40 | gert::TilingContext* context_; | ||
| 41 | uint32_t coreNum_{0}; | ||
| 42 | uint64_t ubSize_{0}; | ||
| 43 | uint64_t sysWorkspaceSize_{0}; | ||
| 44 | uint32_t usedCoreNum_{0}; | ||
| 45 | float weightDecay_{0.0f}; | ||
| 46 | float momentum_{0.0f}; | ||
| 47 | float lr_{0.0f}; | ||
| 48 | float dampening_{0.0f}; | ||
| 49 | uint32_t nesterov_{0}; | ||
| 50 | uint32_t maximize_{0}; | ||
| 51 | uint32_t isFirstStep_{0}; | ||
| 52 | uint32_t useGradScale_{0}; | ||
| 53 | uint32_t useMomentum_{0}; | ||
| 54 | uint32_t tensorsPerCore_{0}; | ||
| 55 | uint32_t dtypeSize_{0}; | ||
| 56 | uint64_t tensorNum_{0}; | ||
| 57 | uint64_t coreCalcMax_{0}; | ||
| 58 | }; | ||
| 59 | } // namespace optiling | ||
| 60 | |||
| @@ -0,0 +1,86 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd.cpp | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | using namespace AscendC; | ||
| 20 | using namespace FusedSgd; | ||
| 21 | |||
| 22 | |||
| 23 | extern "C" __global__ __aicore__ void fused_sgd( | ||
| 24 | GM_ADDR params, GM_ADDR grads, GM_ADDR momentum_buffer_list, | ||
| 25 | GM_ADDR grad_scale, GM_ADDR params_ref, GM_ADDR grads_ref, | ||
| 26 | GM_ADDR momentum_buffer_list_out, GM_ADDR workspace, | ||
| 27 | GM_ADDR tiling) | ||
| 28 | { | ||
| 29 | KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); | ||
| 30 | AscendC::TPipe pipe; | ||
| 31 | GET_TILING_DATA_WITH_STRUCT(FusedSgdTilingData, tilingData, tiling); | ||
| 32 | |||
| 33 | uint32_t blockIdx = GetBlockIdx(); | ||
| 34 | uint64_t tensorStart = static_cast<uint64_t>(blockIdx) * tilingData.tensorsPerCore; | ||
| 35 | uint64_t tensorEnd = tensorStart + tilingData.tensorsPerCore; | ||
| 36 | if (tensorEnd > tilingData.tensorNum) { | ||
| 37 | tensorEnd = tilingData.tensorNum; | ||
| 38 | } | ||
| 39 | |||
| 40 | FusedSgdF32<DTYPE_X> op(&pipe); | ||
| 41 | op.Init(params, grads, momentum_buffer_list, grad_scale, | ||
| 42 | params_ref, grads_ref, momentum_buffer_list_out, | ||
| 43 | tilingData, tensorStart, tensorEnd); | ||
| 44 | op.Process(); | ||
| 45 | } | ||
| 46 | |||
| 47 | extern "C" __global__ __aicore__ void fused_sgd( | ||
| 48 | GM_ADDR params, GM_ADDR grads, GM_ADDR momentum_buffer_list, | ||
| 49 | GM_ADDR grad_scale, | ||
| 50 | GM_ADDR params_ref, GM_ADDR grads_ref, GM_ADDR momentum_buffer_list_out, | ||
| 51 | GM_ADDR workspace, GM_ADDR tiling) | ||
| 52 | { | ||
| 53 | KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY); | ||
| 54 | AscendC::TPipe pipe; | ||
| 55 | REGISTER_TILING_DEFAULT(FusedSgdTilingData); | ||
| 56 | GET_TILING_DATA_WITH_STRUCT(FusedSgdTilingData, tilingData, tiling); | ||
| 57 | |||
| 58 | // 计算当前核处理的tensor范围 [tensorStart, tensorEnd) | ||
| 59 | uint32_t blockIdx = GetBlockIdx(); | ||
| 60 | uint64_t tensorStart = static_cast<uint64_t>(blockIdx) * tilingData.tensorsPerCore; | ||
| 61 | uint64_t tensorEnd = tensorStart + tilingData.tensorsPerCore; | ||
| 62 | if (tensorEnd > tilingData.tensorNum) { | ||
| 63 | tensorEnd = tilingData.tensorNum; | ||
| 64 | } | ||
| 65 | |||
| 66 | |||
| 67 | FusedSgdF16Bf16<bfloat16_t> op(&pipe); | ||
| 68 | op.Init(params, grads, momentum_buffer_list, grad_scale, | ||
| 69 | params_ref, grads_ref, momentum_buffer_list_out, | ||
| 70 | tilingData, tensorStart, tensorEnd); | ||
| 71 | op.Process(); | ||
| 72 | |||
| 73 | FusedSgdF16Bf16<half> op(&pipe); | ||
| 74 | op.Init(params, grads, momentum_buffer_list, grad_scale, | ||
| 75 | params_ref, grads_ref, momentum_buffer_list_out, | ||
| 76 | tilingData, tensorStart, tensorEnd); | ||
| 77 | op.Process(); | ||
| 78 | |||
| 79 | FusedSgdF32<float> op(&pipe); | ||
| 80 | op.Init(params, grads, momentum_buffer_list, grad_scale, | ||
| 81 | params_ref, grads_ref, momentum_buffer_list_out, | ||
| 82 | tilingData, tensorStart, tensorEnd); | ||
| 83 | op.Process(); | ||
| 84 | |||
| 85 | } | ||
| 86 | |||
| @@ -0,0 +1,82 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd_base.h | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | namespace FusedSgd { | ||
| 23 | using namespace AscendC; | ||
| 24 | constexpr int32_t BYTE_ONE_BLOCK = 32; | ||
| 25 | constexpr int32_t BUFFER_NUM = 2; | ||
| 26 | constexpr int32_t INDEX_PARAMS = 0; | ||
| 27 | constexpr int32_t INDEX_GRADS = 1; | ||
| 28 | constexpr int32_t INDEX_MOMENTUM_BUFFER = 2; | ||
| 29 | |||
| 30 | template <typename T> | ||
| 31 | class FusedSgdBase | ||
| 32 | { | ||
| 33 | public: | ||
| 34 | __aicore__ inline FusedSgdBase(){}; | ||
| 35 | __aicore__ inline void InitData(const FusedSgdTilingData& tiling); | ||
| 36 | __aicore__ inline void PipeSync(); | ||
| 37 | |||
| 38 | protected: | ||
| 39 | float weightDecay; | ||
| 40 | float momentum; | ||
| 41 | float lr; | ||
| 42 | float dampening; | ||
| 43 | uint64_t nesterov; | ||
| 44 | uint64_t maximize; | ||
| 45 | uint64_t isFirstStep; | ||
| 46 | uint64_t useGradScale; | ||
| 47 | uint64_t useMomentum; | ||
| 48 | uint64_t tensorNum; | ||
| 49 | uint64_t tensorsPerCore; | ||
| 50 | uint64_t usedCoreNum; | ||
| 51 | uint64_t coreCalcMax; | ||
| 52 | }; | ||
| 53 | |||
| 54 | template <typename T> | ||
| 55 | __aicore__ inline void FusedSgdBase<T>::InitData(const FusedSgdTilingData& tiling) | ||
| 56 | { | ||
| 57 | weightDecay = tiling.weightDecay; | ||
| 58 | momentum = tiling.momentum; | ||
| 59 | lr = tiling.lr; | ||
| 60 | dampening = tiling.dampening; | ||
| 61 | nesterov = tiling.nesterov; | ||
| 62 | maximize = tiling.maximize; | ||
| 63 | isFirstStep = tiling.isFirstStep; | ||
| 64 | useGradScale = tiling.useGradScale; | ||
| 65 | useMomentum = tiling.useMomentum; | ||
| 66 | tensorNum = tiling.tensorNum; | ||
| 67 | tensorsPerCore = tiling.tensorsPerCore; | ||
| 68 | usedCoreNum = tiling.usedCoreNum; | ||
| 69 | coreCalcMax = tiling.coreCalcMax; | ||
| 70 | } | ||
| 71 | |||
| 72 | template <AscendC::HardEvent hardEvent> | ||
| 73 | __aicore__ inline void PipeSync() | ||
| 74 | { | ||
| 75 | int32_t eventID = static_cast<int32_t>(GetTPipePtr()->FetchEventID(hardEvent)); | ||
| 76 | AscendC::SetFlag<hardEvent>(eventID); | ||
| 77 | AscendC::WaitFlag<hardEvent>(eventID); | ||
| 78 | } | ||
| 79 | |||
| 80 | } // namespace FusedSgd | ||
| 81 | |||
| 82 | |||
| @@ -0,0 +1,248 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd_f16_bf16.h | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | namespace FusedSgd { | ||
| 22 | using namespace AscendC; | ||
| 23 | |||
| 24 | template <typename T> | ||
| 25 | class FusedSgdF16Bf16 : public FusedSgdBase<T> | ||
| 26 | { | ||
| 27 | public: | ||
| 28 | __aicore__ inline FusedSgdF16Bf16(TPipe *pipe) : pipe_(pipe) {}; | ||
| 29 | __aicore__ inline void Init( | ||
| 30 | GM_ADDR params, GM_ADDR grads, GM_ADDR momentum_buffer_list, | ||
| 31 | GM_ADDR grad_scale, GM_ADDR params_ref, GM_ADDR grads_ref, | ||
| 32 | GM_ADDR momentum_buffer_list_out, const FusedSgdTilingData& tiling, | ||
| 33 | uint64_t tensorStart, uint64_t tensorEnd); | ||
| 34 | __aicore__ inline void Process(); | ||
| 35 | |||
| 36 | protected: | ||
| 37 | __aicore__ inline void Compute(const uint64_t index, const uint64_t dataCount); | ||
| 38 | |||
| 39 | TQue<QuePosition::VECIN, BUFFER_NUM> inQue; | ||
| 40 | TQue<QuePosition::VECOUT, BUFFER_NUM> outQue; | ||
| 41 | |||
| 42 | GlobalTensor<T> gmParams; | ||
| 43 | GlobalTensor<T> gmGrads; | ||
| 44 | GlobalTensor<T> gmMomentumBuffer; | ||
| 45 | GlobalTensor<float> gmGradScale; | ||
| 46 | GlobalTensor<float> gmFoundInf; | ||
| 47 | GlobalTensor<T> gmParamsRef; | ||
| 48 | GlobalTensor<T> gmGradsRef; | ||
| 49 | GlobalTensor<T> gmMomentumBufferOut; | ||
| 50 | |||
| 51 | ListTensorDesc paramsList_; | ||
| 52 | ListTensorDesc gradsList_; | ||
| 53 | ListTensorDesc momentumList_; | ||
| 54 | ListTensorDesc paramsRefList_; | ||
| 55 | ListTensorDesc gradsRefList_; | ||
| 56 | ListTensorDesc momentumOutList_; | ||
| 57 | TensorDesc<uint64_t> desc_; | ||
| 58 | |||
| 59 | float gradScaleValue; | ||
| 60 | uint64_t hasGradScale; | ||
| 61 | uint64_t tensorStart_; | ||
| 62 | uint64_t tensorEnd_; | ||
| 63 | int64_t paramsOffset; | ||
| 64 | int64_t gradsOffset; | ||
| 65 | int64_t momentumOffset; | ||
| 66 | int64_t paramsOffsetC; | ||
| 67 | int64_t gradsOffsetC; | ||
| 68 | int64_t momentumOffsetC; | ||
| 69 | TPipe *pipe_; | ||
| 70 | const FusedSgdTilingData* tiling_; | ||
| 71 | }; | ||
| 72 | |||
| 73 | template <typename T> | ||
| 74 | __aicore__ inline void FusedSgdF16Bf16<T>::Init( | ||
| 75 | GM_ADDR params, GM_ADDR grads, GM_ADDR momentum_buffer_list, | ||
| 76 | GM_ADDR grad_scale, GM_ADDR params_ref, GM_ADDR grads_ref, | ||
| 77 | GM_ADDR momentum_buffer_list_out, const FusedSgdTilingData& tiling, | ||
| 78 | uint64_t tensorStart, uint64_t tensorEnd) | ||
| 79 | { | ||
| 80 | this->InitData(tiling); | ||
| 81 | tiling_ = &tiling; | ||
| 82 | tensorStart_ = tensorStart; | ||
| 83 | tensorEnd_ = tensorEnd; | ||
| 84 | |||
| 85 | paramsList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(params)); | ||
| 86 | gradsList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(grads)); | ||
| 87 | paramsRefList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(params_ref)); | ||
| 88 | gradsRefList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(grads_ref)); | ||
| 89 | if (this->useMomentum) { | ||
| 90 | momentumList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(momentum_buffer_list)); | ||
| 91 | momentumOutList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(momentum_buffer_list_out)); | ||
| 92 | } | ||
| 93 | |||
| 94 | // UB Buffer布局: inQue = [原始类型(param+grad+momentum)] + [FP32(param+grad+momentum)] | ||
| 95 | // 前半存原始类型，后半(偏移3*sizeof(T))存Cast后的FP32 | ||
| 96 | pipe_->InitBuffer(inQue, BUFFER_NUM, this->coreCalcMax * (sizeof(T) + sizeof(float)) * 3); | ||
| 97 | pipe_->InitBuffer(outQue, BUFFER_NUM, this->coreCalcMax * sizeof(float) * 3); | ||
| 98 | |||
| 99 | paramsOffset = this->coreCalcMax * INDEX_PARAMS; | ||
| 100 | gradsOffset = this->coreCalcMax * INDEX_GRADS; | ||
| 101 | momentumOffset = this->coreCalcMax * INDEX_MOMENTUM_BUFFER; | ||
| 102 | // FP32区域偏移 = 前半3份原始类型 + 对应的FP32偏移 | ||
| 103 | paramsOffsetC = this->coreCalcMax * 3 + paramsOffset; | ||
| 104 | gradsOffsetC = this->coreCalcMax * 3 + gradsOffset; | ||
| 105 | momentumOffsetC = this->coreCalcMax * 3 + momentumOffset; | ||
| 106 | |||
| 107 | hasGradScale = 0; | ||
| 108 | if (this->useGradScale) { | ||
| 109 | gmGradScale.SetGlobalBuffer((__gm__ float*)grad_scale, 1); | ||
| 110 | gradScaleValue = static_cast<float>(gmGradScale.GetValue(0)); | ||
| 111 | hasGradScale = 1; | ||
| 112 | } | ||
| 113 | } | ||
| 114 | |||
| 115 | template <typename T> | ||
| 116 | __aicore__ inline void FusedSgdF16Bf16<T>::Compute(const uint64_t index, const uint64_t dataCount) | ||
| 117 | { | ||
| 118 | uint64_t offset = index * this->coreCalcMax; | ||
| 119 | DataCopyParams copyParams = {1, static_cast<uint16_t>(dataCount * sizeof(T)), 0, 0}; | ||
| 120 | DataCopyPadParams padParams = {false, 0, 0, 0}; | ||
| 121 | |||
| 122 | LocalTensor<T> inLocal = inQue.AllocTensor<T>(); | ||
| 123 | LocalTensor<float> outLocal = outQue.AllocTensor<float>(); | ||
| 124 | |||
| 125 | PipeSync<AscendC::HardEvent::MTE3_MTE2>(); | ||
| 126 | PipeSync<AscendC::HardEvent::S_MTE2>(); | ||
| 127 | PipeSync<AscendC::HardEvent::V_MTE2>(); | ||
| 128 | DataCopyPad(inLocal[paramsOffset], gmParams[offset], copyParams, padParams); | ||
| 129 | DataCopyPad(inLocal[gradsOffset], gmGrads[offset], copyParams, padParams); | ||
| 130 | if (this->useMomentum) { | ||
| 131 | DataCopyPad(inLocal[momentumOffset], gmMomentumBuffer[offset], copyParams, padParams); | ||
| 132 | } | ||
| 133 | PipeSync<AscendC::HardEvent::MTE2_V>(); | ||
| 134 | PipeBarrier<PIPE_V>(); | ||
| 135 | |||
| 136 | LocalTensor<float> inLocalC = inLocal[this->coreCalcMax * 3].template ReinterpretCast<float>(); | ||
| 137 | Cast(inLocalC[paramsOffset], inLocal[paramsOffset], RoundMode::CAST_NONE, dataCount); | ||
| 138 | PipeBarrier<PIPE_V>(); | ||
| 139 | Cast(inLocalC[gradsOffset], inLocal[gradsOffset], RoundMode::CAST_NONE, dataCount); | ||
| 140 | PipeBarrier<PIPE_V>(); | ||
| 141 | Cast(inLocalC[momentumOffset], inLocal[momentumOffset], RoundMode::CAST_NONE, dataCount); | ||
🔵 Low Priority 在 然而，当 功能上无实质影响（结果未被消费），但：
建议添加 变更：第141行新增的无条件 Cast。 | |||
| 142 | PipeBarrier<PIPE_V>(); | ||
| 143 | |||
| 144 | // Step 1: 梯度缩放，并Cast回原始类型写回 | ||
| 145 | if (hasGradScale) { | ||
| 146 | float invGradScale = 1.0f / gradScaleValue; | ||
| 147 | Muls(inLocalC[gradsOffset], inLocalC[gradsOffset], invGradScale, dataCount); | ||
| 148 | PipeBarrier<PIPE_V>(); | ||
| 149 | Cast(inLocal[gradsOffset], inLocalC[gradsOffset], RoundMode::CAST_RINT, dataCount); | ||
| 150 | PipeSync<AscendC::HardEvent::V_MTE3>(); | ||
| 151 | DataCopyPad(gmGradsRef[offset], inLocal[gradsOffset], copyParams); | ||
| 152 | PipeSync<AscendC::HardEvent::MTE3_V>(); | ||
| 153 | } | ||
| 154 | // Step 2: 最大化处理 | ||
| 155 | if (this->maximize) { | ||
| 156 | Muls(inLocalC[gradsOffset], inLocalC[gradsOffset], -1.0f, dataCount); | ||
| 157 | PipeBarrier<PIPE_V>(); | ||
| 158 | } | ||
| 159 | // Step 3: 权重衰减 | ||
| 160 | if (this->weightDecay != 0.0f) { | ||
| 161 | Muls(outLocal[gradsOffset], inLocalC[paramsOffset], this->weightDecay, dataCount); | ||
| 162 | PipeBarrier<PIPE_V>(); | ||
| 163 | Add(inLocalC[gradsOffset], inLocalC[gradsOffset], outLocal[gradsOffset], dataCount); | ||
| 164 | PipeBarrier<PIPE_V>(); | ||
| 165 | } | ||
| 166 | |||
| 167 | // Step 4: 动量更新 (FP32计算，写回时Cast回原始类型) | ||
| 168 | if (this->useMomentum) { | ||
| 169 | if (this->isFirstStep) { | ||
| 170 | Muls(outLocal[momentumOffset], inLocalC[gradsOffset], 1.0f, dataCount); | ||
| 171 | PipeBarrier<PIPE_V>(); | ||
| 172 | } else { | ||
| 173 | Muls(outLocal[momentumOffset], inLocalC[momentumOffset], this->momentum, dataCount); | ||
| 174 | PipeBarrier<PIPE_V>(); | ||
| 175 | Muls(outLocal[paramsOffset], inLocalC[gradsOffset], 1.0f - this->dampening, dataCount); | ||
| 176 | PipeBarrier<PIPE_V>(); | ||
| 177 | Add(outLocal[momentumOffset], outLocal[momentumOffset], outLocal[paramsOffset], dataCount); | ||
| 178 | PipeBarrier<PIPE_V>(); | ||
| 179 | } | ||
| 180 | // 动量Cast回原始类型并写回GM | ||
| 181 | Cast(inLocal[momentumOffset], outLocal[momentumOffset], RoundMode::CAST_RINT, dataCount); | ||
| 182 | PipeBarrier<PIPE_V>(); | ||
| 183 | PipeSync<AscendC::HardEvent::V_MTE3>(); | ||
| 184 | DataCopyPad(gmMomentumBufferOut[offset], inLocal[momentumOffset], copyParams); | ||
| 185 | PipeSync<AscendC::HardEvent::MTE3_V>(); | ||
| 186 | // Nesterov: grad = grad + momentum * buf | ||
| 187 | if (this->nesterov) { | ||
| 188 | Muls(outLocal[momentumOffset], outLocal[momentumOffset], this->momentum, dataCount); | ||
| 189 | PipeBarrier<PIPE_V>(); | ||
| 190 | Add(inLocalC[gradsOffset], outLocal[momentumOffset], inLocalC[gradsOffset], dataCount); | ||
| 191 | PipeBarrier<PIPE_V>(); | ||
| 192 | } else { | ||
| 193 | Muls(inLocalC[gradsOffset], outLocal[momentumOffset], 1.0f, dataCount); | ||
| 194 | PipeBarrier<PIPE_V>(); | ||
| 195 | } | ||
| 196 | } | ||
| 197 | |||
| 198 | // Step 5: 参数更新 (param = param - lr * grad)，Cast回原始类型写回 | ||
| 199 | Muls(inLocalC[gradsOffset], inLocalC[gradsOffset], this->lr, dataCount); | ||
| 200 | PipeBarrier<PIPE_V>(); | ||
| 201 | Sub(inLocalC[gradsOffset], inLocalC[paramsOffset], inLocalC[gradsOffset], dataCount); | ||
| 202 | PipeBarrier<PIPE_V>(); | ||
| 203 | Cast(inLocal[paramsOffset], inLocalC[gradsOffset], RoundMode::CAST_RINT, dataCount); | ||
| 204 | PipeBarrier<PIPE_V>(); | ||
| 205 | PipeSync<AscendC::HardEvent::V_MTE3>(); | ||
| 206 | DataCopyPad(gmParamsRef[offset], inLocal[paramsOffset], copyParams); | ||
| 207 | |||
| 208 | inQue.FreeTensor(inLocal); | ||
| 209 | outQue.FreeTensor(outLocal); | ||
| 210 | } | ||
| 211 | |||
| 212 | template <typename T> | ||
| 213 | __aicore__ inline void FusedSgdF16Bf16<T>::Process() | ||
| 214 | { | ||
| 215 | for (uint64_t idx = tensorStart_; idx < tensorEnd_; idx++) { | ||
| 216 | uint64_t buf[10]; | ||
| 217 | desc_.SetShapeAddr(buf); | ||
| 218 | paramsList_.GetDesc(desc_, static_cast<uint32_t>(idx)); | ||
| 219 | |||
| 220 | uint64_t tensorDataNum = 1; | ||
| 221 | for (uint32_t j = 0; j < desc_.GetDim(); j++) { | ||
| 222 | tensorDataNum *= desc_.GetShape(j); | ||
| 223 | } | ||
| 224 | if (tensorDataNum == 0) { | ||
| 225 | continue; | ||
| 226 | } | ||
| 227 | |||
| 228 | gmParams.SetGlobalBuffer(paramsList_.GetDataPtr<T>(idx), tensorDataNum); | ||
| 229 | gmGrads.SetGlobalBuffer(gradsList_.GetDataPtr<T>(idx), tensorDataNum); | ||
| 230 | gmParamsRef.SetGlobalBuffer(paramsRefList_.GetDataPtr<T>(idx), tensorDataNum); | ||
| 231 | gmGradsRef.SetGlobalBuffer(gradsRefList_.GetDataPtr<T>(idx), tensorDataNum); | ||
| 232 | if (this->useMomentum) { | ||
| 233 | gmMomentumBuffer.SetGlobalBuffer(momentumList_.GetDataPtr<T>(idx), tensorDataNum); | ||
| 234 | gmMomentumBufferOut.SetGlobalBuffer(momentumOutList_.GetDataPtr<T>(idx), tensorDataNum); | ||
| 235 | } | ||
| 236 | |||
| 237 | uint64_t loopNum = (tensorDataNum + this->coreCalcMax - 1) / this->coreCalcMax; | ||
| 238 | for (uint64_t n = 0; n < loopNum - 1; n++) { | ||
| 239 | Compute(n, this->coreCalcMax); | ||
| 240 | } | ||
| 241 | uint64_t lastCount = tensorDataNum - this->coreCalcMax * (loopNum - 1); | ||
| 242 | Compute(loopNum - 1, lastCount); | ||
| 243 | } | ||
| 244 | } | ||
| 245 | |||
| 246 | } // namespace FusedSgd | ||
| 247 | |||
| 248 | |||
| @@ -0,0 +1,226 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd_f32.h | ||
| 13 | * \brief | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | namespace FusedSgd { | ||
| 22 | using namespace AscendC; | ||
| 23 | |||
| 24 | template <typename T> | ||
| 25 | class FusedSgdF32 : public FusedSgdBase<T> | ||
| 26 | { | ||
| 27 | public: | ||
| 28 | __aicore__ inline FusedSgdF32(TPipe *pipe) : pipe_(pipe) {}; | ||
| 29 | __aicore__ inline void Init( | ||
| 30 | GM_ADDR params, GM_ADDR grads, GM_ADDR momentum_buffer_list, | ||
| 31 | GM_ADDR grad_scale, GM_ADDR params_ref, GM_ADDR grads_ref, | ||
| 32 | GM_ADDR momentum_buffer_list_out, const FusedSgdTilingData& tiling, | ||
| 33 | uint64_t tensorStart, uint64_t tensorEnd); | ||
| 34 | __aicore__ inline void Process(); | ||
| 35 | |||
| 36 | protected: | ||
| 37 | __aicore__ inline void Compute(const uint64_t index, const uint64_t dataCount); | ||
| 38 | |||
| 39 | TQue<QuePosition::VECIN, BUFFER_NUM> inQue; | ||
| 40 | TQue<QuePosition::VECOUT, BUFFER_NUM> outQue; | ||
| 41 | |||
| 42 | GlobalTensor<T> gmParams; | ||
| 43 | GlobalTensor<T> gmGrads; | ||
| 44 | GlobalTensor<T> gmMomentumBuffer; | ||
| 45 | GlobalTensor<float> gmGradScale; | ||
| 46 | GlobalTensor<float> gmFoundInf; | ||
| 47 | GlobalTensor<T> gmParamsRef; | ||
| 48 | GlobalTensor<T> gmGradsRef; | ||
| 49 | GlobalTensor<T> gmMomentumBufferOut; | ||
| 50 | |||
| 51 | ListTensorDesc paramsList_; | ||
| 52 | ListTensorDesc gradsList_; | ||
| 53 | ListTensorDesc momentumList_; | ||
| 54 | ListTensorDesc paramsRefList_; | ||
| 55 | ListTensorDesc gradsRefList_; | ||
| 56 | ListTensorDesc momentumOutList_; | ||
| 57 | TensorDesc<uint64_t> desc_; | ||
| 58 | |||
| 59 | float gradScaleValue; | ||
| 60 | uint64_t hasGradScale; | ||
| 61 | uint64_t tensorStart_; | ||
| 62 | uint64_t tensorEnd_; | ||
| 63 | int64_t paramsOffset; | ||
| 64 | int64_t gradsOffset; | ||
| 65 | int64_t momentumOffset; | ||
| 66 | TPipe *pipe_; | ||
| 67 | const FusedSgdTilingData* tiling_; | ||
| 68 | }; | ||
| 69 | |||
| 70 | template <typename T> | ||
| 71 | __aicore__ inline void FusedSgdF32<T>::Init( | ||
| 72 | GM_ADDR params, GM_ADDR grads, GM_ADDR momentum_buffer_list, | ||
| 73 | GM_ADDR grad_scale, GM_ADDR params_ref, GM_ADDR grads_ref, | ||
| 74 | GM_ADDR momentum_buffer_list_out, const FusedSgdTilingData& tiling, | ||
| 75 | uint64_t tensorStart, uint64_t tensorEnd) | ||
| 76 | { | ||
| 77 | this->InitData(tiling); | ||
| 78 | tiling_ = &tiling; | ||
| 79 | tensorStart_ = tensorStart; | ||
| 80 | tensorEnd_ = tensorEnd; | ||
| 81 | |||
| 82 | paramsList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(params)); | ||
| 83 | gradsList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(grads)); | ||
| 84 | paramsRefList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(params_ref)); | ||
| 85 | gradsRefList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(grads_ref)); | ||
| 86 | if (this->useMomentum) { | ||
| 87 | momentumList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(momentum_buffer_list)); | ||
| 88 | momentumOutList_ = ListTensorDesc(reinterpret_cast<__gm__ void*>(momentum_buffer_list_out)); | ||
| 89 | } | ||
| 90 | |||
| 91 | pipe_->InitBuffer(inQue, BUFFER_NUM, this->coreCalcMax * sizeof(T) * 3); | ||
| 92 | pipe_->InitBuffer(outQue, BUFFER_NUM, this->coreCalcMax * sizeof(T) * 3); | ||
| 93 | |||
| 94 | paramsOffset = this->coreCalcMax * INDEX_PARAMS; | ||
| 95 | gradsOffset = this->coreCalcMax * INDEX_GRADS; | ||
| 96 | momentumOffset = this->coreCalcMax * INDEX_MOMENTUM_BUFFER; | ||
| 97 | |||
| 98 | hasGradScale = 0; | ||
| 99 | if (this->useGradScale) { | ||
| 100 | gmGradScale.SetGlobalBuffer((__gm__ float*)grad_scale, 1); | ||
| 101 | gradScaleValue = static_cast<float>(gmGradScale.GetValue(0)); | ||
| 102 | hasGradScale = 1; | ||
| 103 | } | ||
| 104 | } | ||
| 105 | |||
| 106 | template <typename T> | ||
| 107 | __aicore__ inline void FusedSgdF32<T>::Compute(const uint64_t index, const uint64_t dataCount) | ||
| 108 | { | ||
| 109 | uint64_t offset = index * this->coreCalcMax; | ||
| 110 | DataCopyParams copyParams = {1, static_cast<uint16_t>(dataCount * sizeof(T)), 0, 0}; | ||
| 111 | DataCopyPadParams padParams = {false, 0, 0, 0}; | ||
| 112 | |||
| 113 | LocalTensor<T> inLocal = inQue.AllocTensor<T>(); | ||
| 114 | LocalTensor<T> outLocal = outQue.AllocTensor<T>(); | ||
| 115 | |||
| 116 | PipeSync<AscendC::HardEvent::MTE3_MTE2>(); | ||
| 117 | PipeSync<AscendC::HardEvent::S_MTE2>(); | ||
| 118 | PipeSync<AscendC::HardEvent::V_MTE2>(); | ||
| 119 | DataCopyPad(inLocal[paramsOffset], gmParams[offset], copyParams, padParams); | ||
| 120 | DataCopyPad(inLocal[gradsOffset], gmGrads[offset], copyParams, padParams); | ||
| 121 | if (this->useMomentum) { | ||
| 122 | DataCopyPad(inLocal[momentumOffset], gmMomentumBuffer[offset], copyParams, padParams); | ||
| 123 | } | ||
| 124 | PipeSync<AscendC::HardEvent::MTE2_V>(); | ||
| 125 | PipeBarrier<PIPE_V>(); | ||
| 126 | |||
| 127 | // Step 1: 梯度缩放 (grad = grad / grad_scale)，并回写缩放后的梯度 | ||
| 128 | if (hasGradScale) { | ||
| 129 | float invGradScale = 1.0f / gradScaleValue; | ||
| 130 | Muls(inLocal[gradsOffset], inLocal[gradsOffset], invGradScale, dataCount); | ||
| 131 | PipeBarrier<PIPE_V>(); | ||
| 132 | PipeSync<AscendC::HardEvent::V_MTE3>(); | ||
| 133 | DataCopyPad(gmGradsRef[offset], inLocal[gradsOffset], copyParams); | ||
| 134 | PipeSync<AscendC::HardEvent::MTE3_V>(); | ||
| 135 | } | ||
| 136 | // Step 2: 最大化处理 (grad = -grad)，用于最大化目标函数 | ||
| 137 | if (this->maximize) { | ||
| 138 | Muls(inLocal[gradsOffset], inLocal[gradsOffset], -1.0f, dataCount); | ||
| 139 | PipeBarrier<PIPE_V>(); | ||
| 140 | } | ||
| 141 | // Step 3: 权重衰减 (grad = grad + weight_decay * param) | ||
| 142 | if (this->weightDecay != 0.0f) { | ||
| 143 | Muls(outLocal[gradsOffset], inLocal[paramsOffset], this->weightDecay, dataCount); | ||
| 144 | PipeBarrier<PIPE_V>(); | ||
| 145 | Add(inLocal[gradsOffset], inLocal[gradsOffset], outLocal[gradsOffset], dataCount); | ||
| 146 | PipeBarrier<PIPE_V>(); | ||
| 147 | } | ||
| 148 | |||
| 149 | // Step 4: 动量更新 (SGD with Momentum) | ||
| 150 | if (this->useMomentum) { | ||
| 151 | if (this->isFirstStep) { | ||
| 152 | Muls(outLocal[momentumOffset], inLocal[gradsOffset], 1.0f, dataCount); | ||
| 153 | PipeBarrier<PIPE_V>(); | ||
| 154 | } else { | ||
| 155 | Muls(outLocal[momentumOffset], inLocal[momentumOffset], this->momentum, dataCount); | ||
| 156 | PipeBarrier<PIPE_V>(); | ||
| 157 | Muls(outLocal[paramsOffset], inLocal[gradsOffset], 1.0f - this->dampening, dataCount); | ||
| 158 | PipeBarrier<PIPE_V>(); | ||
| 159 | Add(outLocal[momentumOffset], outLocal[momentumOffset], outLocal[paramsOffset], dataCount); | ||
| 160 | PipeBarrier<PIPE_V>(); | ||
| 161 | } | ||
| 162 | // 动量写回GM | ||
| 163 | PipeSync<AscendC::HardEvent::V_MTE3>(); | ||
| 164 | DataCopyPad(gmMomentumBufferOut[offset], outLocal[momentumOffset], copyParams); | ||
| 165 | PipeSync<AscendC::HardEvent::MTE3_V>(); | ||
| 166 | // Nesterov: grad = grad + momentum * buf | ||
| 167 | if (this->nesterov) { | ||
| 168 | Muls(outLocal[momentumOffset], outLocal[momentumOffset], this->momentum, dataCount); | ||
| 169 | PipeBarrier<PIPE_V>(); | ||
| 170 | Add(inLocal[gradsOffset], outLocal[momentumOffset], inLocal[gradsOffset], dataCount); | ||
| 171 | PipeBarrier<PIPE_V>(); | ||
| 172 | } else { | ||
| 173 | Muls(inLocal[gradsOffset], outLocal[momentumOffset], 1.0f, dataCount); | ||
| 174 | PipeBarrier<PIPE_V>(); | ||
| 175 | } | ||
| 176 | } | ||
| 177 | |||
| 178 | // Step 5: 参数更新 (param = param - lr * grad)，并写回GM | ||
| 179 | Muls(outLocal[gradsOffset], inLocal[gradsOffset], this->lr, dataCount); | ||
| 180 | PipeBarrier<PIPE_V>(); | ||
| 181 | Sub(outLocal[paramsOffset], inLocal[paramsOffset], outLocal[gradsOffset], dataCount); | ||
| 182 | PipeBarrier<PIPE_V>(); | ||
| 183 | PipeSync<AscendC::HardEvent::V_MTE3>(); | ||
| 184 | DataCopyPad(gmParamsRef[offset], outLocal[paramsOffset], copyParams); | ||
| 185 | |||
| 186 | inQue.FreeTensor(inLocal); | ||
| 187 | outQue.FreeTensor(outLocal); | ||
| 188 | } | ||
| 189 | |||
| 190 | template <typename T> | ||
| 191 | __aicore__ inline void FusedSgdF32<T>::Process() | ||
| 192 | { | ||
| 193 | for (uint64_t idx = tensorStart_; idx < tensorEnd_; idx++) { | ||
| 194 | uint64_t buf[10]; | ||
| 195 | desc_.SetShapeAddr(buf); | ||
| 196 | paramsList_.GetDesc(desc_, static_cast<uint32_t>(idx)); | ||
| 197 | |||
| 198 | uint64_t tensorDataNum = 1; | ||
| 199 | for (uint32_t j = 0; j < desc_.GetDim(); j++) { | ||
| 200 | tensorDataNum *= desc_.GetShape(j); | ||
| 201 | } | ||
| 202 | if (tensorDataNum == 0) { | ||
| 203 | continue; | ||
| 204 | } | ||
| 205 | |||
| 206 | gmParams.SetGlobalBuffer(paramsList_.GetDataPtr<T>(idx)); | ||
| 207 | gmGrads.SetGlobalBuffer(gradsList_.GetDataPtr<T>(idx)); | ||
| 208 | gmParamsRef.SetGlobalBuffer(paramsRefList_.GetDataPtr<T>(idx)); | ||
| 209 | gmGradsRef.SetGlobalBuffer(gradsRefList_.GetDataPtr<T>(idx)); | ||
| 210 | if (this->useMomentum) { | ||
| 211 | gmMomentumBuffer.SetGlobalBuffer(momentumList_.GetDataPtr<T>(idx)); | ||
| 212 | gmMomentumBufferOut.SetGlobalBuffer(momentumOutList_.GetDataPtr<T>(idx)); | ||
🟠 High Priority 在 对比 F16/BF16 内核（ 在 AscendC 框架中， 变更：第206-212行新增， 建议：在 F32 内核的 majiajian_hw 6月22日 评论： 6月22日 评论： | |||
| 213 | } | ||
| 214 | |||
| 215 | uint64_t loopNum = (tensorDataNum + this->coreCalcMax - 1) / this->coreCalcMax; | ||
| 216 | for (uint64_t n = 0; n < loopNum - 1; n++) { | ||
| 217 | Compute(n, this->coreCalcMax); | ||
| 218 | } | ||
| 219 | uint64_t lastCount = tensorDataNum - this->coreCalcMax * (loopNum - 1); | ||
| 220 | Compute(loopNum - 1, lastCount); | ||
| 221 | } | ||
| 222 | } | ||
| 223 | |||
| 224 | } // namespace FusedSgd | ||
| 225 | |||
| 226 | |||
| @@ -0,0 +1,34 @@ | |||
| 1 | /** | ||
| 2 | * Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | * This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | * CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | * Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | * See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | */ | ||
| 10 | |||
| 11 | /*! | ||
| 12 | * \file fused_sgd_tiling_data.h | ||
| 13 | * \brief tiling data struct | ||
| 14 | */ | ||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | struct FusedSgdTilingData { | ||
| 20 | float weightDecay; | ||
| 21 | float momentum; | ||
| 22 | float lr; | ||
| 23 | float dampening; | ||
| 24 | uint64_t nesterov; | ||
| 25 | uint64_t maximize; | ||
| 26 | uint64_t isFirstStep; | ||
| 27 | uint64_t useGradScale; | ||
| 28 | uint64_t useMomentum; | ||
| 29 | uint64_t tensorNum; | ||
| 30 | uint64_t tensorsPerCore; | ||
| 31 | uint64_t usedCoreNum; | ||
| 32 | uint64_t coreCalcMax; | ||
| 33 | }; | ||
| 34 | |||
| @@ -0,0 +1,18 @@ | |||
| 1 | # ---------------------------------------------------------------------------- | ||
| 2 | # Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | # This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | # CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | # Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | # See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | # ---------------------------------------------------------------------------- | ||
| 10 | |||
| 11 | message(STATUS "=== Debug: start ops.optim.fused_sgd.tests.CMakeLists.txt ") | ||
| 12 | file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*) | ||
| 13 | message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ") | ||
| 14 | foreach(SUB_DIR ${CURRENT_SOURCE_DIRS}) | ||
| 15 | if(EXISTS "${SUB_DIR}/CMakeLists.txt") | ||
| 16 | add_subdirectory(${SUB_DIR}) | ||
| 17 | endif() | ||
| 18 | endforeach() | ||
| @@ -0,0 +1,17 @@ | |||
| 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 | file(GLOB CURRENT_SOURCE_DIRS LIST_DIRECTORIES true ${CMAKE_CURRENT_SOURCE_DIR}/*) | ||
| 12 | message(STATUS "=== Debug: CURRENT_SOURCE_DIRS =${CURRENT_SOURCE_DIRS} ") | ||
| 13 | foreach(SUB_DIR ${CURRENT_SOURCE_DIRS}) | ||
| 14 | if(EXISTS "${SUB_DIR}/CMakeLists.txt") | ||
| 15 | add_subdirectory(${SUB_DIR}) | ||
| 16 | endif() | ||
| 17 | endforeach() | ||
| @@ -0,0 +1,16 @@ | |||
| 1 | # ---------------------------------------------------------------------------- | ||
| 2 | # Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 3 | # This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 4 | # CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 5 | # Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 6 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 7 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 8 | # See LICENSE in the root of the software repository for the full text of the License. | ||
| 9 | # ---------------------------------------------------------------------------- | ||
| 10 | |||
| 11 | message(STATUS "=== Debug: CURRENT_DIRS =${CURRENT_DIRS} ") | ||
| 12 | file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*) | ||
| 13 | if(UT_TEST_ALL OR OP_HOST_UT) | ||
| 14 | add_modules_ut_sources(HOSTNAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR}) | ||
| 15 | add_modules_ut_sources(HOSTNAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR}) | ||
| 16 | endif() | ||
| @@ -0,0 +1,84 @@ | |||
| 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 | |||
| 13 | |||
| 14 | |||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | |||
| 23 | |||
| 24 | |||
| 25 | |||
| 26 | |||
| 27 | |||
| 28 | |||
| 29 | |||
| 30 | |||
| 31 | using namespace ge; | ||
| 32 | using namespace ut_util; | ||
| 33 | using namespace std; | ||
| 34 | |||
| 35 | class FusedSgdInferShape : public testing::Test | ||
| 36 | { | ||
| 37 | protected: | ||
| 38 | static void SetUpTestCase() | ||
| 39 | { | ||
| 40 | std::cout << "FusedSgdInferShape SetUp" << std::endl; | ||
| 41 | } | ||
| 42 | static void TearDownTestCase() | ||
| 43 | { | ||
| 44 | std::cout << "FusedSgdInferShape TearDown" << std::endl; | ||
| 45 | } | ||
| 46 | }; | ||
| 47 | |||
| 48 | std::vector<int64_t> ToVectorForFused(const gert::Shape& shape) { | ||
| 49 | size_t shape_size = shape.GetDimNum(); | ||
| 50 | std::vector<int64_t> shape_vec(shape_size, 0); | ||
| 51 | |||
| 52 | for (size_t i = 0; i < shape_size; i++) { | ||
| 53 | shape_vec[i] = shape.GetDim(i); | ||
| 54 | } | ||
| 55 | return shape_vec; | ||
| 56 | } | ||
| 57 | |||
| 58 | TEST_F(FusedSgdInferShape, test_fused_sgd_infershape_same_shape) | ||
| 59 | { | ||
| 60 | gert::StorageShape paramsShape = {{3, 4, 5}, {3, 6}}; | ||
| 61 | gert::StorageShape gradsShape = {{3, 4, 5}, {3, 6}}; | ||
| 62 | gert::StorageShape momentumShape = {{3, 4, 5}, {3, 6}}; | ||
| 63 | std::string opType("FusedSgd"); | ||
| 64 | auto infershape_func = gert::OpImplRegistry::GetInstance().GetOpImpl(opType)->infer_shape; | ||
| 65 | |||
| 66 | gert::StorageShape paramsRefShape; | ||
| 67 | gert::StorageShape gradsRefShape; | ||
| 68 | gert::StorageShape momentumRefShape; | ||
| 69 | auto holder = gert::InferShapeContextFaker() | ||
| 70 | .NodeIoNum(5, 3) | ||
| 71 | .IrInstanceNum({1,1,1,1,1}) | ||
| 72 | .InputShapes({¶msShape, &gradsShape, &momentumShape}) | ||
| 73 | .OutputShapes({¶msShape, &gradsShape, &momentumShape}) | ||
| 74 | .Build(); | ||
| 75 | |||
| 76 | gert::InferShapeContext* context = holder.GetContext<gert::InferShapeContext>(); | ||
| 77 | EXPECT_EQ(infershape_func(context), ge::GRAPH_SUCCESS); | ||
| 78 | |||
| 79 | std::vector<int64_t> expectedOutputShape1 = {3, 4, 5}; | ||
| 80 | for(int i=0;i<3;i++){ | ||
| 81 | auto tmpOutShape1 = context->GetOutputShape(i); | ||
| 82 | EXPECT_EQ(ToVectorForFused(*tmpOutShape1), expectedOutputShape1); | ||
| 83 | } | ||
| 84 | } | ||
| @@ -0,0 +1,232 @@ | |||
| 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 | |||
| 13 | |||
| 14 | |||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | |||
| 23 | |||
| 24 | |||
| 25 | |||
| 26 | using namespace ut_util; | ||
| 27 | using namespace std; | ||
| 28 | using namespace ge; | ||
| 29 | |||
| 30 | class FusedSgdTiling : public testing::Test | ||
| 31 | { | ||
| 32 | protected: | ||
| 33 | static void SetUpTestCase() | ||
| 34 | { | ||
| 35 | std::cout << "FusedSgdTiling SetUp" << std::endl; | ||
| 36 | } | ||
| 37 | |||
| 38 | static void TearDownTestCase() | ||
| 39 | { | ||
| 40 | std::cout << "FusedSgdTiling TearDown" << std::endl; | ||
| 41 | } | ||
| 42 | }; | ||
| 43 | |||
| 44 | TEST_F(FusedSgdTiling, test_tiling_float32) | ||
| 45 | { | ||
| 46 | gert::StorageShape paramsShape = {{300, 4, 2}, {300, 4, 2}}; | ||
| 47 | gert::StorageShape gradsShape = {{300, 4, 2}, {300, 4, 2}}; | ||
| 48 | gert::StorageShape momentumShape = {{300, 4, 2}, {300, 4, 2}}; | ||
| 49 | gert::StorageShape gradScaleShape = {{1}, {1}}; | ||
| 50 | |||
| 51 | string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1", | ||
| 52 | "Intrinsic_fix_pipe_l0c2out": false, | ||
| 53 | "Intrinsic_data_move_l12ub": true, | ||
| 54 | "Intrinsic_data_move_l0c2ub": true, | ||
| 55 | "Intrinsic_data_move_out2l1_nd2nz": false, | ||
| 56 | "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288, | ||
| 57 | "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072, | ||
| 58 | "CORE_NUM": 48} | ||
| 59 | })"; | ||
| 60 | map<string, string> soc_infos; | ||
| 61 | map<string, string> aicore_spec; | ||
| 62 | map<string, string> intrinsics; | ||
| 63 | |||
| 64 | GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics); | ||
| 65 | |||
| 66 | fe::PlatFormInfos platform_info; | ||
| 67 | platform_info.Init(); | ||
| 68 | |||
| 69 | struct FusedSgdCompileInfo { | ||
| 70 | }; | ||
| 71 | FusedSgdCompileInfo compile_info; | ||
| 72 | |||
| 73 | std::string op_type("FusedSgd"); | ||
| 74 | ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr); | ||
| 75 | auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling; | ||
| 76 | auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse; | ||
| 77 | |||
| 78 | auto kernel_holder = | ||
| 79 | gert::KernelRunContextFaker() | ||
| 80 | .KernelIONum(2, 1) | ||
| 81 | .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)}) | ||
| 82 | .Outputs({&compile_info}) | ||
| 83 | .Build(); | ||
| 84 | |||
| 85 | ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init()); | ||
| 86 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos); | ||
| 87 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec); | ||
| 88 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore"); | ||
| 89 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes( | ||
| 90 | "AICoreintrinsicDtypeMap", intrinsics); | ||
| 91 | |||
| 92 | ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS); | ||
| 93 | |||
| 94 | auto param = gert::TilingData::CreateCap(4096); | ||
| 95 | auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096); | ||
| 96 | auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get()); | ||
| 97 | ASSERT_NE(param, nullptr); | ||
| 98 | auto holder = gert::TilingContextFaker() | ||
| 99 | .NodeIoNum(5, 3) | ||
| 100 | .IrInstanceNum({1, 1, 1, 1, 1}) | ||
| 101 | .InputShapes({¶msShape, &gradsShape, &momentumShape, &gradScaleShape}) | ||
| 102 | .OutputShapes({¶msShape, &gradsShape, &momentumShape}) | ||
| 103 | .CompileInfo(&compile_info) | ||
| 104 | .PlatformInfo(reinterpret_cast<char*>(&platform_info)) | ||
| 105 | .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 106 | .NodeInputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 107 | .NodeInputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 108 | .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 109 | .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 110 | .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 111 | .NodeOutputTd(1, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 112 | .NodeOutputTd(2, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 113 | .NodeAttrs( | ||
| 114 | {{"weight_decay", Ops::NN::AnyValue::CreateFrom(0.01f)}, | ||
| 115 | {"momentum", Ops::NN::AnyValue::CreateFrom(0.9f)}, | ||
| 116 | {"lr", Ops::NN::AnyValue::CreateFrom(0.001f)}, | ||
| 117 | {"dampening", Ops::NN::AnyValue::CreateFrom(0.0f)}, | ||
| 118 | {"nesterov", Ops::NN::AnyValue::CreateFrom<bool>(false)}, | ||
| 119 | {"maximize", Ops::NN::AnyValue::CreateFrom<bool>(false)}, | ||
| 120 | {"is_first_step", Ops::NN::AnyValue::CreateFrom<bool>(false)}}) | ||
| 121 | .TilingData(param.get()) | ||
| 122 | .Workspace(ws_size) | ||
| 123 | .Build(); | ||
| 124 | |||
| 125 | gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>(); | ||
| 126 | ASSERT_NE(tiling_context, nullptr); | ||
| 127 | ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr); | ||
| 128 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos); | ||
| 129 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec); | ||
| 130 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore"); | ||
| 131 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics); | ||
| 132 | |||
| 133 | EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS); | ||
| 134 | |||
| 135 | auto tiling_key = tiling_context->GetTilingKey(); | ||
| 136 | ASSERT_EQ(tiling_key, 0); | ||
| 137 | } | ||
| 138 | |||
| 139 | TEST_F(FusedSgdTiling, test_tiling_float16) | ||
| 140 | { | ||
| 141 | gert::StorageShape paramsShape = {{6400, 42}, {6400, 42}}; | ||
| 142 | gert::StorageShape gradsShape = {{6400, 42}, {6400, 42}}; | ||
| 143 | gert::StorageShape momentumShape = {{6400, 42}, {6400, 42}}; | ||
| 144 | gert::StorageShape gradScaleShape = {{1}, {1}}; | ||
| 145 | |||
| 146 | string compile_info_string = R"({"hardware_info": {"BT_SIZE": 0, "load3d_constraints": "1", | ||
| 147 | "Intrinsic_fix_pipe_l0c2out": false, | ||
| 148 | "Intrinsic_data_move_l12ub": true, | ||
| 149 | "Intrinsic_data_move_l0c2ub": true, | ||
| 150 | "Intrinsic_data_move_out2l1_nd2nz": false, | ||
| 151 | "UB_SIZE": 196608, "L2_SIZE": 33554432, "L1_SIZE": 524288, | ||
| 152 | "L0A_SIZE": 65536, "L0B_SIZE": 65536, "L0C_SIZE": 131072, | ||
| 153 | "CORE_NUM": 48} | ||
| 154 | })"; | ||
| 155 | map<string, string> soc_infos; | ||
| 156 | map<string, string> aicore_spec; | ||
| 157 | map<string, string> intrinsics; | ||
| 158 | |||
| 159 | GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics); | ||
| 160 | |||
| 161 | fe::PlatFormInfos platform_info; | ||
| 162 | platform_info.Init(); | ||
| 163 | |||
| 164 | struct FusedSgdCompileInfo { | ||
| 165 | }; | ||
| 166 | FusedSgdCompileInfo compile_info; | ||
| 167 | |||
| 168 | std::string op_type("FusedSgd"); | ||
| 169 | ASSERT_NE(gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str()), nullptr); | ||
| 170 | auto tiling_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling; | ||
| 171 | auto tiling_parse_func = gert::OpImplRegistry::GetInstance().GetOpImpl(op_type.c_str())->tiling_parse; | ||
| 172 | |||
| 173 | auto kernel_holder = | ||
| 174 | gert::KernelRunContextFaker() | ||
| 175 | .KernelIONum(2, 1) | ||
| 176 | .Inputs({const_cast<char*>(compile_info_string.c_str()), reinterpret_cast<void*>(&platform_info)}) | ||
| 177 | .Outputs({&compile_info}) | ||
| 178 | .Build(); | ||
| 179 | |||
| 180 | ASSERT_TRUE(kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->Init()); | ||
| 181 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos); | ||
| 182 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec); | ||
| 183 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore"); | ||
| 184 | kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes( | ||
| 185 | "AICoreintrinsicDtypeMap", intrinsics); | ||
| 186 | |||
| 187 | ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS); | ||
| 188 | |||
| 189 | auto param = gert::TilingData::CreateCap(4096); | ||
| 190 | auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096); | ||
| 191 | auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get()); | ||
| 192 | ASSERT_NE(param, nullptr); | ||
| 193 | auto holder = gert::TilingContextFaker() | ||
| 194 | .NodeIoNum(5, 3) | ||
| 195 | .IrInstanceNum({1, 1, 1, 1, 1}) | ||
| 196 | .InputShapes({¶msShape, &gradsShape, &momentumShape, &gradScaleShape}) | ||
| 197 | .OutputShapes({¶msShape, &gradsShape, &momentumShape}) | ||
| 198 | .CompileInfo(&compile_info) | ||
| 199 | .PlatformInfo(reinterpret_cast<char*>(&platform_info)) | ||
| 200 | .NodeInputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 201 | .NodeInputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 202 | .NodeInputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 203 | .NodeInputTd(3, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 204 | .NodeInputTd(4, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 205 | .NodeOutputTd(0, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 206 | .NodeOutputTd(1, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 207 | .NodeOutputTd(2, ge::DT_FLOAT16, ge::FORMAT_ND, ge::FORMAT_ND) | ||
| 208 | .NodeAttrs( | ||
| 209 | {{"weight_decay", Ops::NN::AnyValue::CreateFrom(0.01f)}, | ||
| 210 | {"momentum", Ops::NN::AnyValue::CreateFrom(0.9f)}, | ||
| 211 | {"lr", Ops::NN::AnyValue::CreateFrom(0.001f)}, | ||
| 212 | {"dampening", Ops::NN::AnyValue::CreateFrom(0.0f)}, | ||
| 213 | {"nesterov", Ops::NN::AnyValue::CreateFrom<bool>(true)}, | ||
| 214 | {"maximize", Ops::NN::AnyValue::CreateFrom<bool>(false)}, | ||
| 215 | {"is_first_step", Ops::NN::AnyValue::CreateFrom<bool>(true)}}) | ||
| 216 | .TilingData(param.get()) | ||
| 217 | .Workspace(ws_size) | ||
| 218 | .Build(); | ||
| 219 | |||
| 220 | gert::TilingContext* tiling_context = holder.GetContext<gert::TilingContext>(); | ||
| 221 | ASSERT_NE(tiling_context, nullptr); | ||
| 222 | ASSERT_NE(tiling_context->GetPlatformInfo(), nullptr); | ||
| 223 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("SoCInfo", soc_infos); | ||
| 224 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreSpec", aicore_spec); | ||
| 225 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore"); | ||
| 226 | holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics); | ||
| 227 | |||
| 228 | EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS); | ||
| 229 | |||
| 230 | auto tiling_key = tiling_context->GetTilingKey(); | ||
| 231 | ASSERT_EQ(tiling_key, 0); | ||
| 232 | } | ||
| @@ -0,0 +1,13 @@ | |||
| 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 | if ((UT_TEST_ALL OR OP_KERNEL_UT) AND NOT UT_DONE) | ||
| 12 | AddOpTestCase(fused_sgd "ascend910B1" "-DDTYPE_X=float") | ||
| 13 | endif() | ||
| @@ -0,0 +1,57 @@ | |||
| 1 | #!/usr/bin/env python3 | ||
| 2 | # -*- coding: utf-8 -*- | ||
| 3 | # ---------------------------------------------------------------------------- | ||
| 4 | # Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 5 | # This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 6 | # CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 7 | # Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 8 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 9 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 10 | # See LICENSE in the root of the software repository for the full text of the License. | ||
| 11 | # ---------------------------------------------------------------------------- | ||
| 12 | |||
| 13 | import sys | ||
| 14 | import numpy as np | ||
| 15 | import glob | ||
| 16 | import os | ||
| 17 | |||
| 18 | curr_dir = os.path.dirname(os.path.realpath(__file__)) | ||
| 19 | |||
| 20 | |||
| 21 | def compare_data(d_type): | ||
| 22 | if d_type == "float16": | ||
| 23 | np_dtype = np.float16 | ||
| 24 | precision = 1 / 1000 | ||
| 25 | elif d_type == "float32": | ||
| 26 | np_dtype = np.float32 | ||
| 27 | precision = 1 / 10000 | ||
| 28 | else: | ||
| 29 | np_dtype = np.float32 | ||
| 30 | precision = 1 / 1000 | ||
| 31 | |||
| 32 | golden_file_lists = sorted(glob.glob(curr_dir + "/*golden*.bin")) | ||
| 33 | output_file_lists = sorted(glob.glob(curr_dir + "/*output*.bin")) | ||
| 34 | |||
| 35 | data_same = True | ||
| 36 | for gold, out in zip(golden_file_lists, output_file_lists): | ||
| 37 | tmp_out = np.fromfile(out, np_dtype) | ||
| 38 | tmp_gold = np.fromfile(gold, np_dtype) | ||
| 39 | diff_res = np.isclose(tmp_out, tmp_gold, precision, 0, True) | ||
| 40 | diff_idx = np.where(diff_res != True)[0] | ||
| 41 | if len(diff_idx) == 0: | ||
| 42 | print(f"PASSED! {os.path.basename(gold)} vs {os.path.basename(out)}") | ||
| 43 | else: | ||
| 44 | print(f"FAILED! {os.path.basename(gold)} vs {os.path.basename(out)}") | ||
| 45 | for idx in diff_idx[:5]: | ||
| 46 | print(f" index: {idx}, output: {tmp_out[idx]}, golden: {tmp_gold[idx]}") | ||
| 47 | data_same = False | ||
| 48 | |||
| 49 | if not data_same: | ||
| 50 | exit(1) | ||
| 51 | |||
| 52 | |||
| 53 | if __name__ == '__main__': | ||
| 54 | if len(sys.argv) != 2: | ||
| 55 | print("Param num must be 2.") | ||
| 56 | exit(1) | ||
| 57 | compare_data(sys.argv[1]) | ||
| @@ -0,0 +1,96 @@ | |||
| 1 | #!/usr/bin/env python3 | ||
| 2 | # -*- coding: utf-8 -*- | ||
| 3 | # ---------------------------------------------------------------------------- | ||
| 4 | # Copyright (c) 2026 Huawei Technologies Co., Ltd. | ||
| 5 | # This program is free software, you can redistribute it and/or modify it under the terms and conditions of | ||
| 6 | # CANN Open Software License Agreement Version 2.0 (the "License"). | ||
| 7 | # Please refer to the License for details. You may not use this file except in compliance with the License. | ||
| 8 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | ||
| 9 | # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. | ||
| 10 | # See LICENSE in the root of the software repository for the full text of the License. | ||
| 11 | # ---------------------------------------------------------------------------- | ||
| 12 | |||
| 13 | import sys | ||
| 14 | import os | ||
| 15 | import numpy as np | ||
| 16 | import re | ||
| 17 | |||
| 18 | |||
| 19 | def parse_str_to_shape_list(shape_str): | ||
| 20 | shape_list = [] | ||
| 21 | shape_str_arr = re.findall(r"\{([0-9 ,]+)\}", shape_str) | ||
| 22 | for shape_str in shape_str_arr: | ||
| 23 | single_shape = [int(x) for x in shape_str.split(",")] | ||
| 24 | shape_list.append(single_shape) | ||
| 25 | return shape_list | ||
| 26 | |||
| 27 | |||
| 28 | def gen_data_and_golden(shape_str, d_type="float32"): | ||
| 29 | d_type_dict = { | ||
| 30 | "float32": np.float32, | ||
| 31 | "float16": np.float16, | ||
| 32 | } | ||
| 33 | np_type = d_type_dict[d_type] | ||
| 34 | shape_list = parse_str_to_shape_list(shape_str) | ||
| 35 | |||
| 36 | weight_decay = 0.01 | ||
| 37 | momentum = 0.0 | ||
| 38 | lr = 0.001 | ||
| 39 | dampening = 0.0 | ||
| 40 | nesterov = 0 | ||
| 41 | maximize = 0 | ||
| 42 | is_first_step = 1 | ||
| 43 | use_grad_scale = 1 | ||
| 44 | use_momentum = 0 | ||
| 45 | grad_scale = 0.5 | ||
| 46 | |||
| 47 | for index, shape in enumerate(shape_list): | ||
| 48 | params = (np.random.rand(*shape) * 2 - 1).astype(np.float32) * 100 | ||
| 49 | grads = (np.random.rand(*shape) * 2 - 1).astype(np.float32) * 100 | ||
| 50 | momentum_buf = (np.random.rand(*shape) * 2 - 1).astype(np.float32) * 100 | ||
| 51 | |||
| 52 | grads_inv = grads.copy() | ||
| 53 | if use_grad_scale: | ||
| 54 | inv_grad_scale = 1.0 / grad_scale | ||
| 55 | grads_inv = grads_inv * inv_grad_scale | ||
| 56 | |||
| 57 | grads_ref = grads_inv.copy() | ||
| 58 | |||
| 59 | if maximize: | ||
| 60 | grads_inv = -grads_inv | ||
| 61 | |||
| 62 | if weight_decay != 0.0: | ||
| 63 | grads_inv = grads_inv + weight_decay * params | ||
| 64 | |||
| 65 | if use_momentum: | ||
| 66 | if is_first_step: | ||
| 67 | momentum_buf_out = grads_inv.copy() | ||
| 68 | else: | ||
| 69 | momentum_buf_out = momentum * momentum_buf + (1.0 - dampening) * grads_inv | ||
| 70 | if nesterov: | ||
| 71 | grads_inv = grads_inv + momentum * momentum_buf_out | ||
| 72 | else: | ||
| 73 | grads_inv = momentum_buf_out.copy() | ||
| 74 | else: | ||
| 75 | momentum_buf_out = momentum_buf.copy() | ||
| 76 | |||
| 77 | params_ref = params - lr * grads_inv | ||
| 78 | |||
| 79 | params.astype(np_type).tofile(f"{d_type}_input_t_params_{index}.bin") | ||
| 80 | grads.astype(np_type).tofile(f"{d_type}_input_t_grads_{index}.bin") | ||
| 81 | momentum_buf.astype(np_type).tofile(f"{d_type}_input_t_momentum_{index}.bin") | ||
| 82 | params.astype(np_type).tofile(f"{d_type}_input_t_params_ref_{index}.bin") | ||
| 83 | grads.astype(np_type).tofile(f"{d_type}_input_t_grads_ref_{index}.bin") | ||
| 84 | momentum_buf.astype(np_type).tofile(f"{d_type}_input_t_momentum_ref_{index}.bin") | ||
| 85 | |||
| 86 | params_ref.astype(np_type).tofile(f"{d_type}_golden_t_params_ref_{index}.bin") | ||
| 87 | grads_ref.astype(np_type).tofile(f"{d_type}_golden_t_grads_ref_{index}.bin") | ||
| 88 | momentum_buf_out.astype(np_type).tofile(f"{d_type}_golden_t_momentum_ref_{index}.bin") | ||
| 89 | |||
| 90 | |||
| 91 | if __name__ == "__main__": | ||
| 92 | if len(sys.argv) != 3: | ||
| 93 | print("Param num must be 3.") | ||
| 94 | exit(1) | ||
| 95 | os.system("rm -rf *.bin") | ||
🟡 Medium Priority 在 该脚本被测试代码通过 变更：第95行新增 建议：使用 | |||
| 96 | gen_data_and_golden(sys.argv[1], sys.argv[2]) | ||
| @@ -0,0 +1,221 @@ | |||
| 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 | |||
| 13 | |||
| 14 | |||
| 15 | |||
| 16 | |||
| 17 | |||
| 18 | |||
| 19 | |||
| 20 | |||
| 21 | |||
| 22 | |||
| 23 | |||
| 24 | using namespace std; | ||
| 25 | |||
| 26 | extern "C" __global__ __aicore__ void fused_sgd( | ||
| 27 | GM_ADDR params, GM_ADDR grads, GM_ADDR momentum_buffer_list, | ||
| 28 | GM_ADDR grad_scale, | ||
| 29 | GM_ADDR params_ref, GM_ADDR grads_ref, GM_ADDR momentum_buffer_list_ref, | ||
| 30 | GM_ADDR workspace, GM_ADDR tiling); | ||
| 31 | |||
| 32 | class FusedSgdKernelTest : public testing::Test { | ||
| 33 | protected: | ||
| 34 | static void SetUpTestCase() { cout << "FusedSgdKernelTest SetUp" << endl; } | ||
| 35 | static void TearDownTestCase() { cout << "FusedSgdKernelTest TearDown" << endl; } | ||
| 36 | }; | ||
| 37 | |||
| 38 | template <typename T1, typename T2> | ||
| 39 | inline T1 CeilA2B(T1 a, T2 b) | ||
| 40 | { | ||
| 41 | if (b == 0) { | ||
| 42 | return a; | ||
| 43 | } | ||
| 44 | return (a + b - 1) / b; | ||
| 45 | } | ||
| 46 | |||
| 47 | template <typename T> | ||
| 48 | uint8_t* CreateNormTensorList(const std::vector<std::vector<uint64_t>>& shapeInfos, char* d_type, const char* tag = "fused_sgd") | ||
| 49 | { | ||
| 50 | uint64_t tensorListDescCount = 1 + shapeInfos.size() * 2; | ||
| 51 | for (auto s : shapeInfos) { | ||
| 52 | tensorListDescCount += s.size(); | ||
| 53 | } | ||
| 54 | std::vector<uint64_t> shapeSizeList; | ||
| 55 | uint64_t* tensorListDesc = (uint64_t*)AscendC::GmAlloc(tensorListDescCount * sizeof(uint64_t)); | ||
| 56 | *tensorListDesc = (tensorListDescCount - shapeInfos.size()) * sizeof(uint64_t); | ||
| 57 | uint64_t addrIndex = 0; | ||
| 58 | for (size_t i = 0; i < shapeInfos.size(); i++) { | ||
| 59 | addrIndex++; | ||
| 60 | uint16_t dimCount = shapeInfos[i].size(); | ||
| 61 | *(tensorListDesc + addrIndex) = ((uint64_t)(i) << 32) + dimCount; | ||
| 62 | uint64_t shapeSize = 1; | ||
| 63 | for (size_t j = 0; j < dimCount; j++) { | ||
| 64 | addrIndex++; | ||
| 65 | *(tensorListDesc + addrIndex) = shapeInfos[i][j]; | ||
| 66 | shapeSize *= shapeInfos[i][j]; | ||
| 67 | } | ||
| 68 | shapeSizeList.push_back(shapeSize); | ||
| 69 | } | ||
| 70 | for (size_t i = 0; i < shapeInfos.size(); i++) { | ||
| 71 | addrIndex++; | ||
| 72 | uint64_t dataSize = shapeSizeList[i] * sizeof(T); | ||
| 73 | uint8_t* dataPtr = (uint8_t*)AscendC::GmAlloc(CeilA2B(dataSize, 32) * 32); | ||
| 74 | std::stringstream fileName; | ||
| 75 | fileName << "./sgd_data/" << d_type << "_input_t_" << tag << "_" << i << ".bin"; | ||
| 76 | ReadFile(fileName.str(), dataSize, dataPtr, dataSize); | ||
| 77 | *(tensorListDesc + addrIndex) = (uint64_t)dataPtr; | ||
| 78 | } | ||
| 79 | return (uint8_t*)tensorListDesc; | ||
| 80 | } | ||
| 81 | |||
| 82 | template <typename T> | ||
| 83 | void FreeNormTensorList(uint8_t* addr, const std::vector<std::vector<uint64_t>>& shapeInfos, char* d_type, const char* tag = "fused_sgd") | ||
| 84 | { | ||
| 85 | uint64_t dataPtrOffset = *((uint64_t*)addr); | ||
| 86 | uint8_t* dataAddr = addr + dataPtrOffset; | ||
| 87 | for (size_t i = 0; i < shapeInfos.size(); i++) { | ||
| 88 | uint64_t shapeSize = 1; | ||
| 89 | for (size_t j = 0; j < shapeInfos[i].size(); j++) { | ||
| 90 | shapeSize *= shapeInfos[i][j]; | ||
| 91 | } | ||
| 92 | uint8_t* tensorAddr = (uint8_t*)(*((uint64_t*)(dataAddr) + i)); | ||
| 93 | std::stringstream fileName; | ||
| 94 | fileName << "./sgd_data/" << d_type << "_output_t_" << tag << "_" << i << ".bin"; | ||
| 95 | WriteFile(fileName.str(), tensorAddr, shapeSize * sizeof(T)); | ||
| 96 | AscendC::GmFree((void*)(tensorAddr)); | ||
| 97 | } | ||
| 98 | AscendC::GmFree((void*)addr); | ||
| 99 | } | ||
| 100 | |||
| 101 | TEST_F(FusedSgdKernelTest, test_fp32_basic) | ||
| 102 | { | ||
| 103 | size_t tilingSize = sizeof(FusedSgdTilingData); | ||
| 104 | uint32_t blockDim = 1; | ||
| 105 | std::vector<std::vector<uint64_t>> shapeInfos = {{4}}; | ||
| 106 | |||
| 107 | system( | ||
| 108 | "cp -rf " | ||
| 109 | "../../../../optim/fused_sgd/tests/ut/op_kernel/sgd_data ./"); | ||
| 110 | system("chmod -R 755 ./sgd_data/"); | ||
| 111 | system("cd ./sgd_data/ && python3 gen_data.py '{{4}}' 'float32'"); | ||
| 112 | |||
| 113 | uint8_t* paramsBuf = CreateNormTensorList<float>(shapeInfos, "float32", "params"); | ||
| 114 | uint8_t* gradsBuf = CreateNormTensorList<float>(shapeInfos, "float32", "grads"); | ||
| 115 | uint8_t* momentumBuf = CreateNormTensorList<float>(shapeInfos, "float32", "momentum"); | ||
| 116 | uint8_t* paramsRefBuf = CreateNormTensorList<float>(shapeInfos, "float32", "params_ref"); | ||
| 117 | uint8_t* gradsRefBuf = CreateNormTensorList<float>(shapeInfos, "float32", "grads_ref"); | ||
| 118 | uint8_t* momentumRefBuf = CreateNormTensorList<float>(shapeInfos, "float32", "momentum_ref"); | ||
| 119 | uint8_t* gradScaleBuf = (uint8_t*)AscendC::GmAlloc(sizeof(float)); | ||
| 120 | uint8_t* workspace = (uint8_t*)AscendC::GmAlloc(16 * 1024 * 1024); | ||
| 121 | uint8_t* tiling = (uint8_t*)AscendC::GmAlloc(tilingSize); | ||
| 122 | |||
| 123 | float* gradScalePtr = reinterpret_cast<float*>(gradScaleBuf); | ||
| 124 | gradScalePtr[0] = 0.5f; | ||
| 125 | |||
| 126 | FusedSgdTilingData* tilingData = reinterpret_cast<FusedSgdTilingData*>(tiling); | ||
| 127 | tilingData->weightDecay = 0.01f; | ||
| 128 | tilingData->momentum = 0.5f; | ||
| 129 | tilingData->lr = 0.001f; | ||
| 130 | tilingData->dampening = 0.0f; | ||
| 131 | tilingData->nesterov = 0; | ||
| 132 | tilingData->maximize = 0; | ||
| 133 | tilingData->isFirstStep = 1; | ||
| 134 | tilingData->useGradScale = 1; | ||
| 135 | tilingData->useMomentum = 1; | ||
| 136 | tilingData->tensorNum = 1; | ||
| 137 | tilingData->tensorsPerCore = 1; | ||
| 138 | tilingData->usedCoreNum = 1; | ||
| 139 | tilingData->coreCalcMax = 4088; | ||
| 140 | |||
| 141 | AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 142 | ICPU_SET_TILING_KEY(0); | ||
| 143 | ICPU_RUN_KF(fused_sgd, blockDim, | ||
| 144 | paramsBuf, gradsBuf, momentumBuf, | ||
| 145 | gradScaleBuf, | ||
| 146 | paramsRefBuf, gradsRefBuf, momentumRefBuf, | ||
| 147 | workspace, (uint8_t*)(tiling)); | ||
| 148 | |||
| 149 | FreeNormTensorList<float>(paramsRefBuf, shapeInfos, "float32", "params_ref"); | ||
| 150 | FreeNormTensorList<float>(gradsRefBuf, shapeInfos, "float32", "grads_ref"); | ||
| 151 | FreeNormTensorList<float>(momentumRefBuf, shapeInfos, "float32", "momentum_ref"); | ||
| 152 | FreeNormTensorList<float>(paramsBuf, shapeInfos, "float32", "params_ref"); | ||
| 153 | FreeNormTensorList<float>(gradsBuf, shapeInfos, "float32", "grads_ref"); | ||
| 154 | FreeNormTensorList<float>(momentumBuf, shapeInfos, "float32", "momentum_ref"); | ||
| 155 | AscendC::GmFree(gradScaleBuf); | ||
| 156 | AscendC::GmFree(workspace); | ||
| 157 | AscendC::GmFree(tiling); | ||
| 158 | |||
| 159 | system("cd ./sgd_data/ && python3 compare_data.py 'float32'"); | ||
| 160 | } | ||
| 161 | |||
| 162 | TEST_F(FusedSgdKernelTest, test_fp16_basic) | ||
| 163 | { | ||
| 164 | size_t tilingSize = sizeof(FusedSgdTilingData); | ||
| 165 | uint32_t blockDim = 1; | ||
| 166 | std::vector<std::vector<uint64_t>> shapeInfos = {{4}}; | ||
| 167 | |||
| 168 | system( | ||
| 169 | "cp -rf " | ||
| 170 | "../../../../optim/fused_sgd/tests/ut/op_kernel/sgd_data ./"); | ||
| 171 | system("chmod -R 755 ./sgd_data/"); | ||
| 172 | system("cd ./sgd_data/ && python3 gen_data.py '{{4}}' 'float16'"); | ||
| 173 | |||
| 174 | uint8_t* paramsBuf = CreateNormTensorList<half>(shapeInfos, "float16", "params"); | ||
| 175 | uint8_t* gradsBuf = CreateNormTensorList<half>(shapeInfos, "float16", "grads"); | ||
| 176 | uint8_t* momentumBuf = CreateNormTensorList<half>(shapeInfos, "float16", "momentum"); | ||
| 177 | uint8_t* paramsRefBuf = CreateNormTensorList<half>(shapeInfos, "float16", "params_ref"); | ||
| 178 | uint8_t* gradsRefBuf = CreateNormTensorList<half>(shapeInfos, "float16", "grads_ref"); | ||
| 179 | uint8_t* momentumRefBuf = CreateNormTensorList<half>(shapeInfos, "float16", "momentum_ref"); | ||
| 180 | uint8_t* gradScaleBuf = (uint8_t*)AscendC::GmAlloc(sizeof(float)); | ||
| 181 | uint8_t* workspace = (uint8_t*)AscendC::GmAlloc(16 * 1024 * 1024); | ||
| 182 | uint8_t* tiling = (uint8_t*)AscendC::GmAlloc(tilingSize); | ||
| 183 | |||
| 184 | float* gradScalePtr = reinterpret_cast<float*>(gradScaleBuf); | ||
| 185 | gradScalePtr[0] = 0.5f; | ||
| 186 | |||
| 187 | FusedSgdTilingData* tilingData = reinterpret_cast<FusedSgdTilingData*>(tiling); | ||
| 188 | tilingData->weightDecay = 0.01f; | ||
| 189 | tilingData->momentum = 0.5f; | ||
| 190 | tilingData->lr = 0.001f; | ||
| 191 | tilingData->dampening = 0.0f; | ||
| 192 | tilingData->nesterov = 0; | ||
| 193 | tilingData->maximize = 0; | ||
| 194 | tilingData->isFirstStep = 1; | ||
| 195 | tilingData->useGradScale = 1; | ||
| 196 | tilingData->useMomentum = 1; | ||
| 197 | tilingData->tensorNum = 1; | ||
| 198 | tilingData->tensorsPerCore = 1; | ||
| 199 | tilingData->usedCoreNum = 1; | ||
| 200 | tilingData->coreCalcMax = 4088; | ||
| 201 | |||
| 202 | AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 203 | ICPU_SET_TILING_KEY(0); | ||
| 204 | ICPU_RUN_KF(fused_sgd, blockDim, | ||
| 205 | paramsBuf, gradsBuf, momentumBuf, | ||
| 206 | gradScaleBuf, | ||
| 207 | paramsRefBuf, gradsRefBuf, momentumRefBuf, | ||
| 208 | workspace, (uint8_t*)(tiling)); | ||
| 209 | |||
| 210 | FreeNormTensorList<half>(paramsRefBuf, shapeInfos, "float16", "params_ref"); | ||
| 211 | FreeNormTensorList<half>(gradsRefBuf, shapeInfos, "float16", "grads_ref"); | ||
| 212 | FreeNormTensorList<half>(momentumRefBuf, shapeInfos, "float16", "momentum_ref"); | ||
| 213 | FreeNormTensorList<half>(paramsBuf, shapeInfos, "float16", "params_ref"); | ||
| 214 | FreeNormTensorList<half>(gradsBuf, shapeInfos, "float16", "grads_ref"); | ||
| 215 | FreeNormTensorList<half>(momentumBuf, shapeInfos, "float16", "momentum_ref"); | ||
| 216 | AscendC::GmFree(gradScaleBuf); | ||
| 217 | AscendC::GmFree(workspace); | ||
| 218 | AscendC::GmFree(tiling); | ||
| 219 | |||
| 220 | system("cd ./sgd_data/ && python3 compare_data.py 'float16'"); | ||
| 221 | } | ||
🟡 Medium Priority 在 CMakeLists.txt（ 但
该测试用例预期验证 fp16 路径，但实际执行的是 F32 内核，测试结果无意义。 变更：第162-221行新增的 建议：需要在 majiajian_hw 6月22日 评论： 6月22日 评论： | |||
🟠 High Priority
在
docs/zh/op_list.md第3294行，href 路径为../../optim/fusedSgd/README.md，但实际目录名为optim/fused_sgd/（全小写，下划线分割）。在区分大小写的文件系统上，该链接将无法访问，返回404。变更：第 3294 行新增
<a href="../../optim/fusedSgd/README.md">fusedSgd</a>实际路径应为
../../optim/fused_sgd/README.md。建议：将 href 中的
fusedSgd改为fused_sgd，使其与实际目录名一致。_sgd</a></td>