Star936
1856
代码
代码
Issues238
Pull Requests342
Actions
项目讨论
Wiki
分析
Star936
1856
已合并
feat(arch35): add Ascend 950 support for Adam optimizer operators #8078
daiwei18创建于 7月30日
feat(arch35): add Ascend 950 support for Adam optimizer operators #8078
已合并
daiwei18创建于 7月30日
从daiwei18/ops-nn930-optim合入到cann/ops-nnmaster
共 19 个文件变更+3214-30
Moptim/apply_adam_w_quant/README.md+1-1
@@ -4,7 +4,7 @@
4 4 
5|产品 | 是否支持 |5|产品 | 是否支持 |
6|:-------------------------|:----------:|6|:-------------------------|:----------:|
7-| <term>Ascend 950PR/Ascend 950DT</term> | x |7+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
8| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |8| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
9| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |9| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |
10| <term>Atlas 200I/500 A2 推理产品</term> | × |10| <term>Atlas 200I/500 A2 推理产品</term> | × |
Moptim/apply_adam_w_quant/docs/aclnnApplyAdamWQuant.md+5-5
@@ -6,7 +6,7 @@
6 6 
7|产品 | 是否支持 |7|产品 | 是否支持 |
8|:-------------------------|:----------:|8|:-------------------------|:----------:|
9-| <term>Ascend 950PR/Ascend 950DT</term> | x |9+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
10| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |10| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
11| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |11| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |
12| <term>Atlas 200I/500 A2 推理产品</term> | × |12| <term>Atlas 200I/500 A2 推理产品</term> | × |
@@ -164,7 +164,7 @@ aclnnStatus aclnnApplyAdamWQuant(
164 <td class="tg-0pky">absmaxMRef（aclTensor*）</td>164 <td class="tg-0pky">absmaxMRef（aclTensor*）</td>
165 <td class="tg-0pky">输入/输出</td>165 <td class="tg-0pky">输入/输出</td>
166 <td class="tg-0pky">每blockSize(256)个vRef对应一个最大值，用于对mRef索引选择qmapM中的值乘以对应的absmaxMRef进行反量化。再通过更新后的mRef每blockSize(256)个选择出一个最大值，作为absmaxMRef的输出。</td>166 <td class="tg-0pky">每blockSize(256)个vRef对应一个最大值，用于对mRef索引选择qmapM中的值乘以对应的absmaxMRef进行反量化。再通过更新后的mRef每blockSize(256)个选择出一个最大值，作为absmaxMRef的输出。</td>
167- <td class="tg-0pky">shape要求为“absmaxMRef.size = mRef.size/blockSize”。</td>167+ <td class="tg-0pky">shape要求为“absmaxMRef.size = ceil(mRef.size/blockSize)”。</td>
168 <td class="tg-0pky">FLOAT32</td>168 <td class="tg-0pky">FLOAT32</td>
169 <td class="tg-0pky">ND</td>169 <td class="tg-0pky">ND</td>
170 <td class="tg-0pky">1</td>170 <td class="tg-0pky">1</td>
@@ -174,7 +174,7 @@ aclnnStatus aclnnApplyAdamWQuant(
174 <td class="tg-0pky">absmaxVRef（aclTensor*）</td>174 <td class="tg-0pky">absmaxVRef（aclTensor*）</td>
175 <td class="tg-0pky">输入/输出</td>175 <td class="tg-0pky">输入/输出</td>
176 <td class="tg-0pky">每blockSize(256)个vRef对应一个最大值，用于对vRef索引选择qmapV中的值乘以对应的absmaxVRef进行反量化。再通过更新后的vRef每blockSize(256)个选择出一个最大值，作为absmaxVRef的输出。</td>176 <td class="tg-0pky">每blockSize(256)个vRef对应一个最大值，用于对vRef索引选择qmapV中的值乘以对应的absmaxVRef进行反量化。再通过更新后的vRef每blockSize(256)个选择出一个最大值，作为absmaxVRef的输出。</td>
177- <td class="tg-0pky">shape要求为“absmaxVRef.size = vRef.size/blockSize”。</td>177+ <td class="tg-0pky">shape要求为“absmaxVRef.size = ceil(vRef.size/blockSize)”。</td>
178 <td class="tg-0pky">FLOAT32</td>178 <td class="tg-0pky">FLOAT32</td>
179 <td class="tg-0pky">ND</td>179 <td class="tg-0pky">ND</td>
180 <td class="tg-0pky">1</td>180 <td class="tg-0pky">1</td>
@@ -373,8 +373,8 @@ aclnnStatus aclnnApplyAdamWQuant(
373 - varRef.shape = grad.shape373 - varRef.shape = grad.shape
374 - varRef.shape = mRef.shape374 - varRef.shape = mRef.shape
375 - varRef.shape = vRef.shape375 - varRef.shape = vRef.shape
376- - varRef.size/blockSize = absmaxMRef.size376+ - ceil(varRef.size/blockSize) = absmaxMRef.size
377- - varRef.size/blockSize = absmaxVRef.size377+ - ceil(varRef.size/blockSize) = absmaxVRef.size
378 378 
379 确定性计算：379 确定性计算：
380 380 
Moptim/apply_adam_w_quant/op_host/CMakeLists.txt+9-6
@@ -1,12 +1,15 @@
1# ----------------------------------------------------------------------------1# ----------------------------------------------------------------------------
2-# Copyright (c) 2025 Huawei Technologies Co., Ltd.2+# Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4# CANN Open Software License Agreement Version 2.0 (the "License").4# CANN Open Software License Agreement Version 2.0 (the "License").
5# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
6-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# ----------------------------------------------------------------------------9# ----------------------------------------------------------------------------
10-message(STATUS "=== Debug: start ops.optim.apply_adam_w_quant.op_host.CMakeLists.txt ")10+set(SUPPORT_COMPUTE_UNIT "ascend910b" "ascend910_93" "ascend950")
11- 11+# 910b/910_93 走 arch22 通用 tiling，950 走 arch35 regbase tiling。
12-add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR} OPTYPE apply_adam_w_quant ACLNNTYPE aclnn)12+set(SUPPORT_TILING_DIR "arch22" "arch22" "arch35")
13+add_modules_sources(HOSTNAME ${OPHOST_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR}
14+ OPTYPE apply_adam_w_quant ACLNNTYPE aclnn
15+ COMPUTE_UNIT ${SUPPORT_COMPUTE_UNIT} TILING_DIR ${SUPPORT_TILING_DIR})
Moptim/apply_adam_w_quant/op_host/apply_adam_w_quant_def.cpp+15-1
@@ -1,5 +1,5 @@
1/**1/**
2- * Copyright (c) 2025 Huawei Technologies Co., Ltd.2+ * Copyright (c) 2025-2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4 * CANN Open Software License Agreement Version 2.0 (the "License").4 * CANN Open Software License Agreement Version 2.0 (the "License").
5 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
@@ -116,6 +116,20 @@ public:
116 this->Attr("block_size").AttrType(OPTIONAL).Int(BLOCKSIZE);116 this->Attr("block_size").AttrType(OPTIONAL).Int(BLOCKSIZE);
117 this->AICore().AddConfig("ascend910b");117 this->AICore().AddConfig("ascend910b");
118 this->AICore().AddConfig("ascend910_93");118 this->AICore().AddConfig("ascend910_93");
119+ 
120+ // arch35(Ascend950)regbase 出包配置:契约与 A2 完全一致(9 入/5 出、fp32/fp16/bf16),
121+ // 仅走 regbase 编译路径;opFile/opInterface 名与 A2 相同,由 op_kernel CMake 的
122+ // COMPUTE_UNITS 路由区分根 kernel(A2)与 arch35/ kernel(A5)。
123+ OpAICoreConfig aicoreConfig950;
124+ aicoreConfig950.DynamicCompileStaticFlag(true)
125+ .DynamicFormatFlag(false)
126+ .DynamicRankSupportFlag(true)
127+ .DynamicShapeSupportFlag(true)
128+ .NeedCheckSupportFlag(false)
129+ .PrecisionReduceFlag(true)
130+ .ExtendCfgInfo("opFile.value", "apply_adam_w_quant")
131+ .ExtendCfgInfo("opInterface.value", "apply_adam_w_quant");
132+ this->AICore().AddConfig("ascend950", aicoreConfig950);
119 }133 }
120};134};
121 135 
Roptim/apply_adam_w_quant/op_host/apply_adam_w_quant_tiling.cpp→optim/apply_adam_w_quant/op_host/arch22/apply_adam_w_quant_tiling.cpp+1-1
@@ -304,4 +304,4 @@ static ge::graphStatus TilingPrepare4ApplyAdamWQuant(gert::TilingParseContext* c
304IMPL_OP_OPTILING(ApplyAdamWQuant)304IMPL_OP_OPTILING(ApplyAdamWQuant)
305 .Tiling(Tiling4ApplyAdamWQuant)305 .Tiling(Tiling4ApplyAdamWQuant)
306 .TilingParse<Tiling4ApplyAdamWQuantCompileInfo>(TilingPrepare4ApplyAdamWQuant);306 .TilingParse<Tiling4ApplyAdamWQuantCompileInfo>(TilingPrepare4ApplyAdamWQuant);
307-} // namespace optiling307+} // namespace optiling
Aoptim/apply_adam_w_quant/op_host/arch35/apply_adam_w_quant_tiling_arch35.cpp+352-0
@@ -0,0 +1,352 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ *
10+ * NOTE: Portions of this code were AI-generated and have been
11+ * technically reviewed for functional accuracy and security
12+ */
13+ 
14+/*!
15+ * \file apply_adam_w_quant_tiling_arch35.cpp
16+ * \brief ApplyAdamWQuant regbase (Ascend950) tiling implementation.
17+ *
18+ * 与 A2 tiling 完全同算法(blockwise-256 分核 + UB 切分),仅出包路径不同:regbase 走
19+ * RawTilingData 直写裸 buffer。核数/UB 由 PlatformAscendC 运行时取(arch35 自适应)。
20+ */
21+#include "apply_adam_w_quant_tiling_arch35.h"
22+ 
23+#include <cstdint>
24+#include <securec.h>
25+#include <string>
26+#include "log/log.h"
27+#include "error_util.h"
28+#include "register/op_def_registry.h"
29+#include "tiling/platform/platform_ascendc.h"
30+#include "op_host/tiling_templates_registry.h"
31+ 
32+using namespace ge;
33+ 
34+namespace optiling {
35+namespace {
36+constexpr size_t INDEX_IN_VAR = 0;
37+constexpr size_t INDEX_IN_GRAD = 1;
38+constexpr size_t INDEX_IN_M = 2;
39+constexpr size_t INDEX_IN_V = 3;
40+constexpr size_t INDEX_IN_QMAP_M = 4;
41+constexpr size_t INDEX_IN_QMAP_V = 5;
42+constexpr size_t INDEX_IN_ABSMAX_M = 6;
43+constexpr size_t INDEX_IN_ABSMAX_V = 7;
44+constexpr size_t INDEX_IN_STEP = 8;
45+ 
46+constexpr size_t INDEX_ATTR_LR = 0;
47+constexpr size_t INDEX_ATTR_BETA1 = 1;
48+constexpr size_t INDEX_ATTR_BETA2 = 2;
49+constexpr size_t INDEX_ATTR_WEIGHT_DECAY = 3;
50+constexpr size_t INDEX_ATTR_EPS = 4;
51+constexpr size_t INDEX_ATTR_GNORM_SCALE = 5;
52+constexpr size_t INDEX_ATTR_BLOCK_SIZE = 7;
53+ 
54+constexpr uint64_t TILINGKEY_DATA_VAR_FLOAT = 100;
55+constexpr uint64_t TILINGKEY_DATA_VAR_FLOAT16 = 200;
56+constexpr uint64_t TILINGKEY_DATA_VAR_BFLOAT16 = 300;
57+constexpr uint64_t QMAP_SIZE = 256;
58+constexpr uint64_t SIZE_OF_FLOAT = 4;
59+constexpr uint64_t SIZE_OF_FLOAT16 = 2;
60+constexpr uint64_t PER_BLOCK_OF_MAX_NUM = 1;
61+constexpr uint64_t ONE_BLOCK_NEED_BUF = 10;
62+constexpr uint64_t NUM_OF_QMAP = 2;
63+constexpr int64_t BLOCKSIZE = 256;
64+ 
65+inline uint64_t CeilDiv(uint64_t a, uint64_t b) { return (b == 0 ? 0 : ((a + b - 1) / b)); }
66+ 
67+inline bool IsSameShape(const gert::Shape& shape1, const gert::Shape& shape2)
68+{
69+ size_t dimNum = shape1.GetDimNum();
70+ if (shape2.GetDimNum() != dimNum) {
71+ return false;
72+ }
73+ for (size_t i = 0; i < dimNum; ++i) {
74+ if (shape1.GetDim(i) != shape2.GetDim(i)) {
75+ return false;
76+ }
77+ }
78+ return true;
79+}
80+} // namespace
81+ 
82+ge::graphStatus ApplyAdamWQuantRegbaseTiling::GetAttributes()
83+{
84+ auto* attrs = tilingContext_->GetAttrs();
85+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrs);
86+ 
87+ auto* attrLr = attrs->GetAttrPointer<float>(INDEX_ATTR_LR);
88+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrLr);
89+ lr_ = *attrLr;
90+ 
91+ auto* attrBeta1 = attrs->GetAttrPointer<float>(INDEX_ATTR_BETA1);
92+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrBeta1);
93+ beta1_ = *attrBeta1;
94+ 
95+ auto* attrBeta2 = attrs->GetAttrPointer<float>(INDEX_ATTR_BETA2);
96+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrBeta2);
97+ beta2_ = *attrBeta2;
98+ 
99+ auto* attrWeightDecay = attrs->GetAttrPointer<float>(INDEX_ATTR_WEIGHT_DECAY);
100+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrWeightDecay);
101+ weightDecay_ = *attrWeightDecay;
102+ 
103+ auto* attrEps = attrs->GetAttrPointer<float>(INDEX_ATTR_EPS);
104+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrEps);
105+ eps_ = *attrEps;
106+ 
107+ auto* attrGnormScale = attrs->GetAttrPointer<float>(INDEX_ATTR_GNORM_SCALE);
108+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrGnormScale);
109+ gnormScale_ = *attrGnormScale;
110+ 
111+ auto* attrBlockSize = attrs->GetAttrPointer<int64_t>(INDEX_ATTR_BLOCK_SIZE);
112+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, attrBlockSize);
113+ blockSize_ = *attrBlockSize;
114+ OP_CHECK_IF(blockSize_ != BLOCKSIZE,
115+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
116+ "attr block_size should be 256, please check."),
117+ return ge::GRAPH_FAILED);
118+ return ge::GRAPH_SUCCESS;
119+}
120+ 
121+ge::graphStatus ApplyAdamWQuantRegbaseTiling::CheckInputShape()
122+{
123+ auto varShapePtr = tilingContext_->GetInputShape(INDEX_IN_VAR);
124+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, varShapePtr);
125+ auto gradShapePtr = tilingContext_->GetInputShape(INDEX_IN_GRAD);
126+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, gradShapePtr);
127+ auto mShapePtr = tilingContext_->GetInputShape(INDEX_IN_M);
128+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, mShapePtr);
129+ auto vShapePtr = tilingContext_->GetInputShape(INDEX_IN_V);
130+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, vShapePtr);
131+ auto qmapMShapePtr = tilingContext_->GetInputShape(INDEX_IN_QMAP_M);
132+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, qmapMShapePtr);
133+ auto qmapVShapePtr = tilingContext_->GetInputShape(INDEX_IN_QMAP_V);
134+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, qmapVShapePtr);
135+ auto absmaxMShapePtr = tilingContext_->GetInputShape(INDEX_IN_ABSMAX_M);
136+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, absmaxMShapePtr);
137+ auto absmaxVShapePtr = tilingContext_->GetInputShape(INDEX_IN_ABSMAX_V);
138+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, absmaxVShapePtr);
139+ auto stepShapePtr = tilingContext_->GetInputShape(INDEX_IN_STEP);
140+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, stepShapePtr);
141+ 
142+ const gert::Shape& varShape = varShapePtr->GetStorageShape();
143+ const gert::Shape& gradShape = gradShapePtr->GetStorageShape();
144+ const gert::Shape& mShape = mShapePtr->GetStorageShape();
145+ const gert::Shape& vShape = vShapePtr->GetStorageShape();
146+ const gert::Shape& qmapMShape = qmapMShapePtr->GetStorageShape();
147+ const gert::Shape& qmapVShape = qmapVShapePtr->GetStorageShape();
148+ const gert::Shape& absmaxMShape = absmaxMShapePtr->GetStorageShape();
149+ const gert::Shape& absmaxVShape = absmaxVShapePtr->GetStorageShape();
150+ const gert::Shape& stepShape = stepShapePtr->GetStorageShape();
151+ 
152+ bool isDiffShape = !IsSameShape(varShape, gradShape) || !IsSameShape(varShape, mShape) ||
153+ !IsSameShape(varShape, vShape);
154+ bool isQmapDiffShape = !IsSameShape(qmapMShape, qmapVShape) ||
155+ static_cast<uint64_t>(qmapMShape.GetShapeSize()) != QMAP_SIZE;
156+ uint64_t expectedAbsmaxSize = CeilDiv(static_cast<uint64_t>(varShape.GetShapeSize()),
157+ static_cast<uint64_t>(blockSize_));
158+ bool isAbsmaxDiffShape = static_cast<uint64_t>(absmaxMShape.GetShapeSize()) != expectedAbsmaxSize ||
159+ static_cast<uint64_t>(absmaxVShape.GetShapeSize()) != expectedAbsmaxSize;
160+ OP_CHECK_IF(isDiffShape,
161+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
162+ "var,grad,m,v should have same shape, please check."),
163+ return ge::GRAPH_FAILED);
164+ OP_CHECK_IF(
165+ isQmapDiffShape,
166+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
167+ "qmapM and qmapV should be same shape,shape is [256], please check."),
168+ return ge::GRAPH_FAILED);
169+ OP_CHECK_IF(isAbsmaxDiffShape,
170+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(
171+ tilingContext_->GetNodeName(), "parameter", "invalid",
172+ "absmaxM and absmaxV should each contain ceil(var size / block_size) elements, please check."),
173+ return ge::GRAPH_FAILED);
174+ OP_CHECK_IF(stepShape.GetDimNum() != 1 || stepShape.GetDim(0) != 1,
175+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
176+ "step should have only one element, please check."),
177+ return ge::GRAPH_FAILED);
178+ return ge::GRAPH_SUCCESS;
179+}
180+ 
181+ge::graphStatus ApplyAdamWQuantRegbaseTiling::DetermineTilingKey()
182+{
183+ auto dtypePtr = tilingContext_->GetInputDesc(INDEX_IN_VAR);
184+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, dtypePtr);
185+ auto dtype = dtypePtr->GetDataType();
186+ if (dtype == ge::DataType::DT_FLOAT) {
范其瑞
范其瑞范其瑞7月31日

新写代码都有tilingkey中冗余dtype信息问题，请排查整改后再检视

likedislike
daiwei18
daiwei18
8月1日 评论：
187+ tilingKey_ = TILINGKEY_DATA_VAR_FLOAT;
188+ } else if (dtype == ge::DataType::DT_FLOAT16) {
189+ tilingKey_ = TILINGKEY_DATA_VAR_FLOAT16;
190+ } else if (dtype == ge::DataType::DT_BF16) {
191+ tilingKey_ = TILINGKEY_DATA_VAR_BFLOAT16;
192+ } else {
193+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
194+ "var dtype should be float/float16/bfloat16, please check.");
195+ return ge::GRAPH_FAILED;
196+ }
197+ return ge::GRAPH_SUCCESS;
198+}
199+ 
200+ge::graphStatus ApplyAdamWQuantRegbaseTiling::DoTiling()
201+{
202+ auto platformInfoPtr = tilingContext_->GetPlatformInfo();
203+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, platformInfoPtr);
204+ const auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
205+ uint64_t aivNum = ascendcPlatform.GetCoreNumAiv();
206+ OP_CHECK_IF(aivNum == 0,
207+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
208+ "aivNum can not be 0, please check."),
209+ return ge::GRAPH_FAILED);
210+ uint64_t maxUbSize = 0;
211+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, maxUbSize);
212+ OP_CHECK_IF(maxUbSize == 0,
213+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
214+ "ub size can not be 0, please check."),
215+ return ge::GRAPH_FAILED);
216+ 
217+ auto shapePtr = tilingContext_->GetInputShape(INDEX_IN_VAR);
218+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, shapePtr);
219+ 
220+ uint64_t oneBlockSize = static_cast<uint64_t>(blockSize_) *
221+ (SIZE_OF_FLOAT * ONE_BLOCK_NEED_BUF + SIZE_OF_FLOAT16 + SIZE_OF_FLOAT) +
222+ (PER_BLOCK_OF_MAX_NUM + PER_BLOCK_OF_MAX_NUM) * SIZE_OF_FLOAT;
223+ OP_CHECK_IF(oneBlockSize == 0,
224+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
225+ "one core max size can not be 0, please check."),
226+ return ge::GRAPH_FAILED);
227+ uint64_t reservedQmap = QMAP_SIZE * SIZE_OF_FLOAT * NUM_OF_QMAP;
228+ OP_CHECK_IF(maxUbSize <= reservedQmap,
229+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
230+ "ub size too small for qmap, please check."),
231+ return ge::GRAPH_FAILED);
232+ perCoreDoBlockNum_ = (maxUbSize - reservedQmap) / oneBlockSize;
233+ OP_CHECK_IF(perCoreDoBlockNum_ == 0,
234+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
235+ "one core do block num can not be 0, please check."),
236+ return ge::GRAPH_FAILED);
237+ 
238+ uint64_t totalDataNum = static_cast<uint64_t>(shapePtr->GetStorageShape().GetShapeSize());
239+ uint64_t blockNum = CeilDiv(totalDataNum, static_cast<uint64_t>(blockSize_));
240+ uint64_t totalUseNumCore = CeilDiv(blockNum, perCoreDoBlockNum_);
241+ OP_CHECK_IF(totalUseNumCore == 0,
242+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
243+ "total use num core can not be 0, please check."),
244+ return ge::GRAPH_FAILED);
245+ lastCoreLastBlock_ = blockNum - (totalUseNumCore - 1) * perCoreDoBlockNum_;
246+ lastBlockSize_ = totalDataNum - (blockNum - 1) * static_cast<uint64_t>(blockSize_);
247+ useNumCore_ = CeilDiv(totalUseNumCore, CeilDiv(totalUseNumCore, aivNum));
248+ OP_CHECK_IF(useNumCore_ == 0,
249+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
250+ "use_num_core can't be 0."),
251+ return ge::GRAPH_FAILED);
252+ 
253+ lastPreCoreRowWork_ = totalUseNumCore / useNumCore_;
254+ notLastCoreNum_ = totalUseNumCore - lastPreCoreRowWork_ * useNumCore_;
255+ notLastPreCoreRowWork_ = lastPreCoreRowWork_ + 1;
256+ return ge::GRAPH_SUCCESS;
257+}
258+ 
259+ge::graphStatus ApplyAdamWQuantRegbaseTiling::SetTilingData()
260+{
261+ auto* rawTilingData = tilingContext_->GetRawTilingData();
262+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, rawTilingData);
263+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, rawTilingData->GetData());
264+ OP_CHECK_IF(rawTilingData->GetCapacity() < sizeof(ApplyAdamWQuantRegbaseTilingData),
265+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(
266+ tilingContext_->GetNodeName(), "parameter", "invalid",
267+ "tiling data capacity is less than ApplyAdamWQuantRegbaseTilingData size."),
268+ return ge::GRAPH_FAILED);
269+ 
270+ ApplyAdamWQuantRegbaseTilingData tiling;
271+ tiling.use_num_core = useNumCore_;
272+ tiling.last_pre_core_row_work = lastPreCoreRowWork_;
273+ tiling.not_last_core_num = notLastCoreNum_;
274+ tiling.not_last_pre_core_row_work = notLastPreCoreRowWork_;
275+ tiling.last_core_last_block = lastCoreLastBlock_;
276+ tiling.lr = lr_;
277+ tiling.beta1 = beta1_;
278+ tiling.beta2 = beta2_;
279+ tiling.weight_decay = weightDecay_;
280+ tiling.eps = eps_;
281+ tiling.gnorm_scale = gnormScale_;
282+ tiling.block_size = blockSize_;
283+ tiling.one_core_do_block_num_per_row = perCoreDoBlockNum_;
284+ tiling.tiling_key = tilingKey_;
285+ tiling.last_block_size = lastBlockSize_;
286+ 
287+ auto ret = memcpy_s(rawTilingData->GetData(), rawTilingData->GetCapacity(), &tiling,
288+ sizeof(ApplyAdamWQuantRegbaseTilingData));
289+ OP_CHECK_IF(ret != EOK,
290+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "memcpy_s",
291+ std::to_string(ret).c_str(), "copy tiling data failed"),
292+ return ge::GRAPH_FAILED);
293+ rawTilingData->SetDataSize(sizeof(ApplyAdamWQuantRegbaseTilingData));
294+ 
295+ tilingContext_->SetBlockDim(static_cast<uint32_t>(useNumCore_));
296+ tilingContext_->SetTilingKey(tilingKey_);
297+ size_t* workspaces = tilingContext_->GetWorkspaceSizes(1);
298+ OP_CHECK_NULL_WITH_CONTEXT(tilingContext_, workspaces);
299+ workspaces[0] = 0U;
300+ return ge::GRAPH_SUCCESS;
301+}
302+ 
303+ge::graphStatus ApplyAdamWQuantRegbaseTiling::RunTiling()
304+{
305+ OP_CHECK_IF(
306+ tilingContext_ == nullptr,
307+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("ApplyAdamWQuant", "parameter", "invalid", "tiling context is null"),
308+ return ge::GRAPH_FAILED);
309+ OP_CHECK_IF(GetAttributes() != ge::GRAPH_SUCCESS,
310+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
311+ "GetAttributes failed."),
312+ return ge::GRAPH_FAILED);
313+ OP_CHECK_IF(CheckInputShape() != ge::GRAPH_SUCCESS,
314+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
315+ "input shape check failed."),
316+ return ge::GRAPH_FAILED);
317+ OP_CHECK_IF(DetermineTilingKey() != ge::GRAPH_SUCCESS,
318+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
319+ "determine tiling key failed."),
320+ return ge::GRAPH_FAILED);
321+ OP_CHECK_IF(DoTiling() != ge::GRAPH_SUCCESS,
322+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(tilingContext_->GetNodeName(), "parameter", "invalid",
323+ "DoTiling failed."),
324+ return ge::GRAPH_FAILED);
325+ return SetTilingData();
326+}
327+ 
328+ge::graphStatus Tiling4ApplyAdamWQuant(gert::TilingContext* context)
329+{
330+ if (context == nullptr) {
331+ OP_LOGE_FOR_INVALID_VALUE_WITH_REASON("Tiling4ApplyAdamWQuant", "parameter", "invalid",
332+ "Tiling context is null");
333+ return ge::GRAPH_FAILED;
334+ }
335+ OP_LOGD(context, "Tiling4ApplyAdamWQuant running begin");
336+ ApplyAdamWQuantRegbaseTiling regbaseTiling(context);
337+ return regbaseTiling.RunTiling();
338+}
339+ 
340+static ge::graphStatus TilingPrepare4ApplyAdamWQuant(gert::TilingParseContext* context)
341+{
342+ OP_LOGD(context, "TilingPrepare4ApplyAdamWQuant enter.");
343+ auto compileInfo = context->GetCompiledInfo<ApplyAdamWQuantRegbaseCompileInfo>();
344+ OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
345+ OP_LOGD(context, "TilingPrepare4ApplyAdamWQuant exit.");
346+ return ge::GRAPH_SUCCESS;
347+}
348+ 
349+IMPL_OP_OPTILING(ApplyAdamWQuant)
350+ .Tiling(Tiling4ApplyAdamWQuant)
351+ .TilingParse<ApplyAdamWQuantRegbaseCompileInfo>(TilingPrepare4ApplyAdamWQuant);
352+} // namespace optiling
Aoptim/apply_adam_w_quant/op_host/arch35/apply_adam_w_quant_tiling_arch35.h+63-0
@@ -0,0 +1,63 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ *
10+ * NOTE: Portions of this code were AI-generated and have been
11+ * technically reviewed for functional accuracy and security
12+ */
13+ 
14+/*!
15+ * \file apply_adam_w_quant_tiling_arch35.h
16+ * \brief ApplyAdamWQuant regbase (arch35 / Ascend950) tiling class declaration.
17+ */
18+#ifndef APPLY_ADAM_W_QUANT_ARCH35_TILING_H
19+#define APPLY_ADAM_W_QUANT_ARCH35_TILING_H
20+ 
21+#include "register/op_impl_registry.h"
22+#include "../../op_kernel/arch35/apply_adam_w_quant_tiling_data.h"
23+ 
24+namespace optiling {
25+struct ApplyAdamWQuantRegbaseCompileInfo {};
26+ 
27+class ApplyAdamWQuantRegbaseTiling {
28+public:
29+ explicit ApplyAdamWQuantRegbaseTiling(gert::TilingContext* context) : tilingContext_(context) {}
30+ 
31+ ge::graphStatus RunTiling();
32+ 
33+protected:
34+ ge::graphStatus GetAttributes();
35+ ge::graphStatus CheckInputShape();
36+ ge::graphStatus DetermineTilingKey();
37+ ge::graphStatus DoTiling();
38+ ge::graphStatus SetTilingData();
39+ 
40+private:
41+ gert::TilingContext* tilingContext_ = nullptr;
42+ 
43+ // attrs
44+ float lr_ = 0.0f;
45+ float beta1_ = 0.0f;
46+ float beta2_ = 0.0f;
47+ float weightDecay_ = 0.0f;
48+ float eps_ = 0.0f;
49+ float gnormScale_ = 0.0f;
50+ int64_t blockSize_ = 0;
51+ 
52+ // tiling result
53+ uint64_t tilingKey_ = 0;
54+ uint64_t useNumCore_ = 0;
55+ uint64_t lastPreCoreRowWork_ = 0;
56+ uint64_t notLastCoreNum_ = 0;
57+ uint64_t notLastPreCoreRowWork_ = 0;
58+ uint64_t lastCoreLastBlock_ = 0;
59+ uint64_t lastBlockSize_ = 0;
60+ uint64_t perCoreDoBlockNum_ = 0;
61+};
62+} // namespace optiling
63+#endif // APPLY_ADAM_W_QUANT_ARCH35_TILING_H
Aoptim/apply_adam_w_quant/op_host/config/ascend950/apply_adam_w_quant_binary.json+572-0
@@ -0,0 +1,572 @@
1+{
2+ "op_type": "ApplyAdamWQuant",
3+ "op_list": [
4+ {
5+ "bin_filename": "ApplyAdamWQuant_d9384f4e",
6+ "inputs": [
7+ {
8+ "name": "var",
9+ "index": 0,
10+ "dtype": "float32",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ },
17+ {
18+ "name": "grad",
19+ "index": 1,
20+ "dtype": "float32",
21+ "format": "ND",
22+ "paramType": "required",
23+ "shape": [
24+ -2
25+ ]
26+ },
27+ {
28+ "name": "m",
29+ "index": 2,
30+ "dtype": "uint8",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [
34+ -2
35+ ]
36+ },
37+ {
38+ "name": "v",
39+ "index": 3,
40+ "dtype": "uint8",
41+ "format": "ND",
42+ "paramType": "required",
43+ "shape": [
44+ -2
45+ ]
46+ },
47+ {
48+ "name": "qmap_m",
49+ "index": 4,
50+ "dtype": "float32",
51+ "format": "ND",
52+ "paramType": "required",
53+ "shape": [
54+ -2
55+ ]
56+ },
57+ {
58+ "name": "qmap_v",
59+ "index": 5,
60+ "dtype": "float32",
61+ "format": "ND",
62+ "paramType": "required",
63+ "shape": [
64+ -2
65+ ]
66+ },
67+ {
68+ "name": "absmax_m",
69+ "index": 6,
70+ "dtype": "float32",
71+ "format": "ND",
72+ "paramType": "required",
73+ "shape": [
74+ -2
75+ ]
76+ },
77+ {
78+ "name": "absmax_v",
79+ "index": 7,
80+ "dtype": "float32",
81+ "format": "ND",
82+ "paramType": "required",
83+ "shape": [
84+ -2
85+ ]
86+ },
87+ {
88+ "name": "step",
89+ "index": 8,
90+ "dtype": "int64",
91+ "format": "ND",
92+ "paramType": "required",
93+ "shape": [
94+ -2
95+ ]
96+ }
97+ ],
98+ "outputs": [
99+ {
100+ "name": "var",
101+ "index": 0,
102+ "dtype": "float32",
103+ "format": "ND",
104+ "paramType": "required",
105+ "shape": [
106+ -2
107+ ]
108+ },
109+ {
110+ "name": "m",
111+ "index": 1,
112+ "dtype": "uint8",
113+ "format": "ND",
114+ "paramType": "required",
115+ "shape": [
116+ -2
117+ ]
118+ },
119+ {
120+ "name": "v",
121+ "index": 2,
122+ "dtype": "uint8",
123+ "format": "ND",
124+ "paramType": "required",
125+ "shape": [
126+ -2
127+ ]
128+ },
129+ {
130+ "name": "absmax_m",
131+ "index": 3,
132+ "dtype": "float32",
133+ "format": "ND",
134+ "paramType": "required",
135+ "shape": [
136+ -2
137+ ]
138+ },
139+ {
140+ "name": "absmax_v",
141+ "index": 4,
142+ "dtype": "float32",
143+ "format": "ND",
144+ "paramType": "required",
145+ "shape": [
146+ -2
147+ ]
148+ }
149+ ],
150+ "attrs": [
151+ {
152+ "name": "lr",
153+ "dtype": "float",
154+ "value": 0.1
155+ },
156+ {
157+ "name": "beta1",
158+ "dtype": "float",
159+ "value": 0.1
160+ },
161+ {
162+ "name": "beta2",
163+ "dtype": "float",
164+ "value": 0.1
165+ },
166+ {
167+ "name": "weight_decay",
168+ "dtype": "float",
169+ "value": 0.1
170+ },
171+ {
172+ "name": "eps",
173+ "dtype": "float",
174+ "value": 0.1
175+ },
176+ {
177+ "name": "gnorm_scale",
178+ "dtype": "float",
179+ "value": 0.1
180+ },
181+ {
182+ "name": "quant_mode",
183+ "dtype": "string",
184+ "value": null
185+ },
186+ {
187+ "name": "block_size",
188+ "dtype": "int64",
189+ "value": 256
190+ }
191+ ]
192+ },
193+ {
194+ "bin_filename": "ApplyAdamWQuant_896fb64e",
195+ "inputs": [
196+ {
197+ "name": "var",
198+ "index": 0,
199+ "dtype": "float16",
200+ "format": "ND",
201+ "paramType": "required",
202+ "shape": [
203+ -2
204+ ]
205+ },
206+ {
207+ "name": "grad",
208+ "index": 1,
209+ "dtype": "float16",
210+ "format": "ND",
211+ "paramType": "required",
212+ "shape": [
213+ -2
214+ ]
215+ },
216+ {
217+ "name": "m",
218+ "index": 2,
219+ "dtype": "uint8",
220+ "format": "ND",
221+ "paramType": "required",
222+ "shape": [
223+ -2
224+ ]
225+ },
226+ {
227+ "name": "v",
228+ "index": 3,
229+ "dtype": "uint8",
230+ "format": "ND",
231+ "paramType": "required",
232+ "shape": [
233+ -2
234+ ]
235+ },
236+ {
237+ "name": "qmap_m",
238+ "index": 4,
239+ "dtype": "float32",
240+ "format": "ND",
241+ "paramType": "required",
242+ "shape": [
243+ -2
244+ ]
245+ },
246+ {
247+ "name": "qmap_v",
248+ "index": 5,
249+ "dtype": "float32",
250+ "format": "ND",
251+ "paramType": "required",
252+ "shape": [
253+ -2
254+ ]
255+ },
256+ {
257+ "name": "absmax_m",
258+ "index": 6,
259+ "dtype": "float32",
260+ "format": "ND",
261+ "paramType": "required",
262+ "shape": [
263+ -2
264+ ]
265+ },
266+ {
267+ "name": "absmax_v",
268+ "index": 7,
269+ "dtype": "float32",
270+ "format": "ND",
271+ "paramType": "required",
272+ "shape": [
273+ -2
274+ ]
275+ },
276+ {
277+ "name": "step",
278+ "index": 8,
279+ "dtype": "int64",
280+ "format": "ND",
281+ "paramType": "required",
282+ "shape": [
283+ -2
284+ ]
285+ }
286+ ],
287+ "outputs": [
288+ {
289+ "name": "var",
290+ "index": 0,
291+ "dtype": "float16",
292+ "format": "ND",
293+ "paramType": "required",
294+ "shape": [
295+ -2
296+ ]
297+ },
298+ {
299+ "name": "m",
300+ "index": 1,
301+ "dtype": "uint8",
302+ "format": "ND",
303+ "paramType": "required",
304+ "shape": [
305+ -2
306+ ]
307+ },
308+ {
309+ "name": "v",
310+ "index": 2,
311+ "dtype": "uint8",
312+ "format": "ND",
313+ "paramType": "required",
314+ "shape": [
315+ -2
316+ ]
317+ },
318+ {
319+ "name": "absmax_m",
320+ "index": 3,
321+ "dtype": "float32",
322+ "format": "ND",
323+ "paramType": "required",
324+ "shape": [
325+ -2
326+ ]
327+ },
328+ {
329+ "name": "absmax_v",
330+ "index": 4,
331+ "dtype": "float32",
332+ "format": "ND",
333+ "paramType": "required",
334+ "shape": [
335+ -2
336+ ]
337+ }
338+ ],
339+ "attrs": [
340+ {
341+ "name": "lr",
342+ "dtype": "float",
343+ "value": 0.1
344+ },
345+ {
346+ "name": "beta1",
347+ "dtype": "float",
348+ "value": 0.1
349+ },
350+ {
351+ "name": "beta2",
352+ "dtype": "float",
353+ "value": 0.1
354+ },
355+ {
356+ "name": "weight_decay",
357+ "dtype": "float",
358+ "value": 0.1
359+ },
360+ {
361+ "name": "eps",
362+ "dtype": "float",
363+ "value": 0.1
364+ },
365+ {
366+ "name": "gnorm_scale",
367+ "dtype": "float",
368+ "value": 0.1
369+ },
370+ {
371+ "name": "quant_mode",
372+ "dtype": "string",
373+ "value": null
374+ },
375+ {
376+ "name": "block_size",
377+ "dtype": "int64",
378+ "value": 256
379+ }
380+ ]
381+ },
382+ {
383+ "bin_filename": "ApplyAdamWQuant_b14a102b",
384+ "inputs": [
385+ {
386+ "name": "var",
387+ "index": 0,
388+ "dtype": "bfloat16",
389+ "format": "ND",
390+ "paramType": "required",
391+ "shape": [
392+ -2
393+ ]
394+ },
395+ {
396+ "name": "grad",
397+ "index": 1,
398+ "dtype": "bfloat16",
399+ "format": "ND",
400+ "paramType": "required",
401+ "shape": [
402+ -2
403+ ]
404+ },
405+ {
406+ "name": "m",
407+ "index": 2,
408+ "dtype": "uint8",
409+ "format": "ND",
410+ "paramType": "required",
411+ "shape": [
412+ -2
413+ ]
414+ },
415+ {
416+ "name": "v",
417+ "index": 3,
418+ "dtype": "uint8",
419+ "format": "ND",
420+ "paramType": "required",
421+ "shape": [
422+ -2
423+ ]
424+ },
425+ {
426+ "name": "qmap_m",
427+ "index": 4,
428+ "dtype": "float32",
429+ "format": "ND",
430+ "paramType": "required",
431+ "shape": [
432+ -2
433+ ]
434+ },
435+ {
436+ "name": "qmap_v",
437+ "index": 5,
438+ "dtype": "float32",
439+ "format": "ND",
440+ "paramType": "required",
441+ "shape": [
442+ -2
443+ ]
444+ },
445+ {
446+ "name": "absmax_m",
447+ "index": 6,
448+ "dtype": "float32",
449+ "format": "ND",
450+ "paramType": "required",
451+ "shape": [
452+ -2
453+ ]
454+ },
455+ {
456+ "name": "absmax_v",
457+ "index": 7,
458+ "dtype": "float32",
459+ "format": "ND",
460+ "paramType": "required",
461+ "shape": [
462+ -2
463+ ]
464+ },
465+ {
466+ "name": "step",
467+ "index": 8,
468+ "dtype": "int64",
469+ "format": "ND",
470+ "paramType": "required",
471+ "shape": [
472+ -2
473+ ]
474+ }
475+ ],
476+ "outputs": [
477+ {
478+ "name": "var",
479+ "index": 0,
480+ "dtype": "bfloat16",
481+ "format": "ND",
482+ "paramType": "required",
483+ "shape": [
484+ -2
485+ ]
486+ },
487+ {
488+ "name": "m",
489+ "index": 1,
490+ "dtype": "uint8",
491+ "format": "ND",
492+ "paramType": "required",
493+ "shape": [
494+ -2
495+ ]
496+ },
497+ {
498+ "name": "v",
499+ "index": 2,
500+ "dtype": "uint8",
501+ "format": "ND",
502+ "paramType": "required",
503+ "shape": [
504+ -2
505+ ]
506+ },
507+ {
508+ "name": "absmax_m",
509+ "index": 3,
510+ "dtype": "float32",
511+ "format": "ND",
512+ "paramType": "required",
513+ "shape": [
514+ -2
515+ ]
516+ },
517+ {
518+ "name": "absmax_v",
519+ "index": 4,
520+ "dtype": "float32",
521+ "format": "ND",
522+ "paramType": "required",
523+ "shape": [
524+ -2
525+ ]
526+ }
527+ ],
528+ "attrs": [
529+ {
530+ "name": "lr",
531+ "dtype": "float",
532+ "value": 0.1
533+ },
534+ {
535+ "name": "beta1",
536+ "dtype": "float",
537+ "value": 0.1
538+ },
539+ {
540+ "name": "beta2",
541+ "dtype": "float",
542+ "value": 0.1
543+ },
544+ {
545+ "name": "weight_decay",
546+ "dtype": "float",
547+ "value": 0.1
548+ },
549+ {
550+ "name": "eps",
551+ "dtype": "float",
552+ "value": 0.1
553+ },
554+ {
555+ "name": "gnorm_scale",
556+ "dtype": "float",
557+ "value": 0.1
558+ },
559+ {
560+ "name": "quant_mode",
561+ "dtype": "string",
562+ "value": null
563+ },
564+ {
565+ "name": "block_size",
566+ "dtype": "int64",
567+ "value": 256
568+ }
569+ ]
570+ }
571+ ]
572+}
Moptim/apply_adam_w_quant/op_kernel/CMakeLists.txt+7-2
@@ -1,9 +1,9 @@
1# ----------------------------------------------------------------------------1# ----------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# This program is free software, you can redistribute it and/or modify it under the terms and conditions of 3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4# CANN Open Software License Agreement Version 2.0 (the "License").4# CANN Open Software License Agreement Version 2.0 (the "License").
5# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
6-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# ----------------------------------------------------------------------------9# ----------------------------------------------------------------------------
@@ -12,3 +12,8 @@
12add_kernel_sources(12add_kernel_sources(
13 COMPUTE_UNITS ascend910b ascend910_9313 COMPUTE_UNITS ascend910b ascend910_93
14)14)
15+add_kernel_sources(
16+ KERNEL_SRC arch35/apply_adam_w_quant.cpp
17+ COMPUTE_UNITS ascend950
18+ AUTO_SYNC false
19+)
Aoptim/apply_adam_w_quant/op_kernel/arch35/apply_adam_w_quant.cpp+55-0
@@ -0,0 +1,55 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ *
10+ * NOTE: Portions of this code were AI-generated and have been
11+ * technically reviewed for functional accuracy and security
12+ */
13+ 
14+/*!
15+ * \file apply_adam_w_quant.cpp
16+ * \brief ApplyAdamWQuant arch35 (Ascend950 / DAV_3510) kernel entry.
17+ *
18+ * arch35 移植:计算逻辑与 A2 完全一致(blockwise-256 量化 AdamW),复用同族 A2 的
19+ * 手写 TPipe/TBuf + 高层 Vec API 计算头(base/fp32/fp16),仅 tiling 数据从 framework
20+ * GET_TILING_DATA 换成 regbase 的 plain-POD GET_TILING_DATA_WITH_STRUCT。dtype 分发沿用
21+ * A2 的 TilingKey:fp32=100 / fp16=200 / bf16=300。kernel 入口风格对齐 norm/deep_norm。
22+ */
23+#include "kernel_operator.h"
24+#include "apply_adam_w_quant_tiling_data.h"
25+#include "apply_adam_w_quant_fp32.h"
26+#include "apply_adam_w_quant_fp16.h"
27+ 
28+using namespace ApplyAdamWQuantNS;
29+ 
30+extern "C" __global__ __aicore__ void apply_adam_w_quant(GM_ADDR var, GM_ADDR grad, GM_ADDR m, GM_ADDR v,
31+ GM_ADDR qmap_m, GM_ADDR qmap_v, GM_ADDR absmax_m,
32+ GM_ADDR absmax_v, GM_ADDR step, GM_ADDR var_ref, GM_ADDR m_ref,
33+ GM_ADDR v_ref, GM_ADDR absmax_m_ref, GM_ADDR absmax_v_ref,
34+ GM_ADDR workspace, GM_ADDR tiling)
35+{
36+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
37+ REGISTER_TILING_DEFAULT(ApplyAdamWQuantRegbaseTilingData);
38+ GET_TILING_DATA_WITH_STRUCT(ApplyAdamWQuantRegbaseTilingData, tiling_data_in, tiling);
39+ if (TILING_KEY_IS(100)) {
TangPC
TangPCTangPC8月1日

这个分发链没有兜底分支，tiling key 不是 100/200/300 时 kernel 什么都不做就正常返回，输出保持为未初始化的 GM 内容，从外面看就是「算子跑成功了但结果不对」。tiling 侧虽然已经限死了三种 dtype，但两边一旦改歪就是这种最难查的表现。建议补一个 else，至少在 CPU 调试态下 assert 或打一条日志。

likedislike
40+ ApplyAdamWQuant<float, int64_t> op;
41+ op.Init(var, grad, m, v, qmap_m, qmap_v, absmax_m, absmax_v, step, var_ref, m_ref, v_ref, absmax_m_ref,
42+ absmax_v_ref, &tiling_data_in);
43+ op.Process();
44+ } else if (TILING_KEY_IS(200)) {
45+ ApplyAdamWQuant16<float, int64_t, half> op;
46+ op.Init(var, grad, m, v, qmap_m, qmap_v, absmax_m, absmax_v, step, var_ref, m_ref, v_ref, absmax_m_ref,
47+ absmax_v_ref, &tiling_data_in);
48+ op.Process();
49+ } else if (TILING_KEY_IS(300)) {
50+ ApplyAdamWQuant16<float, int64_t, bfloat16_t> op;
51+ op.Init(var, grad, m, v, qmap_m, qmap_v, absmax_m, absmax_v, step, var_ref, m_ref, v_ref, absmax_m_ref,
52+ absmax_v_ref, &tiling_data_in);
53+ op.Process();
54+ }
55+}
Aoptim/apply_adam_w_quant/op_kernel/arch35/apply_adam_w_quant_base.h+146-0
@@ -0,0 +1,146 @@
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 apply_adam_w_quant_base.h
13+ * \brief
14+ */
15+#ifndef _APPLY_ADAM_W_QUANT_BASE_H_
16+#define _APPLY_ADAM_W_QUANT_BASE_H_
17+ 
18+#include "kernel_operator.h"
19+#include "apply_adam_w_quant_tiling_data.h"
20+ 
21+namespace ApplyAdamWQuantNS {
22+using namespace AscendC;
23+ 
24+constexpr int32_t Q_MAP_SIZE = 256;
25+constexpr uint32_t CALC_BUF_NUM = 6;
26+constexpr uint32_t PER_UINT8_8BITS = 8;
27+constexpr uint32_t REPEAT_NUM = 64;
28+constexpr uint32_t REPEAT_7_TIMES = 7;
29+constexpr uint32_t REPEAT_NUM_128 = 128;
30+constexpr uint32_t STRIDE_8 = 8;
31+constexpr uint32_t PER_4NUM_ONEMAX = 4;
32+constexpr uint32_t BROADCAST_DIM2 = 2;
33+constexpr uint32_t BROADCAST_AXIS1 = 1;
34+ 
35+template <typename T>
36+__aicore__ inline void DataCopyIn(const AscendC::LocalTensor<T>& dst, const AscendC::GlobalTensor<T>& src,
37+ uint32_t count)
38+{
39+ AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
40+ AscendC::DataCopyPadExtParams<T> padParams{false, 0, 0, 0};
41+ AscendC::DataCopyPad(dst, src, copyParams, padParams);
42+}
43+ 
44+template <typename T>
45+__aicore__ inline void DataCopyOut(const AscendC::GlobalTensor<T>& dst, const AscendC::LocalTensor<T>& src,
46+ uint32_t count)
47+{
48+ AscendC::DataCopyExtParams copyParams{1, static_cast<uint32_t>(count * sizeof(T)), 0, 0, 0};
49+ 
50+ AscendC::DataCopyPad(dst, src, copyParams);
51+}
52+ 
53+template <typename T, typename T1>
54+__aicore__ inline void CastF16ToFp32(const AscendC::LocalTensor<T>& dst, const AscendC::LocalTensor<T1>& src,
55+ uint32_t count)
56+{
57+ AscendC::Cast(dst, src, AscendC::RoundMode::CAST_NONE, count);
58+}
59+ 
60+template <typename T, typename T1>
61+__aicore__ inline void CastFp32ToF16(const AscendC::LocalTensor<T>& dst, const AscendC::LocalTensor<T1>& src,
62+ uint32_t count)
63+{
64+ if constexpr (AscendC::IsSameType<T, half>::value) {
65+ Cast(dst, src, AscendC::RoundMode::CAST_NONE, count);
66+ } else { // bf16
67+ Cast(dst, src, AscendC::RoundMode::CAST_RINT, count);
68+ }
69+}
70+ 
71+template <AscendC::HardEvent hardEvent>
72+__aicore__ inline void PipeSync()
73+{
74+ int32_t eventID = static_cast<int32_t>(GetTPipePtr()->FetchEventID(hardEvent));
75+ AscendC::SetFlag<hardEvent>(eventID);
76+ AscendC::WaitFlag<hardEvent>(eventID);
77+}
78+ 
79+__aicore__ inline float PowS(const AscendC::LocalTensor<float>& dst, float srcScalar,
80+ const AscendC::LocalTensor<float>& src)
81+{
82+ AscendC::Power<float>(dst, srcScalar, src);
83+ PipeSync<AscendC::HardEvent::V_S>();
84+ float ret = dst.GetValue(0);
85+ AscendC::PipeBarrier<PIPE_ALL>();
86+ return ret;
87+}
88+ 
89+// RegBase VF：AdamW 状态/参数逐元素更新（fp32 寄存器计算 body）。
90+// fp32 与 fp16 路径共用——两者进入本函数时 m/v/grad/var 均为 fp32 UB tile。
91+// 由 UB 层 UpdateStateAndParam(Fp32) 经 asc_vf_call 调用，替代原 tile-API Muls/Mul/Add/Sqrt/Adds/Div 链。
92+// grad' = grad * gnormScale
93+// m_out = m * beta1 + grad' * (1 - beta1)
94+// v_out = v * beta2 + grad'^2 * (1 - beta2)
95+// var_out= (var + (m_out * stepSize) / (sqrt(v_out) + eps*correction2)) * wdFactor
96+// wdFactor = (weight_decay>0) ? (1 - lr*weight_decay) : 1.0f（乘 1.0 为 no-op，等价于原分支跳过）
97+template <typename T>
98+__simd_vf__ inline void AdamWQuantUpdateVF(__ubuf__ T* dqStateMAddr, __ubuf__ T* dqStateVAddr, __ubuf__ T* varAddr,
TangPC
TangPCTangPC8月1日

这个 VF 函数全仓没有任何调用点，fp32.h 的 UpdateStateAndParamFp32 和 fp16.h 的 UpdateStateAndParam 用的都是 Muls/Mul/Add/Sqrt/Adds/Div 这条 tile-API 链，恰恰是上面注释里说「已被替代」的那条。所以现在是 45 行死代码加一段与事实相反的注释，后面接手的人很容易以为计算走的是 VF 路径。要么真的把 UB 层切到 asc_vf_call 走这个 VF body（arch35 上这才是推荐写法，同目录的 adam_apply_one_assign、inplace_apply_centered_rms_prop 都是这么做的），要么把它和注释一起删掉。

likedislike
99+ __ubuf__ T* gradAddr, float gnormScale, float beta1, float oneMinusBeta1,
100+ float beta2, float oneMinusBeta2, float epsCorr, float stepSize,
101+ float wdFactor, uint32_t count, uint32_t oneRepeatSize, uint16_t repeatTimes)
102+{
103+ using namespace AscendC::Reg;
104+ RegTensor<T> regM, regV, regGrad, regVar, regTmp;
105+ MaskReg mask;
106+ AddrReg aReg;
107+ uint32_t remain = count;
108+ for (uint16_t i = 0; i < repeatTimes; ++i) {
109+ aReg = CreateAddrReg<T>(i, oneRepeatSize);
110+ mask = UpdateMask<T>(remain);
111+ 
112+ LoadAlign(regM, dqStateMAddr, aReg);
113+ LoadAlign(regV, dqStateVAddr, aReg);
114+ LoadAlign(regVar, varAddr, aReg);
115+ LoadAlign(regGrad, gradAddr, aReg);
116+ 
117+ // grad = grad * gnormScale
118+ Muls(regGrad, regGrad, gnormScale, mask);
119+ // m = m*beta1 + grad*(1-beta1)
120+ Muls(regM, regM, beta1, mask);
121+ Muls(regTmp, regGrad, oneMinusBeta1, mask);
122+ Add(regM, regM, regTmp, mask);
123+ // v = v*beta2 + grad^2*(1-beta2)
124+ Muls(regV, regV, beta2, mask);
125+ Mul(regGrad, regGrad, regGrad, mask); // grad^2
126+ Muls(regTmp, regGrad, oneMinusBeta2, mask);
127+ Add(regV, regV, regTmp, mask);
128+ // denom = sqrt(v) + eps*correction2
129+ Sqrt(regTmp, regV, mask);
130+ Adds(regTmp, regTmp, epsCorr, mask);
131+ // var = var + (m*stepSize)/denom（regGrad 复用为 m*stepSize，对齐原 tmpVar1 别名 grad 缓存）
132+ Muls(regGrad, regM, stepSize, mask);
133+ Div(regGrad, regGrad, regTmp, mask);
134+ Add(regVar, regVar, regGrad, mask);
135+ // var = var * wdFactor（无 weight_decay 时 wdFactor=1）
136+ Muls(regVar, regVar, wdFactor, mask);
137+ 
138+ StoreAlign(dqStateMAddr, regM, aReg, mask);
139+ StoreAlign(dqStateVAddr, regV, aReg, mask);
140+ StoreAlign(varAddr, regVar, aReg, mask);
141+ }
142+}
143+ 
144+} // namespace ApplyAdamWQuantNS
145+ 
146+#endif // _APPLY_ADAM_W_QUANT_BASE_H_
Aoptim/apply_adam_w_quant/op_kernel/arch35/apply_adam_w_quant_fp16.h+718-0
@@ -0,0 +1,718 @@
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 apply_adam_w_quant_fp16.h
13+ * \brief
14+ */
15+#ifndef APPLY_ADAM_W_QUANT_FP16_H_
16+#define APPLY_ADAM_W_QUANT_FP16_H_
17+ 
18+#include "apply_adam_w_quant_base.h"
19+ 
20+namespace ApplyAdamWQuantNS {
21+using namespace AscendC;
22+ 
23+using namespace ApplyAdamWQuantNS;
24+ 
25+template <typename T, typename U, typename T_VAR_GRAD>
26+class ApplyAdamWQuant16 {
27+public:
28+ __aicore__ inline ApplyAdamWQuant16(){};
29+ __aicore__ inline void Init(GM_ADDR var, GM_ADDR grad, GM_ADDR m, GM_ADDR v, GM_ADDR qmap_m, GM_ADDR qmap_v,
30+ GM_ADDR absmax_m, GM_ADDR absmax_v, GM_ADDR step, GM_ADDR var_ref, GM_ADDR m_ref,
31+ GM_ADDR v_ref, GM_ADDR absmax_m_ref, GM_ADDR absmax_v_ref,
32+ const ApplyAdamWQuantRegbaseTilingData* tilingData)
33+ {
34+ this->blockIdx = AscendC::GetBlockIdx();
35+ ParseTilingData(tilingData);
36+ if (this->blockIdx < notLastCoreNum) {
37+ gmOffset = this->blockIdx * notLastPreCoreRowWork * blockSize * oneCoreDoBlockNumPerRow;
38+ absmaxOffset = this->blockIdx * notLastPreCoreRowWork * oneCoreDoBlockNumPerRow;
39+ } else {
40+ gmOffset = (notLastCoreNum * notLastPreCoreRowWork +
41+ (this->blockIdx - notLastCoreNum) * lastPreCoreRowWork) *
42+ blockSize * oneCoreDoBlockNumPerRow;
43+ absmaxOffset = (notLastCoreNum * notLastPreCoreRowWork +
44+ (this->blockIdx - notLastCoreNum) * lastPreCoreRowWork) *
45+ oneCoreDoBlockNumPerRow;
46+ }
47+ this->singleBlockNum = oneCoreDoBlockNumPerRow;
48+ this->singleSize = blockSize * oneCoreDoBlockNumPerRow;
49+ varGm.SetGlobalBuffer((__gm__ T_VAR_GRAD*)var + gmOffset);
50+ gradGm.SetGlobalBuffer((__gm__ T_VAR_GRAD*)grad + gmOffset);
51+ stateMGm.SetGlobalBuffer((__gm__ uint8_t*)m + gmOffset);
52+ stateVGm.SetGlobalBuffer((__gm__ uint8_t*)v + gmOffset);
53+ qMapMGm.SetGlobalBuffer((__gm__ T*)qmap_m);
54+ qMapVGm.SetGlobalBuffer((__gm__ T*)qmap_v);
55+ absMaxMGm.SetGlobalBuffer((__gm__ T*)absmax_m + absmaxOffset);
56+ absMaxVGm.SetGlobalBuffer((__gm__ T*)absmax_v + absmaxOffset);
57+ stepGm.SetGlobalBuffer((__gm__ U*)step, 1);
58+ varRefGm.SetGlobalBuffer((__gm__ T_VAR_GRAD*)var_ref + gmOffset);
59+ stateMRefGm.SetGlobalBuffer((__gm__ uint8_t*)m_ref + gmOffset);
60+ stateVRefGm.SetGlobalBuffer((__gm__ uint8_t*)v_ref + gmOffset);
61+ absMaxMRefGm.SetGlobalBuffer((__gm__ T*)absmax_m_ref + absmaxOffset);
62+ absMaxVRefGm.SetGlobalBuffer((__gm__ T*)absmax_v_ref + absmaxOffset);
63+ step_ = static_cast<float>(stepGm.GetValue(0));
64+ 
65+ pipe.InitBuffer(varBuf, singleSize * sizeof(T_VAR_GRAD));
66+ pipe.InitBuffer(gradBuf, singleSize * sizeof(T_VAR_GRAD));
67+ pipe.InitBuffer(varFp32Buf, singleSize * sizeof(T));
68+ pipe.InitBuffer(gradFp32Buf, singleSize * sizeof(T));
69+ pipe.InitBuffer(stateMBuf, singleSize * sizeof(T));
70+ pipe.InitBuffer(stateVBuf, singleSize * sizeof(T));
71+ pipe.InitBuffer(qMapMBuf, Q_MAP_SIZE * sizeof(T));
72+ pipe.InitBuffer(qMapVBuf, Q_MAP_SIZE * sizeof(T));
73+ pipe.InitBuffer(absMaxMBuf, oneCoreDoBlockNumPerRow * sizeof(T));
74+ pipe.InitBuffer(absMaxVBuf, oneCoreDoBlockNumPerRow * sizeof(T));
75+ pipe.InitBuffer(calcBuf, CALC_BUF_NUM * singleSize * sizeof(float));
76+ 
77+ this->qMapM = qMapMBuf.Get<T>();
78+ this->qMapV = qMapVBuf.Get<T>();
79+ this->stateM = stateMBuf.Get<uint8_t>();
80+ this->stateV = stateVBuf.Get<uint8_t>();
81+ this->absMaxM = absMaxMBuf.Get<T>();
82+ this->absMaxV = absMaxVBuf.Get<T>();
83+ this->var = varBuf.Get<T_VAR_GRAD>();
84+ this->grad = gradBuf.Get<T_VAR_GRAD>();
85+ this->varFp32 = varFp32Buf.Get<T>();
86+ this->gradFp32 = gradFp32Buf.Get<T>();
87+ 
88+ step_ += 1;
89+ uint32_t count = AscendC::ONE_BLK_SIZE / sizeof(float);
90+ AscendC::LocalTensor<float> stepTensor = calcBuf.GetWithOffset<float>(count, 0);
91+ AscendC::LocalTensor<float> resTensor = calcBuf.GetWithOffset<float>(count, AscendC::ONE_BLK_SIZE);
92+ AscendC::Duplicate<float>(stepTensor, step_, AscendC::ONE_BLK_SIZE / sizeof(float));
93+ AscendC::PipeBarrier<PIPE_V>();
94+ 
95+ float correction1 = 1 - PowS(resTensor, beta1, stepTensor);
96+ this->correction2 = sqrt(1 - PowS(resTensor, beta2, stepTensor));
97+ this->stepSize = -lr * this->correction2 / correction1;
98+ }
99+ 
100+ __aicore__ inline void Process()
101+ {
102+ DataCopyIn<T>(qMapM, qMapMGm, Q_MAP_SIZE);
103+ DataCopyIn<T>(qMapV, qMapVGm, Q_MAP_SIZE);
104+ 
105+ int64_t perCoreRowWork = lastPreCoreRowWork;
106+ if (this->blockIdx < notLastCoreNum) {
107+ perCoreRowWork += 1;
108+ }
109+ uint64_t blockOffset = 0;
110+ uint64_t absMaxBlockOffset = 0;
111+ for (int64_t n = 0; n < perCoreRowWork; n++) {
112+ bool isLastRow = static_cast<uint64_t>(this->blockIdx + 1) == useNumCore && n + 1 == perCoreRowWork;
113+ singleBlockNum = isLastRow ? lastCoreLastBlock : oneCoreDoBlockNumPerRow;
114+ singleSize = blockSize * singleBlockNum;
115+ uint32_t validSize = isLastRow ? singleSize - blockSize + lastBlockSize : singleSize;
116+ 
117+ DataCopyIn<uint8_t>(stateM, stateMGm[blockOffset], validSize);
118+ DataCopyIn<uint8_t>(stateV, stateVGm[blockOffset], validSize);
119+ 
120+ DataCopyIn<T>(absMaxM, absMaxMGm[absMaxBlockOffset], singleBlockNum);
121+ DataCopyIn<T>(absMaxV, absMaxVGm[absMaxBlockOffset], singleBlockNum);
122+ 
123+ PipeSync<AscendC::HardEvent::MTE2_V>();
124+ 
125+ DeQuant(stateM, qMapM, absMaxM, calcBuf, singleBlockNum, validSize);
126+ DeQuant(stateV, qMapV, absMaxV, calcBuf, singleBlockNum, validSize);
127+ DataCopyIn<T_VAR_GRAD>(var, varGm[blockOffset], validSize);
128+ DataCopyIn<T_VAR_GRAD>(grad, gradGm[blockOffset], validSize);
129+ PipeSync<AscendC::HardEvent::MTE2_V>();
130+ CastF16ToFp32<T, T_VAR_GRAD>(varFp32, var, validSize);
131+ CastF16ToFp32<T, T_VAR_GRAD>(gradFp32, grad, validSize);
132+ AscendC::PipeBarrier<PIPE_V>();
133+ UpdateStateAndParam(stateM.template ReinterpretCast<T>(), stateV.template ReinterpretCast<T>(), gradFp32,
134+ varFp32, validSize);
135+ AscendC::PipeBarrier<PIPE_V>();
136+ CastFp32ToF16<T_VAR_GRAD, T>(var, varFp32, validSize);
137+ PipeSync<AscendC::HardEvent::V_MTE3>();
138+ DataCopyOut<T_VAR_GRAD>(varRefGm[blockOffset], var, validSize);
139+ QuantM(stateM.template ReinterpretCast<T>(), qMapM, stateM, absMaxM, gradFp32, calcBuf, singleBlockNum);
140+ QuantV(stateV.template ReinterpretCast<T>(), qMapV, stateV, absMaxV, calcBuf, singleBlockNum);
141+ PipeSync<AscendC::HardEvent::V_MTE3>();
142+ 
143+ DataCopyOut<T>(absMaxMRefGm[absMaxBlockOffset], absMaxM, singleBlockNum);
144+ DataCopyOut<T>(absMaxVRefGm[absMaxBlockOffset], absMaxV, singleBlockNum);
145+ 
146+ DataCopyOut<uint8_t>(stateMRefGm[blockOffset], stateM, validSize);
147+ DataCopyOut<uint8_t>(stateVRefGm[blockOffset], stateV, validSize);
148+ 
149+ PipeSync<AscendC::HardEvent::MTE3_MTE2>();
150+ blockOffset += singleSize;
151+ absMaxBlockOffset += singleBlockNum;
152+ }
153+ }
154+ 
155+private:
156+ __aicore__ inline void ParseTilingData(const ApplyAdamWQuantRegbaseTilingData* tilingData)
157+ {
158+ useNumCore = tilingData->use_num_core;
159+ lastPreCoreRowWork = tilingData->last_pre_core_row_work;
160+ notLastCoreNum = tilingData->not_last_core_num;
161+ notLastPreCoreRowWork = tilingData->not_last_pre_core_row_work;
162+ lastCoreLastBlock = tilingData->last_core_last_block;
163+ lastBlockSize = tilingData->last_block_size;
164+ oneCoreDoBlockNumPerRow = tilingData->one_core_do_block_num_per_row;
165+ 
166+ lr = tilingData->lr;
167+ beta1 = tilingData->beta1;
168+ beta2 = tilingData->beta2;
169+ weightDecay = tilingData->weight_decay;
170+ eps = tilingData->eps;
171+ gnormScale = tilingData->gnorm_scale;
172+ blockSize = tilingData->block_size;
173+ }
174+ 
175+ __aicore__ inline void BinarySearch(const AscendC::LocalTensor<T>& normState, const AscendC::LocalTensor<T>& qMap,
176+ AscendC::LocalTensor<uint8_t>& qState, AscendC::TBuf<>& calcBuf,
177+ int32_t singleBlockNum)
178+ {
179+ uint32_t offset = 0;
180+ AscendC::LocalTensor<int32_t> lowerPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
181+ offset += singleSize * sizeof(int32_t);
182+ AscendC::LocalTensor<int32_t> upperPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
183+ offset += singleSize * sizeof(int32_t);
184+ AscendC::LocalTensor<int32_t> pivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
185+ offset += singleSize * sizeof(int32_t);
186+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
187+ offset += singleSize * sizeof(int32_t);
188+ AscendC::LocalTensor<T> stateTmp = calcBuf.GetWithOffset<T>(singleSize, offset);
189+ offset += singleSize * sizeof(T);
190+ AscendC::LocalTensor<uint8_t> higgerMask = calcBuf.GetWithOffset<uint8_t>(singleSize / PER_UINT8_8BITS, offset);
191+ 
192+ AscendC::Duplicate<int32_t>(lowerPivot, 0, singleSize);
193+ AscendC::Duplicate<int32_t>(upperPivot, Q_MAP_SIZE - 1, singleSize);
194+ AscendC::Duplicate<int32_t>(pivot, (Q_MAP_SIZE - 1) >> 1, singleSize);
195+ 
196+ T midValue = qMap.GetValue((Q_MAP_SIZE - 1) >> 1);
197+ 
198+ PipeSync<AscendC::HardEvent::S_V>();
199+ 
200+ AscendC::Duplicate<T>(stateTmp, midValue, singleSize);
201+ 
202+ AscendC::PipeBarrier<PIPE_V>();
203+ AscendC::Compare<T, uint8_t>(higgerMask, normState, stateTmp, AscendC::CMPMODE::GT, singleSize);
204+ AscendC::PipeBarrier<PIPE_V>();
205+ 
206+ AscendC::Select<float, uint8_t>(
207+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
208+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
209+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
210+ AscendC::Select<float, uint8_t>(
211+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
212+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
213+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
214+ AscendC::PipeBarrier<PIPE_V>();
215+ 
216+ for (int32_t i = 0; i < REPEAT_7_TIMES; ++i) {
217+ AscendC::Add<int32_t>(pivot, lowerPivot, upperPivot, singleSize);
218+ AscendC::PipeBarrier<PIPE_V>();
219+ AscendC::ShiftRight<int32_t>(pivot, pivot, 1, singleSize);
220+ AscendC::PipeBarrier<PIPE_V>();
221+ 
222+ AscendC::Muls<int32_t>(gatherOffset, pivot, sizeof(T), singleSize);
223+ AscendC::PipeBarrier<PIPE_V>();
224+#ifdef ASCENDC_CPU_DEBUG
225+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
226+ AscendC::Gather<T>(stateTmp[i * Q_MAP_SIZE], qMap,
227+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
228+ }
229+#else
230+ AscendC::Gather<T>(stateTmp, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
231+#endif
232+ AscendC::PipeBarrier<PIPE_V>();
233+ 
234+ AscendC::Compare<T, uint8_t>(higgerMask, normState, stateTmp, AscendC::CMPMODE::GT, singleSize);
235+ AscendC::PipeBarrier<PIPE_V>();
236+ 
237+ AscendC::Select<float, uint8_t>(
238+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
239+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
240+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
241+ 
242+ AscendC::Select<float, uint8_t>(
243+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
244+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
245+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
246+ AscendC::PipeBarrier<PIPE_V>();
247+ }
248+ 
249+ AscendC::LocalTensor<T> lowerState = stateTmp;
250+ AscendC::LocalTensor<T> upperState = pivot.template ReinterpretCast<T>();
251+ 
252+ AscendC::Muls<int32_t>(gatherOffset, lowerPivot, sizeof(T), singleSize);
253+ AscendC::PipeBarrier<PIPE_V>();
254+#ifdef ASCENDC_CPU_DEBUG
255+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
256+ AscendC::Gather<T>(lowerState[i * Q_MAP_SIZE], qMap,
257+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
258+ }
259+#else
260+ AscendC::Gather<T>(lowerState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
261+#endif
262+ AscendC::PipeBarrier<PIPE_V>();
263+ AscendC::Sub<T>(lowerState, normState, lowerState, singleSize);
264+ AscendC::PipeBarrier<PIPE_V>();
265+ AscendC::Abs<T>(lowerState, lowerState, singleSize);
266+ 
267+ AscendC::Muls<int32_t>(gatherOffset, upperPivot, sizeof(T), singleSize);
268+ AscendC::PipeBarrier<PIPE_V>();
269+#ifdef ASCENDC_CPU_DEBUG
270+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
271+ AscendC::Gather<T>(upperState[i * Q_MAP_SIZE], qMap,
272+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
273+ }
274+#else
275+ AscendC::Gather<T>(upperState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
276+#endif
277+ AscendC::PipeBarrier<PIPE_V>();
278+ AscendC::Sub<T>(upperState, normState, upperState, singleSize);
279+ AscendC::PipeBarrier<PIPE_V>();
280+ AscendC::Abs<T>(upperState, upperState, singleSize);
281+ AscendC::PipeBarrier<PIPE_V>();
282+ 
283+ AscendC::Compare<T, uint8_t>(higgerMask, lowerState, upperState, AscendC::CMPMODE::LE, singleSize);
284+ AscendC::PipeBarrier<PIPE_V>();
285+ 
286+ AscendC::Select<float, uint8_t>(
287+ pivot.template ReinterpretCast<float>(), higgerMask, lowerPivot.template ReinterpretCast<float>(),
288+ upperPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
289+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
290+ AscendC::PipeBarrier<PIPE_V>();
291+ AscendC::Cast<int16_t, int32_t>(stateTmp.template ReinterpretCast<int16_t>(), pivot,
292+ AscendC::RoundMode::CAST_NONE, singleSize);
293+ AscendC::PipeBarrier<PIPE_V>();
294+ AscendC::Cast<half, int16_t>(pivot.template ReinterpretCast<half>(),
295+ stateTmp.template ReinterpretCast<int16_t>(), AscendC::RoundMode::CAST_NONE,
296+ singleSize);
297+ AscendC::PipeBarrier<PIPE_V>();
298+ AscendC::Cast<uint8_t, half>(qState, pivot.template ReinterpretCast<half>(), AscendC::RoundMode::CAST_NONE,
299+ singleSize);
300+ AscendC::PipeBarrier<PIPE_V>();
301+ }
302+ 
303+ __aicore__ inline void BinarySearchScaled(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& qMap,
304+ const AscendC::LocalTensor<T>& absMax,
305+ AscendC::LocalTensor<uint8_t>& qState, AscendC::TBuf<>& calcBuf,
306+ int32_t singleBlockNum)
307+ {
308+ uint32_t offset = 0;
309+ AscendC::LocalTensor<int32_t> lowerPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
310+ offset += singleSize * sizeof(int32_t);
311+ AscendC::LocalTensor<int32_t> upperPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
312+ offset += singleSize * sizeof(int32_t);
313+ AscendC::LocalTensor<int32_t> pivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
314+ offset += singleSize * sizeof(int32_t);
315+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
316+ offset += singleSize * sizeof(int32_t);
317+ AscendC::LocalTensor<T> stateTmp = calcBuf.GetWithOffset<T>(singleSize, offset);
318+ offset += singleSize * sizeof(T);
319+ AscendC::LocalTensor<uint8_t> higgerMask = calcBuf.GetWithOffset<uint8_t>(singleSize / PER_UINT8_8BITS, offset);
320+ AscendC::LocalTensor<T> absMaxBrcb = gatherOffset.template ReinterpretCast<T>();
321+ 
322+ uint32_t srcShape[2] = {static_cast<uint32_t>(singleBlockNum), 1};
323+ uint32_t dstShape[2] = {static_cast<uint32_t>(singleBlockNum), blockSize};
324+ 
325+ AscendC::Duplicate<int32_t>(lowerPivot, 0, singleSize);
326+ AscendC::Duplicate<int32_t>(upperPivot, Q_MAP_SIZE - 1, singleSize);
327+ AscendC::Duplicate<int32_t>(pivot, (Q_MAP_SIZE - 1) >> 1, singleSize);
328+ 
329+ T midValue = qMap.GetValue((Q_MAP_SIZE - 1) >> 1);
330+ 
331+ PipeSync<AscendC::HardEvent::S_V>();
332+ 
333+ AscendC::Duplicate<T>(stateTmp, midValue, singleSize);
334+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
335+ AscendC::PipeBarrier<PIPE_V>();
336+ AscendC::Mul<T>(stateTmp, stateTmp, absMaxBrcb, singleSize);
337+ AscendC::PipeBarrier<PIPE_V>();
338+ AscendC::Compare<T, uint8_t>(higgerMask, state, stateTmp, AscendC::CMPMODE::GT, singleSize);
339+ AscendC::PipeBarrier<PIPE_V>();
340+ 
341+ AscendC::Select<float, uint8_t>(
342+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
343+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
344+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
345+ AscendC::Select<float, uint8_t>(
346+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
347+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
348+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
349+ AscendC::PipeBarrier<PIPE_V>();
350+ 
351+ for (int32_t i = 0; i < REPEAT_7_TIMES; ++i) {
352+ AscendC::Add<int32_t>(pivot, lowerPivot, upperPivot, singleSize);
353+ AscendC::PipeBarrier<PIPE_V>();
354+ AscendC::ShiftRight<int32_t>(pivot, pivot, 1, singleSize);
355+ AscendC::PipeBarrier<PIPE_V>();
356+ 
357+ AscendC::Muls<int32_t>(gatherOffset, pivot, sizeof(T), singleSize);
358+ AscendC::PipeBarrier<PIPE_V>();
359+#ifdef ASCENDC_CPU_DEBUG
360+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
361+ AscendC::Gather<T>(stateTmp[i * Q_MAP_SIZE], qMap,
362+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
363+ }
364+#else
365+ AscendC::Gather<T>(stateTmp, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
366+#endif
367+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
368+ AscendC::PipeBarrier<PIPE_V>();
369+ AscendC::Mul<T>(stateTmp, stateTmp, absMaxBrcb, singleSize);
370+ AscendC::PipeBarrier<PIPE_V>();
371+ 
372+ AscendC::Compare<T, uint8_t>(higgerMask, state, stateTmp, AscendC::CMPMODE::GT, singleSize);
373+ AscendC::PipeBarrier<PIPE_V>();
374+ 
375+ AscendC::Select<float, uint8_t>(
376+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
377+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
378+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
379+ 
380+ AscendC::Select<float, uint8_t>(
381+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
382+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
383+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
384+ AscendC::PipeBarrier<PIPE_V>();
385+ }
386+ 
387+ AscendC::LocalTensor<T> lowerState = stateTmp;
388+ AscendC::LocalTensor<T> upperState = pivot.template ReinterpretCast<T>();
389+ 
390+ AscendC::Muls<int32_t>(gatherOffset, lowerPivot, sizeof(T), singleSize);
391+ AscendC::PipeBarrier<PIPE_V>();
392+#ifdef ASCENDC_CPU_DEBUG
393+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
394+ AscendC::Gather<T>(lowerState[i * Q_MAP_SIZE], qMap,
395+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
396+ }
397+#else
398+ AscendC::Gather<T>(lowerState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
399+#endif
400+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
401+ AscendC::PipeBarrier<PIPE_V>();
402+ AscendC::Mul<T>(lowerState, lowerState, absMaxBrcb, singleSize);
403+ AscendC::PipeBarrier<PIPE_V>();
404+ AscendC::Sub<T>(lowerState, state, lowerState, singleSize);
405+ AscendC::PipeBarrier<PIPE_V>();
406+ AscendC::Abs<T>(lowerState, lowerState, singleSize);
407+ 
408+ AscendC::Muls<int32_t>(gatherOffset, upperPivot, sizeof(T), singleSize);
409+ AscendC::PipeBarrier<PIPE_V>();
410+#ifdef ASCENDC_CPU_DEBUG
411+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
412+ AscendC::Gather<T>(upperState[i * Q_MAP_SIZE], qMap,
413+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
414+ }
415+#else
416+ AscendC::Gather<T>(upperState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
417+#endif
418+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
419+ AscendC::PipeBarrier<PIPE_V>();
420+ AscendC::Mul<T>(upperState, upperState, absMaxBrcb, singleSize);
421+ AscendC::PipeBarrier<PIPE_V>();
422+ AscendC::Sub<T>(upperState, state, upperState, singleSize);
423+ AscendC::PipeBarrier<PIPE_V>();
424+ AscendC::Abs<T>(upperState, upperState, singleSize);
425+ AscendC::PipeBarrier<PIPE_V>();
426+ 
427+ AscendC::Compare<T, uint8_t>(higgerMask, lowerState, upperState, AscendC::CMPMODE::LE, singleSize);
428+ AscendC::PipeBarrier<PIPE_V>();
429+ 
430+ AscendC::Select<float, uint8_t>(
431+ pivot.template ReinterpretCast<float>(), higgerMask, lowerPivot.template ReinterpretCast<float>(),
432+ upperPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
433+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
434+ AscendC::PipeBarrier<PIPE_V>();
435+ AscendC::Cast<int16_t, int32_t>(stateTmp.template ReinterpretCast<int16_t>(), pivot,
436+ AscendC::RoundMode::CAST_NONE, singleSize);
437+ AscendC::PipeBarrier<PIPE_V>();
438+ AscendC::Cast<half, int16_t>(pivot.template ReinterpretCast<half>(),
439+ stateTmp.template ReinterpretCast<int16_t>(), AscendC::RoundMode::CAST_NONE,
440+ singleSize);
441+ AscendC::PipeBarrier<PIPE_V>();
442+ AscendC::Cast<uint8_t, half>(qState, pivot.template ReinterpretCast<half>(), AscendC::RoundMode::CAST_NONE,
443+ singleSize);
444+ AscendC::PipeBarrier<PIPE_V>();
445+ }
446+ 
447+ __aicore__ inline void CheckSign(const AscendC::LocalTensor<T>& normState, const AscendC::LocalTensor<T>& qMap,
448+ AscendC::LocalTensor<uint8_t>& qState, AscendC::TBuf<>& calcBuf,
449+ int32_t singleBlockNum)
450+ {
451+ // Match golden exactly: if normalized state and selected qmap have opposite
452+ // signs, move the code by one step toward the normalized state's sign.
453+ // Use int16 arithmetic (A2-style) to avoid float rounding in index conversion.
454+ uint32_t offset = 0;
455+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
456+ offset += singleSize * sizeof(int32_t);
457+ AscendC::LocalTensor<T> dq = calcBuf.GetWithOffset<T>(singleSize, offset);
458+ offset += singleSize * sizeof(T);
459+ AscendC::LocalTensor<int16_t> idxI = calcBuf.GetWithOffset<int16_t>(singleSize, offset);
460+ offset += singleSize * sizeof(int16_t);
461+ AscendC::LocalTensor<T> zero = calcBuf.GetWithOffset<T>(singleSize, offset);
462+ offset += singleSize * sizeof(T);
463+ AscendC::LocalTensor<int16_t> adjustedI = calcBuf.GetWithOffset<int16_t>(singleSize, offset);
464+ offset += singleSize * sizeof(int16_t);
465+ const int32_t maskBytes = singleSize / PER_UINT8_8BITS;
466+ AscendC::LocalTensor<uint8_t> normPosMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
467+ offset += maskBytes;
468+ AscendC::LocalTensor<uint8_t> normNegMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
469+ offset += maskBytes;
470+ AscendC::LocalTensor<uint8_t> qPosMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
471+ offset += maskBytes;
472+ AscendC::LocalTensor<uint8_t> qNegMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
473+ AscendC::LocalTensor<half> halfBuf = gatherOffset.template ReinterpretCast<half>();
474+ 
475+ const int32_t selRepeat = singleSize / REPEAT_NUM_128;
476+ 
477+ // idxI = int16(qState) — exact since values are 0-255
478+ AscendC::Cast<half, uint8_t>(halfBuf, qState, AscendC::RoundMode::CAST_NONE, singleSize);
479+ AscendC::PipeBarrier<PIPE_V>();
480+ AscendC::Cast<int16_t, half>(idxI, halfBuf, AscendC::RoundMode::CAST_RINT, singleSize);
481+ AscendC::PipeBarrier<PIPE_V>();
482+ 
483+ // dq = qMap[idxI]
484+ AscendC::Cast<int32_t, int16_t>(gatherOffset, idxI, AscendC::RoundMode::CAST_NONE, singleSize);
485+ AscendC::PipeBarrier<PIPE_V>();
486+ AscendC::Muls<int32_t>(gatherOffset, gatherOffset, sizeof(T), singleSize);
487+ AscendC::PipeBarrier<PIPE_V>();
488+ AscendC::Gather<T>(dq, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
489+ AscendC::PipeBarrier<PIPE_V>();
490+ 
491+ AscendC::Duplicate<T>(zero, static_cast<T>(0.0f), singleSize);
492+ AscendC::PipeBarrier<PIPE_V>();
493+ AscendC::Compare<T, uint8_t>(normPosMask, normState, zero, AscendC::CMPMODE::GT, singleSize);
494+ AscendC::Compare<T, uint8_t>(normNegMask, normState, zero, AscendC::CMPMODE::LT, singleSize);
495+ AscendC::PipeBarrier<PIPE_V>();
496+ AscendC::Compare<T, uint8_t>(qPosMask, dq, zero, AscendC::CMPMODE::GT, singleSize);
497+ AscendC::Compare<T, uint8_t>(qNegMask, dq, zero, AscendC::CMPMODE::LT, singleSize);
498+ AscendC::PipeBarrier<PIPE_V>();
499+ AscendC::And(normPosMask, normPosMask, qNegMask, maskBytes);
500+ AscendC::And(normNegMask, normNegMask, qPosMask, maskBytes);
501+ AscendC::PipeBarrier<PIPE_V>();
502+ 
503+ // Adjust idxI by ±1 in int16 space (avoids float rounding)
504+ AscendC::Adds<int16_t>(adjustedI, idxI, static_cast<int16_t>(1), singleSize);
505+ AscendC::PipeBarrier<PIPE_V>();
506+ AscendC::Select<int16_t, uint8_t>(idxI, normPosMask, adjustedI, idxI, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE,
507+ REPEAT_NUM_128, selRepeat, {1, 1, 1, 8, 8, 8});
508+ AscendC::PipeBarrier<PIPE_V>();
509+ AscendC::Adds<int16_t>(adjustedI, idxI, static_cast<int16_t>(-1), singleSize);
510+ AscendC::PipeBarrier<PIPE_V>();
511+ AscendC::Select<int16_t, uint8_t>(idxI, normNegMask, adjustedI, idxI, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE,
512+ REPEAT_NUM_128, selRepeat, {1, 1, 1, 8, 8, 8});
513+ AscendC::PipeBarrier<PIPE_V>();
514+ 
515+ // Clamp to [0, 255]
516+ AscendC::Maxs<int16_t>(idxI, idxI, static_cast<int16_t>(0), singleSize);
517+ AscendC::PipeBarrier<PIPE_V>();
518+ AscendC::Mins<int16_t>(idxI, idxI, static_cast<int16_t>(Q_MAP_SIZE - 1), singleSize);
519+ AscendC::PipeBarrier<PIPE_V>();
520+ 
521+ // int16 → half → uint8 (exact, no rounding)
522+ AscendC::Cast<half, int16_t>(halfBuf, idxI, AscendC::RoundMode::CAST_NONE, singleSize);
523+ AscendC::PipeBarrier<PIPE_V>();
524+ AscendC::Cast<uint8_t, half>(qState, halfBuf, AscendC::RoundMode::CAST_NONE, singleSize);
525+ AscendC::PipeBarrier<PIPE_V>();
526+ }
527+ 
528+ __aicore__ inline void Normlize(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& absMax,
529+ AscendC::TBuf<>& calcBuf, int32_t singleBlockNum)
530+ {
531+ AscendC::LocalTensor<T> absState = calcBuf.GetWithOffset<T>(singleSize, 0);
532+ AscendC::LocalTensor<T> absMaxTmp = calcBuf.GetWithOffset<T>(singleSize, singleSize * sizeof(T));
533+ AscendC::LocalTensor<T> absMaxBrcb = calcBuf.GetWithOffset<T>(singleSize, 2 * singleSize * sizeof(T));
534+ 
535+ AscendC::Abs<T>(absState, state, singleSize);
536+ AscendC::PipeBarrier<PIPE_V>();
537+ 
538+ uint32_t mask = AscendC::ONE_REPEAT_BYTE_SIZE / sizeof(T);
539+ uint32_t repeat = singleSize / mask;
540+ AscendC::BlockReduceMax<T>(absMaxTmp, absState, repeat, mask, 1, 1, STRIDE_8);
541+ AscendC::PipeBarrier<PIPE_V>();
542+ AscendC::WholeReduceMax<T>(absMax, absMaxTmp, blockSize / (AscendC::ONE_BLK_SIZE / sizeof(T)), singleBlockNum,
543+ 1, 1, PER_4NUM_ONEMAX, AscendC::ReduceOrder::ORDER_ONLY_VALUE);
544+ AscendC::PipeBarrier<PIPE_V>();
545+ 
546+ uint32_t srcShape[2] = {static_cast<uint32_t>(singleBlockNum), 1};
547+ uint32_t dstShape[2] = {static_cast<uint32_t>(singleBlockNum), blockSize};
548+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
549+ AscendC::PipeBarrier<PIPE_V>();
550+ AscendC::Div<T>(state, state, absMaxBrcb, singleSize);
551+ AscendC::PipeBarrier<PIPE_V>();
552+ }
553+ 
554+ __aicore__ inline void CalcAbsMax(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& absMax,
555+ AscendC::TBuf<>& calcBuf, int32_t singleBlockNum)
556+ {
557+ AscendC::LocalTensor<T> absState = calcBuf.GetWithOffset<T>(singleSize, 0);
558+ AscendC::LocalTensor<T> absMaxTmp = calcBuf.GetWithOffset<T>(singleSize, singleSize * sizeof(T));
559+ 
560+ AscendC::Abs<T>(absState, state, singleSize);
561+ AscendC::PipeBarrier<PIPE_V>();
562+ 
563+ uint32_t mask = AscendC::ONE_REPEAT_BYTE_SIZE / sizeof(T);
564+ uint32_t repeat = singleSize / mask;
565+ AscendC::BlockReduceMax<T>(absMaxTmp, absState, repeat, mask, 1, 1, STRIDE_8);
566+ AscendC::PipeBarrier<PIPE_V>();
567+ AscendC::WholeReduceMax<T>(absMax, absMaxTmp, blockSize / (AscendC::ONE_BLK_SIZE / sizeof(T)), singleBlockNum,
568+ 1, 1, PER_4NUM_ONEMAX, AscendC::ReduceOrder::ORDER_ONLY_VALUE);
569+ AscendC::PipeBarrier<PIPE_V>();
570+ }
571+ 
572+ __aicore__ inline void QuantM(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& qMap,
573+ AscendC::LocalTensor<uint8_t>& qState, const AscendC::LocalTensor<T>& absMax,
574+ AscendC::LocalTensor<T>& normStateTmp, AscendC::TBuf<>& calcBuf,
575+ int32_t singleBlockNum)
576+ {
577+ Normlize(state, absMax, calcBuf, singleBlockNum);
578+ 
579+ AscendC::Adds<T>(normStateTmp, state, static_cast<T>(0.0f), singleSize);
580+ AscendC::PipeBarrier<PIPE_V>();
581+ 
582+ BinarySearch(state, qMap, qState, calcBuf, singleBlockNum);
583+ 
584+ CheckSign(normStateTmp, qMap, qState, calcBuf, singleBlockNum);
585+ }
586+ 
587+ __aicore__ inline void QuantV(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& qMap,
588+ AscendC::LocalTensor<uint8_t>& qState, const AscendC::LocalTensor<T>& absMax,
589+ AscendC::TBuf<>& calcBuf, int32_t singleBlockNum)
590+ {
591+ Normlize(state, absMax, calcBuf, singleBlockNum);
592+ BinarySearch(state, qMap, qState, calcBuf, singleBlockNum);
593+ }
594+ 
595+ __aicore__ inline void DeQuant(const AscendC::LocalTensor<uint8_t>& state, const AscendC::LocalTensor<T>& qMap,
596+ AscendC::LocalTensor<T>& absMax, AscendC::TBuf<>& calcBuf, int32_t singleBlockNum,
597+ uint32_t validSize)
598+ {
599+ uint32_t offset = 0;
600+ AscendC::LocalTensor<T> dqState = calcBuf.GetWithOffset<T>(singleSize, offset);
601+ offset += singleSize * sizeof(T);
602+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
603+ offset += singleSize * sizeof(int32_t);
604+ AscendC::LocalTensor<T> absMaxBrcb = calcBuf.GetWithOffset<T>(singleSize, offset);
605+ 
606+ AscendC::Cast<half, uint8_t>(gatherOffset[singleSize >> 1].template ReinterpretCast<half>(), state,
607+ AscendC::RoundMode::CAST_NONE, singleSize);
608+ AscendC::PipeBarrier<PIPE_V>();
609+ AscendC::Cast<int32_t, half>(gatherOffset, gatherOffset[singleSize >> 1].template ReinterpretCast<half>(),
610+ AscendC::RoundMode::CAST_RINT, singleSize);
611+ AscendC::PipeBarrier<PIPE_V>();
612+ AscendC::Muls<int32_t>(gatherOffset, gatherOffset, sizeof(T), singleSize);
613+ AscendC::PipeBarrier<PIPE_V>();
614+#ifdef ASCENDC_CPU_DEBUG
615+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
616+ AscendC::Gather<T>(dqState[i * Q_MAP_SIZE], qMap,
617+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
618+ }
619+#else
620+ AscendC::Gather<T>(dqState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
621+#endif
622+ uint32_t srcShape[2] = {static_cast<uint32_t>(singleBlockNum), 1};
623+ uint32_t dstShape[2] = {static_cast<uint32_t>(singleBlockNum), blockSize};
624+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
625+ AscendC::PipeBarrier<PIPE_V>();
626+ if (validSize < static_cast<uint32_t>(singleSize)) {
627+ AscendC::Duplicate<T>(state.template ReinterpretCast<T>(), static_cast<T>(0), singleSize);
628+ AscendC::PipeBarrier<PIPE_V>();
629+ }
630+ AscendC::Mul<T>(state.template ReinterpretCast<T>(), dqState, absMaxBrcb, validSize);
631+ AscendC::PipeBarrier<PIPE_V>();
632+ }
633+ 
634+ __aicore__ inline void UpdateStateAndParam(const AscendC::LocalTensor<T>& dqStateM,
635+ const AscendC::LocalTensor<T>& dqStateV,
636+ const AscendC::LocalTensor<T>& grad, const AscendC::LocalTensor<T>& var,
637+ uint32_t validSize)
638+ {
639+ constexpr uint32_t oneRepeatSize = AscendC::GetVecLen() / sizeof(T);
640+ uint16_t repeatTimes = AscendC::CeilDivision(validSize, oneRepeatSize);
641+ float weightDecayFactor = weightDecay > 0.0f ? 1.0f - lr * weightDecay : 1.0f;
642+ asc_vf_call<AdamWQuantUpdateVF<T>>((__ubuf__ T*)dqStateM.GetPhyAddr(), (__ubuf__ T*)dqStateV.GetPhyAddr(),
643+ (__ubuf__ T*)var.GetPhyAddr(), (__ubuf__ T*)grad.GetPhyAddr(), gnormScale,
644+ beta1, 1.0f - beta1, beta2, 1.0f - beta2, eps * correction2, stepSize,
645+ weightDecayFactor, validSize, oneRepeatSize, repeatTimes);
646+ }
647+ 
648+private:
649+ TPipe pipe;
650+ 
651+ AscendC::TBuf<AscendC::QuePosition::VECCALC> varBuf;
652+ AscendC::TBuf<AscendC::QuePosition::VECCALC> gradBuf;
653+ AscendC::TBuf<AscendC::QuePosition::VECCALC> varFp32Buf;
654+ AscendC::TBuf<AscendC::QuePosition::VECCALC> gradFp32Buf;
655+ AscendC::TBuf<AscendC::QuePosition::VECCALC> stateMBuf;
656+ AscendC::TBuf<AscendC::QuePosition::VECCALC> stateVBuf;
657+ AscendC::TBuf<AscendC::QuePosition::VECCALC> qMapMBuf;
658+ AscendC::TBuf<AscendC::QuePosition::VECCALC> qMapVBuf;
659+ AscendC::TBuf<AscendC::QuePosition::VECCALC> absMaxMBuf;
660+ AscendC::TBuf<AscendC::QuePosition::VECCALC> absMaxVBuf;
661+ AscendC::TBuf<AscendC::QuePosition::VECCALC> calcBuf;
662+ 
663+ AscendC::GlobalTensor<T_VAR_GRAD> varGm;
664+ AscendC::GlobalTensor<T_VAR_GRAD> varRefGm;
665+ AscendC::GlobalTensor<T_VAR_GRAD> gradGm;
666+ AscendC::GlobalTensor<uint8_t> stateMGm;
667+ AscendC::GlobalTensor<uint8_t> stateMRefGm;
668+ AscendC::GlobalTensor<uint8_t> stateVGm;
669+ AscendC::GlobalTensor<uint8_t> stateVRefGm;
670+ AscendC::GlobalTensor<T> qMapMGm;
671+ 
672+ AscendC::GlobalTensor<T> qMapVGm;
673+ AscendC::GlobalTensor<T> absMaxMGm;
674+ AscendC::GlobalTensor<T> absMaxMRefGm;
675+ AscendC::GlobalTensor<T> absMaxVGm;
676+ AscendC::GlobalTensor<T> absMaxVRefGm;
677+ AscendC::GlobalTensor<U> stepGm;
678+ 
679+ AscendC::LocalTensor<T> qMapM;
680+ AscendC::LocalTensor<T> qMapV;
681+ AscendC::LocalTensor<uint8_t> stateM;
682+ AscendC::LocalTensor<uint8_t> stateV;
683+ AscendC::LocalTensor<T> absMaxM;
684+ AscendC::LocalTensor<T> absMaxV;
685+ AscendC::LocalTensor<T_VAR_GRAD> var;
686+ AscendC::LocalTensor<T_VAR_GRAD> grad;
687+ AscendC::LocalTensor<T> varFp32;
688+ AscendC::LocalTensor<T> gradFp32;
689+ 
690+ int32_t blockIdx;
691+ float step_ = 0;
692+ 
693+ int32_t singleBlockNum;
694+ int32_t singleSize;
695+ 
696+ uint64_t useNumCore = 0;
697+ uint64_t lastPreCoreRowWork = 0;
698+ uint64_t notLastCoreNum = 0;
699+ uint64_t notLastPreCoreRowWork = 0;
700+ uint64_t lastCoreLastBlock = 0;
701+ uint64_t lastBlockSize = 0;
702+ int64_t oneCoreDoBlockNumPerRow = 0;
703+ 
704+ float stepSize;
705+ float correction2;
706+ float lr;
707+ float beta1;
708+ float beta2;
709+ float weightDecay;
710+ float eps;
711+ float gnormScale;
712+ uint32_t blockSize;
713+ 
714+ uint64_t gmOffset = 0;
715+ uint64_t absmaxOffset = 0;
716+};
717+} // namespace ApplyAdamWQuantNS
718+#endif // APPLY_ADAM_W_QUANT_FP16_H_
Aoptim/apply_adam_w_quant/op_kernel/arch35/apply_adam_w_quant_fp32.h+711-0
@@ -0,0 +1,711 @@
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 apply_adam_w_quant_fp32.h
13+ * \brief
14+ */
15+#ifndef APPLY_ADAM_W_QUANT_FP32_H_
16+#define APPLY_ADAM_W_QUANT_FP32_H_
17+ 
18+#include "apply_adam_w_quant_base.h"
19+ 
20+namespace ApplyAdamWQuantNS {
21+using namespace AscendC;
22+ 
23+using namespace ApplyAdamWQuantNS;
24+ 
25+template <typename T, typename U>
26+class ApplyAdamWQuant {
27+public:
28+ __aicore__ inline ApplyAdamWQuant(){};
29+ __aicore__ inline void Init(GM_ADDR var, GM_ADDR grad, GM_ADDR m, GM_ADDR v, GM_ADDR qmap_m, GM_ADDR qmap_v,
30+ GM_ADDR absmax_m, GM_ADDR absmax_v, GM_ADDR step, GM_ADDR var_ref, GM_ADDR m_ref,
31+ GM_ADDR v_ref, GM_ADDR absmax_m_ref, GM_ADDR absmax_v_ref,
32+ const ApplyAdamWQuantRegbaseTilingData* tilingData)
33+ {
34+ this->blockIdx = AscendC::GetBlockIdx();
35+ ParseTilingDataFp32(tilingData);
36+ 
37+ if (this->blockIdx < notLastCoreNum) {
38+ gmOffset = this->blockIdx * notLastPreCoreRowWork * blockSize * oneCoreDoBlockNumPerRow;
39+ absmaxOffset = this->blockIdx * notLastPreCoreRowWork * oneCoreDoBlockNumPerRow;
40+ } else {
41+ gmOffset = (notLastCoreNum * notLastPreCoreRowWork +
42+ (this->blockIdx - notLastCoreNum) * lastPreCoreRowWork) *
43+ blockSize * oneCoreDoBlockNumPerRow;
44+ absmaxOffset = (notLastCoreNum * notLastPreCoreRowWork +
45+ (this->blockIdx - notLastCoreNum) * lastPreCoreRowWork) *
46+ oneCoreDoBlockNumPerRow;
47+ }
48+ this->singleBlockNum = oneCoreDoBlockNumPerRow;
49+ this->singleSize = blockSize * oneCoreDoBlockNumPerRow;
50+ varGm.SetGlobalBuffer((__gm__ T*)var + gmOffset);
51+ gradGm.SetGlobalBuffer((__gm__ T*)grad + gmOffset);
52+ stateMGm.SetGlobalBuffer((__gm__ uint8_t*)m + gmOffset);
53+ stateVGm.SetGlobalBuffer((__gm__ uint8_t*)v + gmOffset);
54+ qMapMGm.SetGlobalBuffer((__gm__ T*)qmap_m);
55+ qMapVGm.SetGlobalBuffer((__gm__ T*)qmap_v);
56+ absMaxMGm.SetGlobalBuffer((__gm__ T*)absmax_m + absmaxOffset);
57+ absMaxVGm.SetGlobalBuffer((__gm__ T*)absmax_v + absmaxOffset);
58+ stepGm.SetGlobalBuffer((__gm__ U*)step, 1);
59+ varRefGm.SetGlobalBuffer((__gm__ T*)var_ref + gmOffset);
60+ stateMRefGm.SetGlobalBuffer((__gm__ uint8_t*)m_ref + gmOffset);
61+ stateVRefGm.SetGlobalBuffer((__gm__ uint8_t*)v_ref + gmOffset);
62+ absMaxMRefGm.SetGlobalBuffer((__gm__ T*)absmax_m_ref + absmaxOffset);
63+ absMaxVRefGm.SetGlobalBuffer((__gm__ T*)absmax_v_ref + absmaxOffset);
64+ step_ = static_cast<float>(stepGm.GetValue(0));
65+ 
66+ pipe.InitBuffer(varBuf, singleSize * sizeof(T));
67+ pipe.InitBuffer(gradBuf, singleSize * sizeof(T));
68+ pipe.InitBuffer(stateMBuf, singleSize * sizeof(T));
69+ pipe.InitBuffer(stateVBuf, singleSize * sizeof(T));
70+ pipe.InitBuffer(qMapMBuf, Q_MAP_SIZE * sizeof(T));
71+ pipe.InitBuffer(qMapVBuf, Q_MAP_SIZE * sizeof(T));
72+ pipe.InitBuffer(absMaxMBuf, oneCoreDoBlockNumPerRow * sizeof(T));
73+ pipe.InitBuffer(absMaxVBuf, oneCoreDoBlockNumPerRow * sizeof(T));
74+ pipe.InitBuffer(calcBuf, CALC_BUF_NUM * singleSize * sizeof(float));
75+ 
76+ this->qMapM = qMapMBuf.Get<T>();
77+ this->qMapV = qMapVBuf.Get<T>();
78+ this->stateM = stateMBuf.Get<uint8_t>();
79+ this->stateV = stateVBuf.Get<uint8_t>();
80+ this->absMaxM = absMaxMBuf.Get<T>();
81+ this->absMaxV = absMaxVBuf.Get<T>();
82+ this->var = varBuf.Get<T>();
83+ this->grad = gradBuf.Get<T>();
84+ 
85+ step_ += 1;
86+ uint32_t count = AscendC::ONE_BLK_SIZE / sizeof(float);
87+ AscendC::LocalTensor<float> stepTensor = calcBuf.GetWithOffset<float>(count, 0);
88+ AscendC::LocalTensor<float> resTensor = calcBuf.GetWithOffset<float>(count, AscendC::ONE_BLK_SIZE);
89+ AscendC::Duplicate<float>(stepTensor, step_, AscendC::ONE_BLK_SIZE / sizeof(float));
90+ AscendC::PipeBarrier<PIPE_V>();
91+ 
92+ float correction1 = 1 - PowS(resTensor, beta1, stepTensor);
93+ this->correction2 = sqrt(1 - PowS(resTensor, beta2, stepTensor));
94+ this->stepSize = -lr * this->correction2 / correction1;
95+ }
96+ __aicore__ inline void Process()
97+ {
98+ DataCopyIn<T>(qMapM, qMapMGm, Q_MAP_SIZE);
99+ DataCopyIn<T>(qMapV, qMapVGm, Q_MAP_SIZE);
100+ 
101+ int64_t perCoreRowWork = lastPreCoreRowWork;
102+ if (this->blockIdx < notLastCoreNum) {
103+ perCoreRowWork += 1;
104+ }
105+ uint64_t blockOffset = 0;
106+ uint64_t absMaxBlockOffset = 0;
107+ for (int64_t n = 0; n < perCoreRowWork; n++) {
108+ bool isLastRow = static_cast<uint64_t>(this->blockIdx + 1) == useNumCore && n + 1 == perCoreRowWork;
109+ singleBlockNum = isLastRow ? lastCoreLastBlock : oneCoreDoBlockNumPerRow;
110+ singleSize = blockSize * singleBlockNum;
111+ uint32_t validSize = isLastRow ? singleSize - blockSize + lastBlockSize : singleSize;
112+ 
113+ DataCopyIn<uint8_t>(stateM, stateMGm[blockOffset], validSize);
114+ DataCopyIn<uint8_t>(stateV, stateVGm[blockOffset], validSize);
115+ 
116+ DataCopyIn<T>(absMaxM, absMaxMGm[absMaxBlockOffset], singleBlockNum);
117+ DataCopyIn<T>(absMaxV, absMaxVGm[absMaxBlockOffset], singleBlockNum);
118+ 
119+ PipeSync<AscendC::HardEvent::MTE2_V>();
120+ 
121+ DeQuantFp32(stateM, qMapM, absMaxM, calcBuf, singleBlockNum, validSize);
122+ DeQuantFp32(stateV, qMapV, absMaxV, calcBuf, singleBlockNum, validSize);
123+ 
124+ DataCopyIn<T>(var, varGm[blockOffset], validSize);
125+ DataCopyIn<T>(grad, gradGm[blockOffset], validSize);
126+ 
127+ PipeSync<AscendC::HardEvent::MTE2_V>();
128+ 
129+ UpdateStateAndParamFp32(stateM.template ReinterpretCast<T>(), stateV.template ReinterpretCast<T>(), grad,
130+ var, validSize);
131+ PipeSync<AscendC::HardEvent::V_MTE3>();
132+ 
133+ DataCopyOut<T>(varRefGm[blockOffset], var, validSize);
134+ 
135+ QuantMFp32(stateM.template ReinterpretCast<T>(), qMapM, stateM, absMaxM, grad, calcBuf, singleBlockNum);
136+ QuantVFp32(stateV.template ReinterpretCast<T>(), qMapV, stateV, absMaxV, calcBuf, singleBlockNum);
137+ PipeSync<AscendC::HardEvent::V_MTE3>();
138+ 
139+ DataCopyOut<T>(absMaxMRefGm[absMaxBlockOffset], absMaxM, singleBlockNum);
140+ DataCopyOut<T>(absMaxVRefGm[absMaxBlockOffset], absMaxV, singleBlockNum);
141+ DataCopyOut<uint8_t>(stateMRefGm[blockOffset], stateM, validSize);
142+ DataCopyOut<uint8_t>(stateVRefGm[blockOffset], stateV, validSize);
143+ PipeSync<AscendC::HardEvent::MTE3_MTE2>();
144+ 
145+ blockOffset += singleSize;
146+ absMaxBlockOffset += singleBlockNum;
147+ }
148+ }
149+ 
150+private:
151+ __aicore__ inline void ParseTilingDataFp32(const ApplyAdamWQuantRegbaseTilingData* tilingData)
152+ {
153+ useNumCore = tilingData->use_num_core;
154+ lastPreCoreRowWork = tilingData->last_pre_core_row_work;
155+ notLastCoreNum = tilingData->not_last_core_num;
156+ notLastPreCoreRowWork = tilingData->not_last_pre_core_row_work;
157+ lastCoreLastBlock = tilingData->last_core_last_block;
158+ lastBlockSize = tilingData->last_block_size;
159+ oneCoreDoBlockNumPerRow = tilingData->one_core_do_block_num_per_row;
160+ 
161+ lr = tilingData->lr;
162+ beta1 = tilingData->beta1;
163+ beta2 = tilingData->beta2;
164+ weightDecay = tilingData->weight_decay;
165+ eps = tilingData->eps;
166+ gnormScale = tilingData->gnorm_scale;
167+ blockSize = tilingData->block_size;
168+ }
169+ 
170+ __aicore__ inline void BinarySearchFp32(const AscendC::LocalTensor<T>& normState,
171+ const AscendC::LocalTensor<T>& qMap, AscendC::LocalTensor<uint8_t>& qState,
172+ AscendC::TBuf<>& calcBuf, int32_t singleBlockNum)
173+ {
174+ uint32_t offset = 0;
175+ AscendC::LocalTensor<int32_t> lowerPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
176+ offset += singleSize * sizeof(int32_t);
177+ AscendC::LocalTensor<int32_t> upperPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
178+ offset += singleSize * sizeof(int32_t);
179+ AscendC::LocalTensor<int32_t> pivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
180+ offset += singleSize * sizeof(int32_t);
181+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
182+ offset += singleSize * sizeof(int32_t);
183+ AscendC::LocalTensor<T> stateTmp = calcBuf.GetWithOffset<T>(singleSize, offset);
184+ offset += singleSize * sizeof(T);
185+ AscendC::LocalTensor<uint8_t> higgerMask = calcBuf.GetWithOffset<uint8_t>(singleSize / PER_UINT8_8BITS, offset);
186+ 
187+ AscendC::Duplicate<int32_t>(lowerPivot, 0, singleSize);
188+ AscendC::Duplicate<int32_t>(upperPivot, Q_MAP_SIZE - 1, singleSize);
189+ AscendC::Duplicate<int32_t>(pivot, (Q_MAP_SIZE - 1) >> 1, singleSize);
190+ 
191+ T midValue = qMap.GetValue((Q_MAP_SIZE - 1) >> 1);
192+ 
193+ PipeSync<AscendC::HardEvent::S_V>();
194+ 
195+ AscendC::Duplicate<T>(stateTmp, midValue, singleSize);
196+ AscendC::PipeBarrier<PIPE_V>();
197+ AscendC::Compare<T, uint8_t>(higgerMask, normState, stateTmp, AscendC::CMPMODE::GT, singleSize);
198+ AscendC::PipeBarrier<PIPE_V>();
199+ AscendC::Select<float, uint8_t>(
200+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
201+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
202+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
203+ AscendC::Select<float, uint8_t>(
204+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
205+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
206+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
207+ AscendC::PipeBarrier<PIPE_V>();
208+ 
209+ for (int32_t i = 0; i < REPEAT_7_TIMES; ++i) {
210+ AscendC::Add<int32_t>(pivot, lowerPivot, upperPivot, singleSize);
211+ AscendC::PipeBarrier<PIPE_V>();
212+ AscendC::ShiftRight<int32_t>(pivot, pivot, 1, singleSize);
213+ AscendC::PipeBarrier<PIPE_V>();
214+ 
215+ AscendC::Muls<int32_t>(gatherOffset, pivot, sizeof(T), singleSize);
216+ AscendC::PipeBarrier<PIPE_V>();
217+ 
218+#ifdef ASCENDC_CPU_DEBUG
219+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
220+ AscendC::Gather<T>(stateTmp[i * Q_MAP_SIZE], qMap,
221+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
222+ }
223+#else
224+ AscendC::Gather<T>(stateTmp, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
225+#endif
226+ AscendC::PipeBarrier<PIPE_V>();
227+ 
228+ AscendC::Compare<T, uint8_t>(higgerMask, normState, stateTmp, AscendC::CMPMODE::GT, singleSize);
229+ AscendC::PipeBarrier<PIPE_V>();
230+ 
231+ AscendC::Select<float, uint8_t>(
232+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
233+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
234+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
235+ 
236+ AscendC::Select<float, uint8_t>(
237+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
238+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
239+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
240+ AscendC::PipeBarrier<PIPE_V>();
241+ }
242+ 
243+ AscendC::LocalTensor<T> lowerState = stateTmp;
244+ AscendC::LocalTensor<T> upperState = pivot.template ReinterpretCast<T>();
245+ 
246+ AscendC::Muls<int32_t>(gatherOffset, lowerPivot, sizeof(T), singleSize);
247+ AscendC::PipeBarrier<PIPE_V>();
248+#ifdef ASCENDC_CPU_DEBUG
249+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
250+ AscendC::Gather<T>(lowerState[i * Q_MAP_SIZE], qMap,
251+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
252+ }
253+#else
254+ AscendC::Gather<T>(lowerState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
255+#endif
256+ AscendC::PipeBarrier<PIPE_V>();
257+ AscendC::Sub<T>(lowerState, normState, lowerState, singleSize);
258+ AscendC::PipeBarrier<PIPE_V>();
259+ AscendC::Abs<T>(lowerState, lowerState, singleSize);
260+ 
261+ AscendC::Muls<int32_t>(gatherOffset, upperPivot, sizeof(T), singleSize);
262+ AscendC::PipeBarrier<PIPE_V>();
263+ 
264+#ifdef ASCENDC_CPU_DEBUG
265+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
266+ AscendC::Gather<T>(upperState[i * Q_MAP_SIZE], qMap,
267+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
268+ }
269+#else
270+ AscendC::Gather<T>(upperState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
271+#endif
272+ AscendC::PipeBarrier<PIPE_V>();
273+ AscendC::Sub<T>(upperState, normState, upperState, singleSize);
274+ AscendC::PipeBarrier<PIPE_V>();
275+ AscendC::Abs<T>(upperState, upperState, singleSize);
276+ AscendC::PipeBarrier<PIPE_V>();
277+ 
278+ AscendC::Compare<T, uint8_t>(higgerMask, lowerState, upperState, AscendC::CMPMODE::LE, singleSize);
279+ AscendC::PipeBarrier<PIPE_V>();
280+ 
281+ AscendC::Select<float, uint8_t>(
282+ pivot.template ReinterpretCast<float>(), higgerMask, lowerPivot.template ReinterpretCast<float>(),
283+ upperPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
284+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
285+ AscendC::PipeBarrier<PIPE_V>();
286+ AscendC::Cast<int16_t, int32_t>(stateTmp.template ReinterpretCast<int16_t>(), pivot,
287+ AscendC::RoundMode::CAST_NONE, singleSize);
288+ AscendC::PipeBarrier<PIPE_V>();
289+ AscendC::Cast<half, int16_t>(pivot.template ReinterpretCast<half>(),
290+ stateTmp.template ReinterpretCast<int16_t>(), AscendC::RoundMode::CAST_NONE,
291+ singleSize);
292+ AscendC::PipeBarrier<PIPE_V>();
293+ AscendC::Cast<uint8_t, half>(qState, pivot.template ReinterpretCast<half>(), AscendC::RoundMode::CAST_NONE,
294+ singleSize);
295+ AscendC::PipeBarrier<PIPE_V>();
296+ }
297+ 
298+ __aicore__ inline void BinarySearchScaledFp32(const AscendC::LocalTensor<T>& state,
299+ const AscendC::LocalTensor<T>& qMap,
300+ const AscendC::LocalTensor<T>& absMax,
301+ AscendC::LocalTensor<uint8_t>& qState, AscendC::TBuf<>& calcBuf,
302+ int32_t singleBlockNum)
303+ {
304+ uint32_t offset = 0;
305+ AscendC::LocalTensor<int32_t> lowerPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
306+ offset += singleSize * sizeof(int32_t);
307+ AscendC::LocalTensor<int32_t> upperPivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
308+ offset += singleSize * sizeof(int32_t);
309+ AscendC::LocalTensor<int32_t> pivot = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
310+ offset += singleSize * sizeof(int32_t);
311+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
312+ offset += singleSize * sizeof(int32_t);
313+ AscendC::LocalTensor<T> stateTmp = calcBuf.GetWithOffset<T>(singleSize, offset);
314+ offset += singleSize * sizeof(T);
315+ AscendC::LocalTensor<uint8_t> higgerMask = calcBuf.GetWithOffset<uint8_t>(singleSize / PER_UINT8_8BITS, offset);
316+ AscendC::LocalTensor<T> absMaxBrcb = gatherOffset.template ReinterpretCast<T>();
317+ 
318+ uint32_t srcShape[2] = {static_cast<uint32_t>(singleBlockNum), 1};
319+ uint32_t dstShape[2] = {static_cast<uint32_t>(singleBlockNum), blockSize};
320+ 
321+ AscendC::Duplicate<int32_t>(lowerPivot, 0, singleSize);
322+ AscendC::Duplicate<int32_t>(upperPivot, Q_MAP_SIZE - 1, singleSize);
323+ AscendC::Duplicate<int32_t>(pivot, (Q_MAP_SIZE - 1) >> 1, singleSize);
324+ 
325+ T midValue = qMap.GetValue((Q_MAP_SIZE - 1) >> 1);
326+ 
327+ PipeSync<AscendC::HardEvent::S_V>();
328+ 
329+ AscendC::Duplicate<T>(stateTmp, midValue, singleSize);
330+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
331+ AscendC::PipeBarrier<PIPE_V>();
332+ AscendC::Mul<T>(stateTmp, stateTmp, absMaxBrcb, singleSize);
333+ AscendC::PipeBarrier<PIPE_V>();
334+ AscendC::Compare<T, uint8_t>(higgerMask, state, stateTmp, AscendC::CMPMODE::GT, singleSize);
335+ AscendC::PipeBarrier<PIPE_V>();
336+ 
337+ AscendC::Select<float, uint8_t>(
338+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
339+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
340+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
341+ AscendC::Select<float, uint8_t>(
342+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
343+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
344+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
345+ AscendC::PipeBarrier<PIPE_V>();
346+ 
347+ for (int32_t i = 0; i < REPEAT_7_TIMES; ++i) {
348+ AscendC::Add<int32_t>(pivot, lowerPivot, upperPivot, singleSize);
349+ AscendC::PipeBarrier<PIPE_V>();
350+ AscendC::ShiftRight<int32_t>(pivot, pivot, 1, singleSize);
351+ AscendC::PipeBarrier<PIPE_V>();
352+ 
353+ AscendC::Muls<int32_t>(gatherOffset, pivot, sizeof(T), singleSize);
354+ AscendC::PipeBarrier<PIPE_V>();
355+ 
356+#ifdef ASCENDC_CPU_DEBUG
357+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
358+ AscendC::Gather<T>(stateTmp[i * Q_MAP_SIZE], qMap,
359+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
360+ }
361+#else
362+ AscendC::Gather<T>(stateTmp, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
363+#endif
364+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
365+ AscendC::PipeBarrier<PIPE_V>();
366+ AscendC::Mul<T>(stateTmp, stateTmp, absMaxBrcb, singleSize);
367+ AscendC::PipeBarrier<PIPE_V>();
368+ 
369+ AscendC::Compare<T, uint8_t>(higgerMask, state, stateTmp, AscendC::CMPMODE::GT, singleSize);
370+ AscendC::PipeBarrier<PIPE_V>();
371+ 
372+ AscendC::Select<float, uint8_t>(
373+ lowerPivot.template ReinterpretCast<float>(), higgerMask, pivot.template ReinterpretCast<float>(),
374+ lowerPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
375+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
376+ 
377+ AscendC::Select<float, uint8_t>(
378+ upperPivot.template ReinterpretCast<float>(), higgerMask, upperPivot.template ReinterpretCast<float>(),
379+ pivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
380+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
381+ AscendC::PipeBarrier<PIPE_V>();
382+ }
383+ 
384+ AscendC::LocalTensor<T> lowerState = stateTmp;
385+ AscendC::LocalTensor<T> upperState = pivot.template ReinterpretCast<T>();
386+ 
387+ AscendC::Muls<int32_t>(gatherOffset, lowerPivot, sizeof(T), singleSize);
388+ AscendC::PipeBarrier<PIPE_V>();
389+#ifdef ASCENDC_CPU_DEBUG
390+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
391+ AscendC::Gather<T>(lowerState[i * Q_MAP_SIZE], qMap,
392+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
393+ }
394+#else
395+ AscendC::Gather<T>(lowerState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
396+#endif
397+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
398+ AscendC::PipeBarrier<PIPE_V>();
399+ AscendC::Mul<T>(lowerState, lowerState, absMaxBrcb, singleSize);
400+ AscendC::PipeBarrier<PIPE_V>();
401+ AscendC::Sub<T>(lowerState, state, lowerState, singleSize);
402+ AscendC::PipeBarrier<PIPE_V>();
403+ AscendC::Abs<T>(lowerState, lowerState, singleSize);
404+ 
405+ AscendC::Muls<int32_t>(gatherOffset, upperPivot, sizeof(T), singleSize);
406+ AscendC::PipeBarrier<PIPE_V>();
407+ 
408+#ifdef ASCENDC_CPU_DEBUG
409+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
410+ AscendC::Gather<T>(upperState[i * Q_MAP_SIZE], qMap,
411+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
412+ }
413+#else
414+ AscendC::Gather<T>(upperState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
415+#endif
416+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
417+ AscendC::PipeBarrier<PIPE_V>();
418+ AscendC::Mul<T>(upperState, upperState, absMaxBrcb, singleSize);
419+ AscendC::PipeBarrier<PIPE_V>();
420+ AscendC::Sub<T>(upperState, state, upperState, singleSize);
421+ AscendC::PipeBarrier<PIPE_V>();
422+ AscendC::Abs<T>(upperState, upperState, singleSize);
423+ AscendC::PipeBarrier<PIPE_V>();
424+ 
425+ AscendC::Compare<T, uint8_t>(higgerMask, lowerState, upperState, AscendC::CMPMODE::LE, singleSize);
426+ AscendC::PipeBarrier<PIPE_V>();
427+ 
428+ AscendC::Select<float, uint8_t>(
429+ pivot.template ReinterpretCast<float>(), higgerMask, lowerPivot.template ReinterpretCast<float>(),
430+ upperPivot.template ReinterpretCast<float>(), AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE, REPEAT_NUM,
431+ singleSize / REPEAT_NUM, {1, 1, 1, 8, 8, 8});
432+ AscendC::PipeBarrier<PIPE_V>();
433+ AscendC::Cast<int16_t, int32_t>(stateTmp.template ReinterpretCast<int16_t>(), pivot,
434+ AscendC::RoundMode::CAST_NONE, singleSize);
435+ AscendC::PipeBarrier<PIPE_V>();
436+ AscendC::Cast<half, int16_t>(pivot.template ReinterpretCast<half>(),
437+ stateTmp.template ReinterpretCast<int16_t>(), AscendC::RoundMode::CAST_NONE,
438+ singleSize);
439+ AscendC::PipeBarrier<PIPE_V>();
440+ AscendC::Cast<uint8_t, half>(qState, pivot.template ReinterpretCast<half>(), AscendC::RoundMode::CAST_NONE,
441+ singleSize);
442+ AscendC::PipeBarrier<PIPE_V>();
443+ }
444+ 
445+ __aicore__ inline void CheckSignFp32(const AscendC::LocalTensor<T>& normState, const AscendC::LocalTensor<T>& qMap,
446+ AscendC::LocalTensor<uint8_t>& qState, AscendC::TBuf<>& calcBuf,
447+ int32_t singleBlockNum)
448+ {
449+ // Match golden exactly: if normalized state and selected qmap have opposite
450+ // signs, move the code by one step toward the normalized state's sign.
451+ // Use int16 arithmetic (A2-style) to avoid float rounding in index conversion.
452+ uint32_t offset = 0;
453+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
454+ offset += singleSize * sizeof(int32_t);
455+ AscendC::LocalTensor<T> dq = calcBuf.GetWithOffset<T>(singleSize, offset);
456+ offset += singleSize * sizeof(T);
457+ AscendC::LocalTensor<int16_t> idxI = calcBuf.GetWithOffset<int16_t>(singleSize, offset);
458+ offset += singleSize * sizeof(int16_t);
459+ AscendC::LocalTensor<T> zero = calcBuf.GetWithOffset<T>(singleSize, offset);
460+ offset += singleSize * sizeof(T);
461+ AscendC::LocalTensor<int16_t> adjustedI = calcBuf.GetWithOffset<int16_t>(singleSize, offset);
462+ offset += singleSize * sizeof(int16_t);
463+ const int32_t maskBytes = singleSize / PER_UINT8_8BITS;
464+ AscendC::LocalTensor<uint8_t> normPosMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
465+ offset += maskBytes;
466+ AscendC::LocalTensor<uint8_t> normNegMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
467+ offset += maskBytes;
468+ AscendC::LocalTensor<uint8_t> qPosMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
469+ offset += maskBytes;
470+ AscendC::LocalTensor<uint8_t> qNegMask = calcBuf.GetWithOffset<uint8_t>(maskBytes, offset);
471+ AscendC::LocalTensor<half> halfBuf = gatherOffset.template ReinterpretCast<half>();
472+ 
473+ const int32_t selRepeat = singleSize / REPEAT_NUM_128;
474+ 
475+ // idxI = int16(qState) — exact since values are 0-255
476+ AscendC::Cast<half, uint8_t>(halfBuf, qState, AscendC::RoundMode::CAST_NONE, singleSize);
477+ AscendC::PipeBarrier<PIPE_V>();
478+ AscendC::Cast<int16_t, half>(idxI, halfBuf, AscendC::RoundMode::CAST_RINT, singleSize);
479+ AscendC::PipeBarrier<PIPE_V>();
480+ 
481+ // dq = qMap[idxI]
482+ AscendC::Cast<int32_t, int16_t>(gatherOffset, idxI, AscendC::RoundMode::CAST_NONE, singleSize);
483+ AscendC::PipeBarrier<PIPE_V>();
484+ AscendC::Muls<int32_t>(gatherOffset, gatherOffset, sizeof(T), singleSize);
485+ AscendC::PipeBarrier<PIPE_V>();
486+ AscendC::Gather<T>(dq, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
487+ AscendC::PipeBarrier<PIPE_V>();
488+ 
489+ AscendC::Duplicate<T>(zero, static_cast<T>(0.0f), singleSize);
490+ AscendC::PipeBarrier<PIPE_V>();
491+ AscendC::Compare<T, uint8_t>(normPosMask, normState, zero, AscendC::CMPMODE::GT, singleSize);
492+ AscendC::Compare<T, uint8_t>(normNegMask, normState, zero, AscendC::CMPMODE::LT, singleSize);
493+ AscendC::PipeBarrier<PIPE_V>();
494+ AscendC::Compare<T, uint8_t>(qPosMask, dq, zero, AscendC::CMPMODE::GT, singleSize);
495+ AscendC::Compare<T, uint8_t>(qNegMask, dq, zero, AscendC::CMPMODE::LT, singleSize);
496+ AscendC::PipeBarrier<PIPE_V>();
497+ AscendC::And(normPosMask, normPosMask, qNegMask, maskBytes);
498+ AscendC::And(normNegMask, normNegMask, qPosMask, maskBytes);
499+ AscendC::PipeBarrier<PIPE_V>();
500+ 
501+ // Adjust idxI by ±1 in int16 space (avoids float rounding)
502+ AscendC::Adds<int16_t>(adjustedI, idxI, static_cast<int16_t>(1), singleSize);
503+ AscendC::PipeBarrier<PIPE_V>();
504+ AscendC::Select<int16_t, uint8_t>(idxI, normPosMask, adjustedI, idxI, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE,
505+ REPEAT_NUM_128, selRepeat, {1, 1, 1, 8, 8, 8});
506+ AscendC::PipeBarrier<PIPE_V>();
507+ AscendC::Adds<int16_t>(adjustedI, idxI, static_cast<int16_t>(-1), singleSize);
508+ AscendC::PipeBarrier<PIPE_V>();
509+ AscendC::Select<int16_t, uint8_t>(idxI, normNegMask, adjustedI, idxI, AscendC::SELMODE::VSEL_TENSOR_TENSOR_MODE,
510+ REPEAT_NUM_128, selRepeat, {1, 1, 1, 8, 8, 8});
511+ AscendC::PipeBarrier<PIPE_V>();
512+ 
513+ // Clamp to [0, 255]
514+ AscendC::Maxs<int16_t>(idxI, idxI, static_cast<int16_t>(0), singleSize);
515+ AscendC::PipeBarrier<PIPE_V>();
516+ AscendC::Mins<int16_t>(idxI, idxI, static_cast<int16_t>(Q_MAP_SIZE - 1), singleSize);
517+ AscendC::PipeBarrier<PIPE_V>();
518+ 
519+ // int16 → half → uint8 (exact, no rounding)
520+ AscendC::Cast<half, int16_t>(halfBuf, idxI, AscendC::RoundMode::CAST_NONE, singleSize);
521+ AscendC::PipeBarrier<PIPE_V>();
522+ AscendC::Cast<uint8_t, half>(qState, halfBuf, AscendC::RoundMode::CAST_NONE, singleSize);
523+ AscendC::PipeBarrier<PIPE_V>();
524+ }
525+ 
526+ __aicore__ inline void NormlizeFp32(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& absMax,
527+ AscendC::TBuf<>& calcBuf, int32_t singleBlockNum)
528+ {
529+ AscendC::LocalTensor<T> absState = calcBuf.GetWithOffset<T>(singleSize, 0);
530+ AscendC::LocalTensor<T> absMaxTmp = calcBuf.GetWithOffset<T>(singleSize, singleSize * sizeof(T));
531+ AscendC::LocalTensor<T> absMaxBrcb = calcBuf.GetWithOffset<T>(singleSize, 2 * singleSize * sizeof(T));
532+ 
533+ AscendC::Abs<T>(absState, state, singleSize);
534+ AscendC::PipeBarrier<PIPE_V>();
535+ 
536+ uint32_t mask = AscendC::ONE_REPEAT_BYTE_SIZE / sizeof(T);
537+ uint32_t repeat = singleSize / mask;
538+ AscendC::BlockReduceMax<T>(absMaxTmp, absState, repeat, mask, 1, 1, STRIDE_8);
539+ AscendC::PipeBarrier<PIPE_V>();
540+ AscendC::WholeReduceMax<T>(absMax, absMaxTmp, blockSize / (AscendC::ONE_BLK_SIZE / sizeof(T)), singleBlockNum,
541+ 1, 1, PER_4NUM_ONEMAX, AscendC::ReduceOrder::ORDER_ONLY_VALUE);
542+ AscendC::PipeBarrier<PIPE_V>();
543+ 
544+ uint32_t srcShape[2] = {static_cast<uint32_t>(singleBlockNum), 1};
545+ uint32_t dstShape[2] = {static_cast<uint32_t>(singleBlockNum), blockSize};
546+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
547+ AscendC::PipeBarrier<PIPE_V>();
548+ AscendC::Div<T>(state, state, absMaxBrcb, singleSize);
TangPC
TangPCTangPC8月1日

这里按 absmax 归一化没有做零保护，而本 PR 新增的 golden 里 _quantize_state 明确写了 denom = where(absmax == 0, 1, absmax)。整个 256 元素的 block 全为零时（首步 m 全零、稀疏梯度都会出现），kernel 这边会 inf 或 nan，随后进二分搜索得到的量化码完全是垃圾，golden 那边得到的是 0，两边直接对不上。fp16.h 的 Normlize 同一处也一样。建议按 golden 的写法先把 absmax 为 0 的位置换成 1 再除。

likedislike
549+ AscendC::PipeBarrier<PIPE_V>();
550+ }
551+ 
552+ __aicore__ inline void CalcAbsMaxFp32(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& absMax,
553+ AscendC::TBuf<>& calcBuf, int32_t singleBlockNum)
554+ {
555+ AscendC::LocalTensor<T> absState = calcBuf.GetWithOffset<T>(singleSize, 0);
556+ AscendC::LocalTensor<T> absMaxTmp = calcBuf.GetWithOffset<T>(singleSize, singleSize * sizeof(T));
557+ 
558+ AscendC::Abs<T>(absState, state, singleSize);
559+ AscendC::PipeBarrier<PIPE_V>();
560+ 
561+ uint32_t mask = AscendC::ONE_REPEAT_BYTE_SIZE / sizeof(T);
562+ uint32_t repeat = singleSize / mask;
563+ AscendC::BlockReduceMax<T>(absMaxTmp, absState, repeat, mask, 1, 1, STRIDE_8);
564+ AscendC::PipeBarrier<PIPE_V>();
565+ AscendC::WholeReduceMax<T>(absMax, absMaxTmp, blockSize / (AscendC::ONE_BLK_SIZE / sizeof(T)), singleBlockNum,
566+ 1, 1, PER_4NUM_ONEMAX, AscendC::ReduceOrder::ORDER_ONLY_VALUE);
567+ AscendC::PipeBarrier<PIPE_V>();
568+ }
569+ 
570+ __aicore__ inline void QuantMFp32(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& qMap,
571+ AscendC::LocalTensor<uint8_t>& qState, const AscendC::LocalTensor<T>& absMax,
572+ AscendC::LocalTensor<T>& normStateTmp, AscendC::TBuf<>& calcBuf,
573+ int32_t singleBlockNum)
574+ {
575+ NormlizeFp32(state, absMax, calcBuf, singleBlockNum);
576+ 
577+ AscendC::Adds<T>(normStateTmp, state, static_cast<T>(0.0f), singleSize);
578+ AscendC::PipeBarrier<PIPE_V>();
579+ 
580+ BinarySearchFp32(state, qMap, qState, calcBuf, singleBlockNum);
581+ 
582+ CheckSignFp32(normStateTmp, qMap, qState, calcBuf, singleBlockNum);
583+ }
584+ 
585+ __aicore__ inline void QuantVFp32(const AscendC::LocalTensor<T>& state, const AscendC::LocalTensor<T>& qMap,
586+ AscendC::LocalTensor<uint8_t>& qState, const AscendC::LocalTensor<T>& absMax,
587+ AscendC::TBuf<>& calcBuf, int32_t singleBlockNum)
588+ {
589+ NormlizeFp32(state, absMax, calcBuf, singleBlockNum);
590+ BinarySearchFp32(state, qMap, qState, calcBuf, singleBlockNum);
591+ }
592+ 
593+ __aicore__ inline void DeQuantFp32(const AscendC::LocalTensor<uint8_t>& state, const AscendC::LocalTensor<T>& qMap,
594+ AscendC::LocalTensor<T>& absMax, AscendC::TBuf<>& calcBuf,
595+ int32_t singleBlockNum, uint32_t validSize)
596+ {
597+ uint32_t offset = 0;
598+ AscendC::LocalTensor<T> dqState = calcBuf.GetWithOffset<T>(singleSize, offset);
599+ offset += singleSize * sizeof(T);
600+ AscendC::LocalTensor<int32_t> gatherOffset = calcBuf.GetWithOffset<int32_t>(singleSize, offset);
601+ offset += singleSize * sizeof(int32_t);
602+ AscendC::LocalTensor<T> absMaxBrcb = calcBuf.GetWithOffset<T>(singleSize, offset);
603+ 
604+ AscendC::Cast<half, uint8_t>(gatherOffset[singleSize >> 1].template ReinterpretCast<half>(), state,
605+ AscendC::RoundMode::CAST_NONE, singleSize);
606+ AscendC::PipeBarrier<PIPE_V>();
607+ AscendC::Cast<int32_t, half>(gatherOffset, gatherOffset[singleSize >> 1].template ReinterpretCast<half>(),
608+ AscendC::RoundMode::CAST_RINT, singleSize);
609+ AscendC::PipeBarrier<PIPE_V>();
610+ AscendC::Muls<int32_t>(gatherOffset, gatherOffset, sizeof(T), singleSize);
611+ AscendC::PipeBarrier<PIPE_V>();
612+#ifdef ASCENDC_CPU_DEBUG
613+ for (int i = 0; i < singleSize / Q_MAP_SIZE; i++) {
614+ AscendC::Gather<T>(dqState[i * Q_MAP_SIZE], qMap,
615+ gatherOffset.template ReinterpretCast<uint32_t>()[i * Q_MAP_SIZE], 0, Q_MAP_SIZE);
616+ }
617+#else
618+ AscendC::Gather<T>(dqState, qMap, gatherOffset.template ReinterpretCast<uint32_t>(), 0, singleSize);
619+#endif
620+ 
621+ uint32_t srcShape[2] = {static_cast<uint32_t>(singleBlockNum), 1};
622+ uint32_t dstShape[2] = {static_cast<uint32_t>(singleBlockNum), blockSize};
623+ AscendC::Broadcast<T, BROADCAST_DIM2, BROADCAST_AXIS1>(absMaxBrcb, absMax, dstShape, srcShape);
624+ AscendC::PipeBarrier<PIPE_V>();
625+ if (validSize < static_cast<uint32_t>(singleSize)) {
626+ AscendC::Duplicate<T>(state.template ReinterpretCast<T>(), static_cast<T>(0), singleSize);
627+ AscendC::PipeBarrier<PIPE_V>();
628+ }
629+ AscendC::Mul<T>(state.template ReinterpretCast<T>(), dqState, absMaxBrcb, validSize);
630+ AscendC::PipeBarrier<PIPE_V>();
631+ }
632+ 
633+ __aicore__ inline void UpdateStateAndParamFp32(const AscendC::LocalTensor<T>& dqStateM,
634+ const AscendC::LocalTensor<T>& dqStateV,
635+ AscendC::LocalTensor<T>& grad, AscendC::LocalTensor<T>& var,
636+ uint32_t validSize)
637+ {
638+ constexpr uint32_t oneRepeatSize = AscendC::GetVecLen() / sizeof(T);
639+ uint16_t repeatTimes = AscendC::CeilDivision(validSize, oneRepeatSize);
640+ float weightDecayFactor = weightDecay > 0.0f ? 1.0f - lr * weightDecay : 1.0f;
641+ asc_vf_call<AdamWQuantUpdateVF<T>>((__ubuf__ T*)dqStateM.GetPhyAddr(), (__ubuf__ T*)dqStateV.GetPhyAddr(),
642+ (__ubuf__ T*)var.GetPhyAddr(), (__ubuf__ T*)grad.GetPhyAddr(), gnormScale,
643+ beta1, 1.0f - beta1, beta2, 1.0f - beta2, eps * correction2, stepSize,
644+ weightDecayFactor, validSize, oneRepeatSize, repeatTimes);
645+ }
646+ 
647+private:
648+ TPipe pipe;
TangPC
TangPCTangPC8月1日

TPipe 放在算子类内部是有明确代价的，实测 scalar 耗时会多出一成七左右，规范里要求在 kernel 入口创建、类里只存指针。apply_adam_w_quant.cpp 那个入口现在没有创建 TPipe，建议把 pipe 提到入口，Init 改成收一个 TPipe*。fp16.h 同样。文件头注释写着「kernel 入口风格对齐 norm/deep_norm」，但 deep_norm 是在入口建 pipe 的。

likedislike
649+ 
650+ AscendC::TBuf<AscendC::QuePosition::VECCALC> varBuf;
651+ AscendC::TBuf<AscendC::QuePosition::VECCALC> gradBuf;
652+ AscendC::TBuf<AscendC::QuePosition::VECCALC> stateMBuf;
653+ AscendC::TBuf<AscendC::QuePosition::VECCALC> stateVBuf;
654+ AscendC::TBuf<AscendC::QuePosition::VECCALC> qMapMBuf;
655+ AscendC::TBuf<AscendC::QuePosition::VECCALC> qMapVBuf;
656+ AscendC::TBuf<AscendC::QuePosition::VECCALC> absMaxMBuf;
657+ AscendC::TBuf<AscendC::QuePosition::VECCALC> absMaxVBuf;
658+ AscendC::TBuf<AscendC::QuePosition::VECCALC> calcBuf;
659+ 
660+ AscendC::GlobalTensor<T> varGm;
661+ AscendC::GlobalTensor<T> varRefGm;
662+ AscendC::GlobalTensor<T> gradGm;
663+ AscendC::GlobalTensor<uint8_t> stateMGm;
664+ AscendC::GlobalTensor<uint8_t> stateMRefGm;
665+ AscendC::GlobalTensor<uint8_t> stateVGm;
666+ AscendC::GlobalTensor<uint8_t> stateVRefGm;
667+ AscendC::GlobalTensor<T> qMapMGm;
668+ AscendC::GlobalTensor<T> qMapVGm;
669+ AscendC::GlobalTensor<T> absMaxMGm;
670+ AscendC::GlobalTensor<T> absMaxMRefGm;
671+ AscendC::GlobalTensor<T> absMaxVGm;
672+ AscendC::GlobalTensor<T> absMaxVRefGm;
673+ AscendC::GlobalTensor<U> stepGm;
674+ 
675+ AscendC::LocalTensor<T> qMapM;
676+ AscendC::LocalTensor<T> qMapV;
677+ AscendC::LocalTensor<uint8_t> stateM;
678+ AscendC::LocalTensor<uint8_t> stateV;
679+ AscendC::LocalTensor<T> absMaxM;
680+ AscendC::LocalTensor<T> absMaxV;
681+ AscendC::LocalTensor<T> var;
682+ AscendC::LocalTensor<T> grad;
683+ 
684+ int32_t blockIdx;
685+ float step_ = 0;
686+ 
687+ int32_t singleBlockNum;
688+ int32_t singleSize;
689+ uint64_t useNumCore = 0;
690+ uint64_t lastPreCoreRowWork = 0;
691+ uint64_t notLastCoreNum = 0;
692+ uint64_t notLastPreCoreRowWork = 0;
693+ uint64_t lastCoreLastBlock = 0;
694+ uint64_t lastBlockSize = 0;
695+ int64_t oneCoreDoBlockNumPerRow = 0;
696+ 
697+ float stepSize;
698+ float correction2;
699+ float lr;
700+ float beta1;
701+ float beta2;
702+ float weightDecay;
703+ float eps;
704+ float gnormScale;
705+ uint32_t blockSize;
706+ 
707+ uint64_t gmOffset = 0;
708+ uint64_t absmaxOffset = 0;
709+};
710+} // namespace ApplyAdamWQuantNS
711+#endif // APPLY_ADAM_W_QUANT_FP32_H_
Aoptim/apply_adam_w_quant/op_kernel/arch35/apply_adam_w_quant_tiling_data.h+47-0
@@ -0,0 +1,47 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ *
10+ * NOTE: Portions of this code were AI-generated and have been
11+ * technically reviewed for functional accuracy and security
12+ */
13+ 
14+/*!
15+ * \file apply_adam_w_quant_tiling_data.h
16+ * \brief ApplyAdamWQuant arch35 (Ascend950) plain tiling-data struct.
17+ *
18+ * arch35 走 regbase 出包路径(host GetTilingData<T> 直写裸 buffer + kernel
19+ * GET_TILING_DATA_WITH_STRUCT 直读),故用 plain POD,字段名/顺序/类型与 A2 的
20+ * framework 版 ApplyAdamWQuantTilingData 一一对应,使 arch35 的计算头(从 A2 拷贝)
21+ * 无需改动成员访问即可复用。参考 deep_norm/apply_adam_w_v2 的 arch35 tiling_data 布局。
22+ */
23+ 
24+#ifndef APPLY_ADAM_W_QUANT_ARCH35_TILING_DATA_H
25+#define APPLY_ADAM_W_QUANT_ARCH35_TILING_DATA_H
26+ 
27+#include <cstdint>
28+ 
29+struct ApplyAdamWQuantRegbaseTilingData {
30+ uint64_t use_num_core = 0; // 总共使用的核数
范其瑞
范其瑞范其瑞7月31日

基础命名风格问题，请遵守华为编程规范

likedislike
daiwei18
daiwei18
8月1日 评论：
31+ uint64_t last_pre_core_row_work = 0; // 尾核一个核循环的个数
32+ uint64_t not_last_core_num = 0; // 非尾核的个数
33+ uint64_t not_last_pre_core_row_work = 0; // 非尾核一个核循环的个数
34+ uint64_t last_core_last_block = 0; // 最后一个核最后一次循环的 block 个数
35+ float lr = 0.0f;
36+ float beta1 = 0.0f;
37+ float beta2 = 0.0f;
38+ float weight_decay = 0.0f;
39+ float eps = 0.0f;
40+ float gnorm_scale = 0.0f;
41+ int64_t block_size = 0;
42+ uint64_t one_core_do_block_num_per_row = 0;
43+ uint64_t tiling_key = 0;
范其瑞
范其瑞范其瑞7月31日

不要把tilingKey放到tilingData中

likedislike
daiwei18
daiwei18
8月1日 评论：
44+ uint64_t last_block_size = 0; // 最后一个量化 block 的有效元素数
45+};
46+ 
47+#endif // APPLY_ADAM_W_QUANT_ARCH35_TILING_DATA_H
Aoptim/apply_adam_w_quant/tests/assets/golden.py+139-0
@@ -0,0 +1,139 @@
1+#!/usr/bin/env python3
2+# -*- coding: UTF-8 -*-
3+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
4+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5+# CANN Open Software License Agreement Version 2.0 (the "License").
6+# Please refer to the License for details. You may not use this file except in compliance with the License.
7+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9+# See LICENSE in the root of the software repository for the full text of the License.
10+ 
11+"""Torch CPU golden for ApplyAdamWQuant.
12+ 
13+NumPy is used only for TTK array conversion and input sanitization. The golden
14+math itself is expressed with Torch tensor operations.
15+"""
16+ 
17+import numpy as np
18+import torch
19+ 
20+BLOCK = 256
21+ 
22+__spec__ = {"apply_adam_w_quant": "ApplyAdamWQuantTestSpec"}
23+__golden__ = {"kernel": {"apply_adam_w_quant": "apply_adam_w_quant_golden"}}
24+ 
25+ 
26+def _to_f32(x):
27+ return torch.as_tensor(np.asarray(x).astype(np.float32), dtype=torch.float32)
28+ 
29+ 
30+def _quantize_state(state, qmap, absmax):
31+ one = torch.tensor(1.0, dtype=torch.float32)
32+ denom = torch.where(absmax == 0, one, absmax)
33+ normalized = state / denom
34+ distances = torch.abs(normalized[:, None] - qmap[None, :])
35+ indices = torch.argmin(distances, dim=1).to(torch.int64)
36+ if bool(torch.any(qmap < 0)):
37+ quantized = qmap[indices]
38+ flip = normalized * quantized < 0
39+ adjusted = indices.clone()
40+ adjusted = torch.where(flip & (normalized > 0), adjusted + 1, adjusted)
41+ adjusted = torch.where(flip & (normalized < 0), adjusted - 1, adjusted)
42+ indices = torch.clamp(adjusted, 0, qmap.numel() - 1)
43+ return indices
44+ 
45+ 
46+def apply_adam_w_quant_golden(
47+ var, grad, m, v, qmap_m, qmap_v, absmax_m, absmax_v, step, **kwargs
48+):
49+ lr = torch.tensor(float(kwargs.get("lr", 0.001)), dtype=torch.float32)
50+ beta1 = torch.tensor(float(kwargs.get("beta1", 0.9)), dtype=torch.float32)
51+ beta2 = torch.tensor(float(kwargs.get("beta2", 0.999)), dtype=torch.float32)
52+ weight_decay = torch.tensor(
53+ float(kwargs.get("weight_decay", 1.0)), dtype=torch.float32
54+ )
55+ eps = torch.tensor(float(kwargs.get("eps", 1e-8)), dtype=torch.float32)
56+ gnorm_scale = torch.tensor(
57+ float(kwargs.get("gnorm_scale", 1.0)), dtype=torch.float32
58+ )
59+ block_size = int(kwargs.get("block_size", BLOCK))
60+ 
61+ out_dtype = np.asarray(var).dtype
62+ one = torch.tensor(1.0, dtype=torch.float32)
63+ step_v = _to_f32(step).reshape(-1)[0] + one
64+ bias_c1 = one - torch.pow(beta1, step_v)
65+ bias_c2_sqrt = torch.sqrt(one - torch.pow(beta2, step_v))
66+ step_size = -lr * bias_c2_sqrt / bias_c1
67+ one_minus_beta1 = one - beta1
68+ one_minus_beta2 = one - beta2
69+ weight_decay_factor = one - lr * weight_decay
70+ 
71+ var_f = _to_f32(var).reshape(-1).clone()
72+ grad_f = _to_f32(grad).reshape(-1)
73+ m_codes = torch.as_tensor(np.asarray(m).reshape(-1).copy(), dtype=torch.int64)
74+ v_codes = torch.as_tensor(np.asarray(v).reshape(-1).copy(), dtype=torch.int64)
75+ qmap_m_t = _to_f32(qmap_m).reshape(-1)
76+ qmap_v_t = _to_f32(qmap_v).reshape(-1)
77+ absmax_m_t = _to_f32(absmax_m).reshape(-1)
78+ absmax_v_t = _to_f32(absmax_v).reshape(-1)
79+ new_am = torch.zeros_like(absmax_m_t, dtype=torch.float32)
80+ new_av = torch.zeros_like(absmax_v_t, dtype=torch.float32)
81+ 
82+ n = var_f.numel()
83+ num_blocks = (n + block_size - 1) // block_size
84+ for block_idx in range(num_blocks):
85+ start = block_idx * block_size
86+ end = min((block_idx + 1) * block_size, n)
87+ param = var_f[start:end]
88+ grad_block = grad_f[start:end] * gnorm_scale
89+ s1 = qmap_m_t[m_codes[start:end]] * absmax_m_t[block_idx]
90+ s2 = qmap_v_t[v_codes[start:end]] * absmax_v_t[block_idx]
91+ s1_update = s1 * beta1 + one_minus_beta1 * grad_block
92+ s2_update = s2 * beta2 + one_minus_beta2 * (grad_block * grad_block)
93+ denom = torch.sqrt(s2_update) + eps * bias_c2_sqrt
94+ param = param + step_size * s1_update / denom
95+ if float(weight_decay) > 0.0:
96+ param = param * weight_decay_factor
97+ 
98+ abs_m = torch.max(torch.abs(s1_update))
99+ abs_v = torch.max(s2_update)
100+ var_f[start:end] = param
101+ m_codes[start:end] = _quantize_state(s1_update, qmap_m_t, abs_m)
102+ v_codes[start:end] = _quantize_state(s2_update, qmap_v_t, abs_v)
103+ new_am[block_idx] = abs_m
104+ new_av[block_idx] = abs_v
105+ 
106+ var_out = var_f.reshape(np.asarray(var).shape).numpy().astype(out_dtype, copy=False)
107+ return [
108+ var_out,
109+ m_codes.numpy()
110+ .astype(np.asarray(m).dtype, copy=False)
111+ .reshape(np.asarray(m).shape),
112+ v_codes.numpy()
113+ .astype(np.asarray(v).dtype, copy=False)
114+ .reshape(np.asarray(v).shape),
115+ new_am.numpy()
116+ .astype(np.float32, copy=False)
117+ .reshape(np.asarray(absmax_m).shape),
118+ new_av.numpy()
119+ .astype(np.float32, copy=False)
120+ .reshape(np.asarray(absmax_v).shape),
121+ ]
122+ 
123+ 
124+def customize_inputs(
125+ var, grad, m, v, qmap_m, qmap_v, absmax_m, absmax_v, step, **kwargs
126+):
127+ qmap_m = np.linspace(-1.0, 1.0, 256, dtype=np.float32)
128+ qmap_v = np.linspace(0.0, 1.0, 256, dtype=np.float32)
129+ m = np.clip(m, 0, 255).astype(np.uint8)
130+ v = np.clip(v, 0, 255).astype(np.uint8)
131+ absmax_m = np.abs(absmax_m).astype(np.float32) + 0.1
132+ absmax_v = np.abs(absmax_v).astype(np.float32) + 0.1
133+ step = np.maximum(np.asarray(step), 1)
134+ return var, grad, m, v, qmap_m, qmap_v, absmax_m, absmax_v, step
135+ 
136+ 
137+class ApplyAdamWQuantTestSpec:
138+ golden = apply_adam_w_quant_golden
139+ customize_inputs = customize_inputs
Roptim/apply_adam_w_quant/tests/ut/op_kernel/test_apply_adam_w_quant_tiling_def.h→optim/apply_adam_w_quant/tests/ut/op_kernel/apply_adam_w_quant_tiling_def.h+17-9