Line data Source code
1 : /**
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
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 quant_update_scatter_tiling_arch35.cpp
13 : * \brief quant_update_scatter_regbase tiling file
14 : */
15 :
16 : #include "quant_update_scatter_tiling_arch35.h"
17 : #include "register/op_impl_registry.h"
18 : #include "log/log.h"
19 : #include "platform/platform_info.h"
20 : #include "index/quant_update_scatter/op_kernel/arch35/quant_update_scatter_struct.h"
21 : #include "util/math_util.h"
22 : #include "atvoss/broadcast/broadcast_tiling.h"
23 :
24 : namespace optiling {
25 : using namespace QuantUpdateScatter;
26 : using namespace std;
27 :
28 : const set<ge::DataType> INPUT_VAR_SUPPORT_DTYPE_SET = {ge::DT_INT8, ge::DT_HIFLOAT8, ge::DT_FLOAT8_E4M3FN,
29 : ge::DT_FLOAT8_E5M2};
30 : const set<ge::DataType> INPUT_INDICES_SUPPORT_DTYPE_SET = {ge::DT_INT32, ge::DT_INT64};
31 : const set<ge::DataType> INPUT_UPDATES_SUPPORT_DTYPE_SET = {ge::DT_BF16, ge::DT_FLOAT16};
32 : const set<ge::DataType> INPUT_SCALE_SUPPORT_DTYPE_SET = {ge::DT_BF16, ge::DT_FLOAT};
33 : const set<ge::DataType> INPUT_ZERO_POINT_SUPPORT_DTYPE_SET = {ge::DT_BF16, ge::DT_INT32};
34 : const set<ge::DataType> OUTPUT_VAR_SUPPORT_DTYPE_SET = {ge::DT_INT8, ge::DT_HIFLOAT8, ge::DT_FLOAT8_E4M3FN,
35 : ge::DT_FLOAT8_E5M2};
36 : const map<ge::DataType, vector<string>> DTYPE_ROUND_MODE_MAP = {{ge::DT_INT8, {"rint"}},
37 : {ge::DT_HIFLOAT8, {"round", "hybrid"}},
38 : {ge::DT_FLOAT8_E4M3FN, {"rint"}},
39 : {ge::DT_FLOAT8_E5M2, {"rint"}}};
40 :
41 : const map<ge::DataType, string> DTYPE_ROUND_MODE_LOG_MAP = {
42 : {ge::DT_INT8, "int8 datatype only support 'rint', currently is: "},
43 : {ge::DT_HIFLOAT8, "hifloat8 datatype only support 'round' and 'hybrid', currently is: "},
44 : {ge::DT_FLOAT8_E4M3FN, "float8_e4m3fn datatype only support 'rint', currently is: "},
45 : {ge::DT_FLOAT8_E5M2, "float8_e5m2 datatype only support 'rint', currently is: "}};
46 :
47 : const map<string, uint64_t> ROUND_MODE_TPL_MAP = {
48 : {"rint", TPL_ROUND_MODE_RINT}, {"round", TPL_ROUND_MODE_ROUND}, {"hybrid", TPL_ROUND_MODE_HYBRID}};
49 :
50 : int64_t QuantUpdateScatterRegbaseTiling::NewAxis(int64_t axis) const
51 : {
52 : int64_t newAxis = axis < 0 ? (oldDims_ + axis) : axis;
53 : if (0 < newAxis && newAxis < oldDims_ - 1) {
54 : newAxis = static_cast<int64_t>(DIM_2);
55 : }
56 : return newAxis;
57 : }
58 :
59 : double QuantUpdateScatterRegbaseTiling::GetUpdateUbRatio(bool isLittleQuant) const
60 : {
61 : int64_t totalPart = varDtypeSize_ + updateDtypeSize_;
62 : if (isLittleQuant) {
63 : totalPart += quantScalesDtypeSize_;
64 : if (zeroPointsType_ != TPL_NONE) {
65 : totalPart += quantZeroPointsDtypeSize_;
66 : }
67 : }
68 : double ratio = updateDtypeSize_ * 1.0 / totalPart;
69 : return ratio;
70 : }
71 :
72 : bool QuantUpdateScatterRegbaseTiling::CheckRoundMode(ge::DataType type, string mode) const
73 : {
74 : auto it = DTYPE_ROUND_MODE_MAP.find(type);
75 : if (it == DTYPE_ROUND_MODE_MAP.end()) {
76 : return false;
77 : }
78 : if (find(it->second.begin(), it->second.end(), mode) == it->second.end()) {
79 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "round_mode", mode, GetErrMsg(type));
80 : return false;
81 : }
82 : return true;
83 : }
84 :
85 : string QuantUpdateScatterRegbaseTiling::GetErrMsg(ge::DataType type) const
86 : {
87 : auto it = DTYPE_ROUND_MODE_LOG_MAP.find(type);
88 : if (it != DTYPE_ROUND_MODE_LOG_MAP.end()) {
89 : return it->second;
90 : } else {
91 : return "Wrong data type, round mode: %s";
92 : }
93 : }
94 :
95 : void QuantUpdateScatterRegbaseTiling::CalcTilingDataForLargeBatchLargeQuant()
96 : {
97 : tilingData_.set_innerLoopEle(maxUpdatesSize_ / BYTES_ONE_BLOCK * BYTES_ONE_BLOCK / updateDtypeSize_ / BUFFER_NUM);
98 : tilingData_.set_innerLoopFullRpt(0);
99 : if (tilingData_.get_innerLoopEle() == 0) {
100 : OP_LOGE(context_->GetNodeName(), "innerLoopEle is 0");
101 : return;
102 : }
103 : // 核内切update[3]
104 : tilingData_.set_innerLoopTimes(tilingData_.get_updateOriLastDim() / tilingData_.get_innerLoopEle());
105 : tilingData_.set_innerLoopTail(tilingData_.get_updateOriLastDim() % tilingData_.get_innerLoopEle());
106 : tilingData_.set_innerLoopTailRpt(0);
107 : tilingData_.set_innerLoopTimesLastCore(0);
108 : tilingData_.set_innerLoopTailLastCore(0);
109 : tilingData_.set_innerLoopTailRptLastCore(0);
110 :
111 : return;
112 : }
113 :
114 : void QuantUpdateScatterRegbaseTiling::CalcTilingDataForLargeBatchLittleQuant()
115 : {
116 : int64_t updateDim3Align = tilingData_.get_updateDim3() / tilingData_.get_updateOriLastDim() *
117 : tilingData_.get_updateOriLastDimAlign();
118 : int64_t innerLoopEle = maxUpdatesSize_ / BYTES_ONE_BLOCK * BYTES_ONE_BLOCK / updateDtypeSize_ / BUFFER_NUM /
119 : updateDim3Align * updateDim3Align;
120 : tilingData_.set_innerLoopFullRpt(innerLoopEle / updateDim3Align);
121 : if (tilingData_.get_innerLoopFullRpt() == 0) {
122 : OP_LOGE(context_->GetNodeName(), "innerLoopFullRpt is 0");
123 : return;
124 : }
125 : tilingData_.set_innerLoopEle(tilingData_.get_innerLoopFullRpt() * tilingData_.get_updateOriLastDim());
126 : tilingData_.set_innerLoopTimes(updateNewShape_.GetDim(DIM_2) * updateNewShape_.GetDim(DIM_3) /
127 : tilingData_.get_updateOriLastDim() / tilingData_.get_innerLoopFullRpt());
128 : tilingData_.set_innerLoopTailRpt(updateNewShape_.GetDim(DIM_2) * updateNewShape_.GetDim(DIM_3) /
129 : tilingData_.get_updateOriLastDim() % tilingData_.get_innerLoopFullRpt());
130 : tilingData_.set_innerLoopTail(tilingData_.get_innerLoopTailRpt() * tilingData_.get_updateOriLastDim());
131 : tilingData_.set_innerLoopTimesLastCore(0);
132 : tilingData_.set_innerLoopTailLastCore(0);
133 : tilingData_.set_innerLoopTailRptLastCore(0);
134 :
135 : return;
136 : }
137 :
138 : void QuantUpdateScatterRegbaseTiling::CalcTilingDataForLargeEleLargeQuant()
139 : {
140 : int64_t innerLoopEle = maxUpdatesSize_ / BYTES_ONE_BLOCK * BYTES_ONE_BLOCK / updateDtypeSize_ / BUFFER_NUM;
141 : tilingData_.set_innerLoopEle(innerLoopEle);
142 : tilingData_.set_innerLoopFullRpt(0);
143 : if (tilingData_.get_innerLoopEle() == 0) {
144 : OP_LOGE(context_->GetNodeName(), "innerLoopEle is 0");
145 : return;
146 : }
147 : // 核内切update[3]
148 : tilingData_.set_innerLoopTimes(tilingData_.get_updateOriLastDim() / tilingData_.get_innerLoopEle());
149 : tilingData_.set_innerLoopTail(tilingData_.get_updateOriLastDim() % tilingData_.get_innerLoopEle());
150 : tilingData_.set_innerLoopTailRpt(0);
151 : tilingData_.set_innerLoopTimesLastCore(0);
152 : tilingData_.set_innerLoopTailLastCore(0);
153 : tilingData_.set_innerLoopTailRptLastCore(0);
154 :
155 : return;
156 : }
157 :
158 : ge::graphStatus QuantUpdateScatterRegbaseTiling::CalcTilingDataForLargeEleLittleQuant()
159 : {
160 : int64_t updateDim3Align = tilingData_.get_updateDim3() / tilingData_.get_updateOriLastDim() *
161 : tilingData_.get_updateOriLastDimAlign();
162 : int64_t innerLoopEle = maxUpdatesSize_ / updateDtypeSize_ / BUFFER_NUM / updateDim3Align *
163 : updateDim3Align; // 一次可处理的update数
164 : if (innerLoopEle == 0) {
165 : OP_LOGE(context_->GetNodeName(), "innerLoopEle is 0");
166 : return ge::GRAPH_FAILED;
167 : }
168 : tilingData_.set_innerLoopEle(innerLoopEle);
169 : // 核内切update[2]
170 : tilingData_.set_innerLoopFullRpt(tilingData_.get_innerLoopEle() / updateDim3Align); // 一次可搬入的update[2]个数
171 : tilingData_.set_innerLoopTimes(tilingData_.get_eachCoreBsNum() * updateDim3Align /
172 : tilingData_.get_innerLoopEle()); // 循环次数
173 : tilingData_.set_innerLoopTail(tilingData_.get_eachCoreBsNum() * updateDim3Align % tilingData_.get_innerLoopEle());
174 : tilingData_.set_innerLoopTailRpt(tilingData_.get_innerLoopTail() / updateDim3Align);
175 :
176 : tilingData_.set_innerLoopTimesLastCore(tilingData_.get_lastCoreBsNum() * updateDim3Align /
177 : tilingData_.get_innerLoopEle()); // 循环次数
178 : tilingData_.set_innerLoopFullRptLastCore(tilingData_.get_innerLoopFullRpt());
179 : tilingData_.set_innerLoopTailLastCore(tilingData_.get_lastCoreBsNum() * updateDim3Align %
180 : tilingData_.get_innerLoopEle());
181 : tilingData_.set_innerLoopTailRptLastCore(tilingData_.get_innerLoopTailLastCore() / updateDim3Align);
182 :
183 : return ge::GRAPH_SUCCESS;
184 : }
185 :
186 : ge::graphStatus QuantUpdateScatterRegbaseTiling::GetTilingNeg2()
187 : {
188 : int64_t updateDim3Align = tilingData_.get_updateDim3() / tilingData_.get_updateOriLastDim() *
189 : tilingData_.get_updateOriLastDimAlign();
190 : int64_t updateDim23Align = tilingData_.get_updateDim2() * updateDim3Align;
191 : // UB放不下一个-1（aligned）& -2轴的场景
192 : if ((updateDim23Align * updateDtypeSize_ * BUFFER_NUM) > maxUpdatesSize_) {
193 : int64_t indicesNeededUb = BYTES_ONE_BLOCK;
194 : maxUpdatesSize_ = static_cast<int64_t>(
195 : (calcUbSize_ - indicesNeededUb - quantScalesUbSize_ - quantZeroPointsUbSize_) * GetUpdateUbRatio(false));
196 : if (maxUpdatesSize_ < 0) {
197 : OP_LOGD(context_->GetNodeName(), "GetTilingNeg2 maxUpdatesSize is 0");
198 : maxUpdatesSize_ = 0;
199 : }
200 : int64_t maxInnerLoopEle = maxUpdatesSize_ / BYTES_ONE_BLOCK * BYTES_ONE_BLOCK / updateDtypeSize_ / BUFFER_NUM;
201 : OP_LOGD(context_->GetNodeName(), "maxInnerLoopEle: %ld", maxInnerLoopEle);
202 : int64_t updateNewShapeDim2 = updateNewShape_.GetDim(DIM_2);
203 : if (updateNewShape_.GetDim(DIM_0) * updateNewShape_.GetDim(DIM_1) < updateNewShapeDim2) {
204 : tilingData_.set_eachCoreBsNum(Ops::Base::CeilDiv(updateNewShapeDim2, actualCoreNum_));
205 : tilingData_.set_coreNum(Ops::Base::CeilDiv(updateNewShapeDim2, tilingData_.get_eachCoreBsNum()));
206 : tilingData_.set_lastCoreBsNum(updateNewShapeDim2 -
207 : tilingData_.get_eachCoreBsNum() * (tilingData_.get_coreNum() - 1));
208 :
209 : if (maxInnerLoopEle > updateDim3Align) {
210 : splitMode_ = TPL_MODE_LARGE_ELE_LITTLE_QUANT;
211 : OP_CHECK_IF(ge::GRAPH_SUCCESS != CalcTilingDataForLargeEleLittleQuant(),
212 : OP_LOGE(context_->GetNodeName(), "CalcTilingDataForLargeEleLittleQuant failed."),
213 : return ge::GRAPH_FAILED);
214 : } else {
215 : maxUpdatesSize_ = static_cast<int64_t>((calcUbSize_ - indicesNeededUb) * GetUpdateUbRatio(true));
216 : splitMode_ = TPL_MODE_LARGE_ELE_LARGE_QUANT;
217 : CalcTilingDataForLargeEleLargeQuant();
218 : }
219 : } else {
220 : if (maxInnerLoopEle > updateDim3Align) {
221 : splitMode_ = TPL_MODE_LARGE_BATCH_LITTLE_QUANT;
222 : CalcTilingDataForLargeBatchLittleQuant();
223 : } else {
224 : maxUpdatesSize_ = static_cast<int64_t>((calcUbSize_ - indicesNeededUb) * GetUpdateUbRatio(true));
225 : splitMode_ = TPL_MODE_LARGE_BATCH_LARGE_QUANT;
226 : CalcTilingDataForLargeBatchLargeQuant();
227 : }
228 : }
229 : } else if (updateUbSize_ > maxUpdatesSize_) {
230 : splitMode_ = TPL_MODE_LARGE_BATCH;
231 : } else {
232 : splitMode_ = TPL_MODE_LITTLE_ELE_LITTLE_QUANT;
233 : }
234 :
235 : return ge::GRAPH_SUCCESS;
236 : }
237 :
238 : void QuantUpdateScatterRegbaseTiling::UpdateTilingParam()
239 : {
240 : int64_t indexBlockSize = BYTES_ONE_BLOCK / indexDtypeSize_;
241 : int64_t varBlockSize = BYTES_ONE_BLOCK / varDtypeSize_;
242 :
243 : auto totalBs = updateNewShape_.GetDim(DIM_0) * updateNewShape_.GetDim(DIM_1);
244 : tilingData_.set_eachCoreBsNum(Ops::Base::CeilDiv(totalBs, actualCoreNum_));
245 : tilingData_.set_coreNum(Ops::Base::CeilDiv(totalBs, tilingData_.get_eachCoreBsNum()));
246 : tilingData_.set_lastCoreBsNum(totalBs - tilingData_.get_eachCoreBsNum() * (tilingData_.get_coreNum() - 1));
247 : tilingData_.set_srcBsStride(updateNewShape_.GetDim(DIM_2) * updateNewShape_.GetDim(DIM_3));
248 :
249 : tilingData_.set_indexElements(indexElements_);
250 : indexUbSize_ = Ops::Base::CeilDiv(indexElements_, indexBlockSize) * indexBlockSize * indexDtypeSize_ * BUFFER_NUM;
251 : // 量化后的type对齐，防止vst时不对齐
252 : int64_t updateOriLastDim = updateOriginShape_.GetDim(updateOriginShape_.GetDimNum() - 1);
253 : if (updateOriLastDim == 0) {
254 : OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context_->GetNodeName(), "updates", std::to_string(updateOriLastDim),
255 : "last dim of updates must not be 0");
256 : return;
257 : }
258 : int64_t updateOriLastDimAligned = Ops::Base::CeilAlign(updateOriLastDim, varBlockSize);
259 : updateUbSize_ = updateNewShape_.GetDim(absAxis_) * updateNewShape_.GetDim(absQuantAxis_) / updateOriLastDim *
260 : updateOriLastDimAligned * tilingData_.get_eachCoreBsNum() * updateDtypeSize_ * BUFFER_NUM;
261 : OP_LOGD(context_->GetNodeName(), "updateUbSize_: %ld", updateUbSize_);
262 : quantScalesUbSize_ = updateOriLastDimAligned * quantScalesDtypeSize_ * BUFFER_NUM;
263 : quantZeroPointsUbSize_ = updateOriLastDimAligned * quantZeroPointsDtypeSize_ * BUFFER_NUM;
264 : tilingData_.set_updateOriLastDim(updateOriLastDim);
265 : tilingData_.set_updateOriLastDimAlign(updateOriLastDimAligned);
266 : tilingData_.set_quantScalesElements(quantScalesElements_);
267 : tilingData_.set_quantZeroPointsElements(quantZeroPointsElements_);
268 : tilingData_.set_updateDim0(updateNewShape_.GetDim(DIM_0));
269 : tilingData_.set_updateDim1(updateNewShape_.GetDim(DIM_1));
270 : tilingData_.set_updateDim2(updateNewShape_.GetDim(DIM_2));
271 : tilingData_.set_updateDim3(updateNewShape_.GetDim(DIM_3));
272 : tilingData_.set_indicesShapeRank(indicesShapeRank_);
273 : tilingData_.set_dstBsStride(varNewShape_.GetDim(DIM_2) * varNewShape_.GetDim(DIM_3));
274 : tilingData_.set_varDim1(varNewShape_.GetDim(DIM_1));
275 : tilingData_.set_varDim2(varNewShape_.GetDim(DIM_2));
276 : tilingData_.set_varDim3(varNewShape_.GetDim(DIM_3));
277 :
278 : tilingData_.set_srcFirBsStride(updateNewShape_.GetDim(DIM_1) * updateNewShape_.GetDim(DIM_2) *
279 : updateNewShape_.GetDim(DIM_3));
280 : tilingData_.set_dstFirSecBsStride(varNewShape_.GetDim(DIM_1) * varNewShape_.GetDim(DIM_2) *
281 : varNewShape_.GetDim(DIM_3));
282 :
283 : if (quantZeroPointsElements_ == 0) {
284 : zeroPointsType_ = TPL_NONE;
285 : } else {
286 : if (quantZeroPointsDtype_ == ge::DT_INT32) {
287 : zeroPointsType_ = TPL_INT32;
288 : } else if (quantZeroPointsDtype_ == ge::DT_BF16) {
289 : zeroPointsType_ = TPL_BF16;
290 : } else {
291 : zeroPointsType_ = TPL_NONE;
292 : OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "quant_zero_points",
293 : ge::TypeUtils::DataTypeToSerialString(quantZeroPointsDtype_),
294 : "[DT_BF16, DT_INT32]");
295 : }
296 : }
297 : return;
298 : }
299 :
300 : ge::graphStatus QuantUpdateScatterRegbaseTiling::GetTilingParam()
301 : {
302 : UpdateTilingParam();
303 : int64_t tilingUbReserved = Ops::Base::CeilAlign(static_cast<int64_t>(sizeof(QuantUpdateScatterTilingData)),
304 : BYTES_ONE_BLOCK) +
305 : RESERVED_BYTES;
306 : calcUbSize_ = ubSize_ - tilingUbReserved;
307 : maxUpdatesSize_ = static_cast<int64_t>((calcUbSize_ - indexUbSize_ - quantScalesUbSize_ - quantZeroPointsUbSize_) *
308 : GetUpdateUbRatio(false));
309 : if (maxUpdatesSize_ < 0) {
310 : OP_LOGD(context_->GetNodeName(), "GetTilingParam maxUpdatesSize is 0");
311 : maxUpdatesSize_ = 0;
312 : }
313 : OP_LOGD(context_->GetNodeName(), "maxUpdatesSize_: %ld", maxUpdatesSize_);
314 :
315 11 : OP_CHECK_IF(ge::GRAPH_SUCCESS != GetTilingNeg2(),
316 : OP_LOGE(context_->GetNodeName(), "GetTilingNeg2 failed, updateDim2: %ld, updateDim3: %ld.",
317 : updateNewShape_.GetDim(DIM_2), updateNewShape_.GetDim(DIM_3)),
318 : return ge::GRAPH_FAILED);
319 : return ge::GRAPH_SUCCESS;
320 : }
321 :
322 : ge::graphStatus QuantUpdateScatterRegbaseTiling::PrepareTilingParams()
323 : {
324 : // get coreNum and ubSize
325 : auto platformInfo = context_->GetPlatformInfo();
326 : OP_CHECK_NULL_WITH_CONTEXT(context_, platformInfo);
327 : auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
328 : actualCoreNum_ = ascendcPlatform.GetCoreNumAiv();
329 : uint64_t ubSizePlatform = 0;
330 : ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizePlatform);
331 : ubSize_ = ubSizePlatform;
332 :
333 : // get input_shape
334 : auto dataShape = context_->GetInputShape(INDEX_DATA);
335 : OP_CHECK_NULL_WITH_CONTEXT(context_, dataShape);
336 : auto indicesShape = context_->GetInputShape(INDEX_INDICES);
337 : OP_CHECK_NULL_WITH_CONTEXT(context_, indicesShape);
338 : auto updatesShape = context_->GetInputShape(INDEX_UPDATES);
339 : OP_CHECK_NULL_WITH_CONTEXT(context_, updatesShape);
340 : auto quantScalesShape = context_->GetInputShape(INDEX_QUANT_SCALES);
341 : OP_CHECK_NULL_WITH_CONTEXT(context_, quantScalesShape);
342 :
343 : varOriginShape_ = Ops::Base::EnsureNotScalar(dataShape->GetOriginShape());
344 : indicesOriginShape_ = Ops::Base::EnsureNotScalar(indicesShape->GetOriginShape());
345 : updateOriginShape_ = Ops::Base::EnsureNotScalar(updatesShape->GetOriginShape());
346 : quantScalesShape_ = Ops::Base::EnsureNotScalar(quantScalesShape->GetOriginShape());
347 :
348 : indexElements_ = indicesOriginShape_.GetShapeSize();
349 : quantScalesElements_ = quantScalesShape_.GetShapeSize();
350 :
351 : auto quantZeroPointsShape = context_->GetOptionalInputShape(INDEX_QUANT_ZERO_POINTS);
352 : if (quantZeroPointsShape == nullptr) {
353 : quantZeroPointsElements_ = 0;
354 : } else {
355 : quantZeroPointsShape_ = Ops::Base::EnsureNotScalar(quantZeroPointsShape->GetOriginShape());
356 : quantZeroPointsElements_ = quantZeroPointsShape_.GetShapeSize();
357 : }
358 :
359 : // get varDtypeSize and indexDtypeSize
360 : auto dataDesc = context_->GetInputDesc(INDEX_DATA);
361 : varDtype_ = dataDesc->GetDataType();
362 : varDtypeSize_ = ge::GetSizeByDataType(varDtype_);
363 :
364 : auto indicesDesc = context_->GetInputDesc(INDEX_INDICES);
365 : indexDtype_ = indicesDesc->GetDataType();
366 : indexDtypeSize_ = ge::GetSizeByDataType(indexDtype_);
367 :
368 : auto updateDesc = context_->GetInputDesc(INDEX_UPDATES);
369 : updateDtype_ = updateDesc->GetDataType();
370 : updateDtypeSize_ = ge::GetSizeByDataType(updateDtype_);
371 :
372 : auto quantScalesDesc = context_->GetInputDesc(INDEX_QUANT_SCALES);
373 : quantScalesDtype_ = quantScalesDesc->GetDataType();
374 : quantScalesDtypeSize_ = ge::GetSizeByDataType(quantScalesDtype_);
375 :
376 : auto quantZeroPointDesc = context_->GetOptionalInputDesc(INDEX_QUANT_ZERO_POINTS);
377 : if (quantZeroPointDesc == nullptr) {
378 : quantZeroPointsDtypeSize_ = 0;
379 : } else {
380 : quantZeroPointsDtype_ = quantZeroPointDesc->GetDataType();
381 : quantZeroPointsDtypeSize_ = ge::GetSizeByDataType(quantZeroPointsDtype_);
382 : }
383 :
384 : indicesShapeRank_ = indicesOriginShape_.GetDimNum();
385 : return ge::GRAPH_SUCCESS;
386 : }
387 :
388 : ge::graphStatus QuantUpdateScatterRegbaseTiling::VerifyNullTenosr() const
389 : {
390 : OP_CHECK_IF(varOriginShape_.GetDimNum() != updateOriginShape_.GetDimNum(),
391 : OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
392 : context_->GetNodeName(), "var, updates",
393 : std::to_string(varOriginShape_.GetDimNum()) + ", " + std::to_string(updateOriginShape_.GetDimNum()),
394 : "The shape dim of var must be the same as the shape dim of updates"),
395 : return ge::GRAPH_FAILED);
396 :
397 : OP_CHECK_IF(varOriginShape_.GetDimNum() * indicesShapeRank_ == 0,
398 : OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(
399 : context_->GetNodeName(), "var, indices",
400 : std::to_string(varOriginShape_.GetDimNum()) + ", " + std::to_string(indicesShapeRank_),
401 : "The shape dim of var and indices must not be 0"),
402 : return ge::GRAPH_FAILED);
403 :
404 : int64_t dataNum = varOriginShape_.GetShapeSize();
405 : int64_t indicesNum = indicesOriginShape_.GetShapeSize();
406 : int64_t updateNum = updateOriginShape_.GetShapeSize();
407 : int64_t quantScalesNum = quantScalesElements_;
408 : OP_CHECK_IF(dataNum == 0 || indicesNum == 0 || updateNum == 0 || quantScalesNum == 0,
409 : OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(
410 : context_->GetNodeName(), "var, indices, updates, quant_scales",
411 : std::to_string(dataNum) + ", " + std::to_string(indicesNum) + ", " + std::to_string(updateNum) +
412 : ", " + std::to_string(quantScalesNum),
413 : "var, indices, updates, and quant_scales do not support empty tensor"),
414 : return ge::GRAPH_FAILED);
415 :
416 : return ge::GRAPH_SUCCESS;
417 : }
418 :
419 : ge::graphStatus QuantUpdateScatterRegbaseTiling::VerifyParamsDtype() const
420 : {
421 : OP_CHECK_IF(
422 : INPUT_VAR_SUPPORT_DTYPE_SET.count(varDtype_) == 0,
423 : OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "var", ge::TypeUtils::DataTypeToSerialString(varDtype_),
424 : "[DT_INT8, DT_HIFLOAT8, DT_FLOAT8_E4M3FN, DT_FLOAT8_E5M2]"),
425 : return ge::GRAPH_FAILED);
426 : OP_CHECK_IF(INPUT_INDICES_SUPPORT_DTYPE_SET.count(indexDtype_) == 0,
427 : OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "indices",
428 : ge::TypeUtils::DataTypeToSerialString(indexDtype_), "[DT_INT32, DT_INT64]"),
429 : return ge::GRAPH_FAILED);
430 : OP_CHECK_IF(INPUT_UPDATES_SUPPORT_DTYPE_SET.count(updateDtype_) == 0,
431 : OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "updates",
432 : ge::TypeUtils::DataTypeToSerialString(updateDtype_), "[DT_BF16, DT_FLOAT16]"),
433 : return ge::GRAPH_FAILED);
434 : OP_CHECK_IF(
435 : INPUT_SCALE_SUPPORT_DTYPE_SET.count(quantScalesDtype_) == 0,
436 : OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "quant_scales",
437 : ge::TypeUtils::DataTypeToSerialString(quantScalesDtype_), "[DT_BF16, DT_FLOAT]"),
438 : return ge::GRAPH_FAILED);
439 : if (quantZeroPointsDtypeSize_ != 0) {
440 : OP_CHECK_IF(INPUT_ZERO_POINT_SUPPORT_DTYPE_SET.count(quantZeroPointsDtype_) == 0,
441 : OP_LOGE_FOR_INVALID_DTYPE(context_->GetNodeName(), "quant_zero_points",
442 : ge::TypeUtils::DataTypeToSerialString(quantZeroPointsDtype_),
443 : "[DT_BF16, DT_INT32]"),
444 : return ge::GRAPH_FAILED);
445 : }
446 :
447 : return ge::GRAPH_SUCCESS;
448 : }
449 :
450 : ge::graphStatus QuantUpdateScatterRegbaseTiling::VerifyTilingQuantParams()
451 : {
452 : int64_t updateDimNum = updateOriginShape_.GetDimNum();
453 : OP_CHECK_IF(
454 : (updateDimNum < 3 || updateDimNum > 8),
455 : OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context_->GetNodeName(), "updates", std::to_string(updateDimNum),
456 : "The shape dim of updates must be within the range [3, 8]"),
457 : return ge::GRAPH_FAILED);
458 : int64_t dataDimNum = varOriginShape_.GetDimNum();
459 : OP_CHECK_IF(
460 : (updateDimNum != dataDimNum),
461 : OP_LOGE_FOR_INVALID_SHAPEDIMS_WITH_REASON(context_->GetNodeName(), "updates, var",
462 : std::to_string(updateDimNum) + ", " + std::to_string(dataDimNum),
463 : "The shape dim of updates must be the same as the shape dim of var"),
464 : return ge::GRAPH_FAILED);
465 :
466 : auto attrs = context_->GetAttrs();
467 : OP_CHECK_NULL_WITH_CONTEXT(context_, attrs);
468 : // reduce attribute
469 : const char* reduceAxisPtr = attrs->GetAttrPointer<char>(ATTR_REDUCE_INDEX);
470 : string reduceAxis(reduceAxisPtr);
471 :
472 : OP_CHECK_IF(reduceAxis != "update",
473 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "reduce", reduceAxis,
474 : "The value of reduce must be [update]"),
475 : return ge::GRAPH_FAILED);
476 :
477 : // axis attribute
478 : const auto axisPtr = attrs->GetAttrPointer<int64_t>(ATTR_AXIS_INDEX);
479 : int64_t axis = (axisPtr == nullptr) ? -2 : *axisPtr;
480 : absAxis_ = (axis < 0) ? axis + updateDimNum : axis;
481 : OP_CHECK_IF((absAxis_ > updateDimNum - 2 || absAxis_ < 1),
482 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "axis", std::to_string(axis),
483 : "The value of axis must be within the range [1, " +
484 : std::to_string(updateDimNum - 2) + "] or [" +
485 : std::to_string(1 - updateDimNum) + ", -2]"),
486 : return ge::GRAPH_FAILED);
487 :
488 : // quant_axis attribute
489 : const auto quantAxisPtr = attrs->GetAttrPointer<int64_t>(ATTR_QUANT_AXIS_INDEX);
490 : int64_t quantAxis = (quantAxisPtr == nullptr) ? -1 : *quantAxisPtr;
491 :
492 : absQuantAxis_ = (quantAxis < 0) ? quantAxis + updateDimNum : quantAxis;
493 : OP_CHECK_IF((absQuantAxis_ != updateDimNum - 1),
494 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(
495 : context_->GetNodeName(), "quant_axis", std::to_string(quantAxis),
496 : "The value of quant_axis must be -1 or " + std::to_string(updateDimNum - 1)),
497 : return ge::GRAPH_FAILED);
498 :
499 : // reciprocal attribute
500 : const auto reciprocalPtr = attrs->GetAttrPointer<bool>(ATTR_RECIPROCAL_INDEX);
501 : bool reciprocal = (reciprocalPtr == nullptr) ? false : *reciprocalPtr;
502 : divMode_ = reciprocal ? TPL_DIV_MODE_MUL : TPL_DIV_MODE_DIV;
503 :
504 : // round_mode attribute
505 : const char* roundModePtr = attrs->GetAttrPointer<char>(ATTR_ROUND_MODE_INDEX);
506 : string roundMode((roundModePtr) ? roundModePtr : "rint");
507 : OP_CHECK_IF(!CheckRoundMode(varDtype_, roundMode),
508 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "round_mode", roundMode,
509 : "The value of round_mode must be [rint, round, hybrid]"),
510 : return ge::GRAPH_FAILED);
511 :
512 : auto it = ROUND_MODE_TPL_MAP.find(roundMode);
513 : if (it != ROUND_MODE_TPL_MAP.end()) {
514 : castRoundMode_ = it->second;
515 : } else {
516 : OP_LOGE_FOR_INVALID_VALUE_WITH_REASON(context_->GetNodeName(), "round_mode", roundMode,
517 : "The value of round_mode must be [rint, round, hybrid]");
518 : return ge::GRAPH_FAILED;
519 : }
520 : int64_t quantScalesNum = quantScalesElements_;
521 : int64_t quantZeroPointsElements = quantZeroPointsElements_;
522 : OP_CHECK_IF((quantScalesNum != updateOriginShape_.GetDim(absQuantAxis_)),
523 : OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(
524 : context_->GetNodeName(), "quant_scales, updates",
525 : std::to_string(quantScalesNum) + ", " + std::to_string(updateOriginShape_.GetDim(absQuantAxis_)),
526 : "The shape size of quant_scales must equal the last dim of updates"),
527 : return ge::GRAPH_FAILED);
528 : OP_CHECK_IF((quantScalesNum != quantZeroPointsElements) && (quantZeroPointsElements != 0),
529 : OP_LOGE_FOR_INVALID_SHAPESIZES_WITH_REASON(
530 : context_->GetNodeName(), "quant_scales, quant_zero_points",
531 : std::to_string(quantScalesNum) + ", " + std::to_string(quantZeroPointsElements),
532 : "The shape size of quant_scales must be the same as the shape size of quant_zero_points"),
533 : return ge::GRAPH_FAILED);
534 :
535 : return ge::GRAPH_SUCCESS;
536 : }
537 :
538 : ge::graphStatus QuantUpdateScatterRegbaseTiling::MergeDims()
539 : {
540 : oldDims_ = varOriginShape_.GetDimNum();
541 :
542 : varNewShape_.SetDimNum(0);
543 : updateNewShape_.SetDimNum(0);
544 :
545 : varNewShape_.AppendDim(varOriginShape_[0]);
546 : updateNewShape_.AppendDim(updateOriginShape_[0]);
547 :
548 : size_t dataSecondDims = 1;
549 : size_t updataSecondDims = 1;
550 : for (int64_t i = 1; i < absAxis_; i++) {
551 : dataSecondDims *= varOriginShape_[i];
552 : updataSecondDims *= updateOriginShape_[i];
553 : }
554 : varNewShape_.AppendDim(dataSecondDims);
555 : updateNewShape_.AppendDim(updataSecondDims);
556 :
557 : varNewShape_.AppendDim(varOriginShape_[absAxis_]);
558 : updateNewShape_.AppendDim(updateOriginShape_[absAxis_]);
559 :
560 : size_t dataFourthDims = 1;
561 : size_t updataFourthDims = 1;
562 : for (int64_t i = absAxis_ + 1; i < oldDims_; i++) {
563 : dataFourthDims *= varOriginShape_[i];
564 : updataFourthDims *= updateOriginShape_[i];
565 : }
566 : absAxis_ = DIM_2;
567 : absQuantAxis_ = DIM_3;
568 : varNewShape_.AppendDim(dataFourthDims);
569 : updateNewShape_.AppendDim(updataFourthDims);
570 :
571 : return ge::GRAPH_SUCCESS;
572 : }
573 :
574 : ge::graphStatus QuantUpdateScatterRegbaseTiling::VerifyTilingParams() const
575 : {
576 : OP_CHECK_IF(updateOriginShape_[0] != indicesOriginShape_[0],
577 : OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(
578 : context_->GetNodeName(), "updates, indices",
579 : Ops::Base::ToString(updateOriginShape_) + ", " + Ops::Base::ToString(indicesOriginShape_),
580 : "dim[0] of updates must be equal to dim[0] of indices"),
581 : return ge::GRAPH_FAILED);
582 : OP_CHECK_IF(updateOriginShape_[0] > varOriginShape_[0],
583 : OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "updates",
584 : Ops::Base::ToString(updateOriginShape_),
585 : "dim[0] of updates must be less than dim[0] of var"),
586 : return ge::GRAPH_FAILED);
587 : OP_CHECK_IF(updateOriginShape_[absAxis_] > varOriginShape_[absAxis_],
588 : OP_LOGE_FOR_INVALID_SHAPE_WITH_REASON(context_->GetNodeName(), "updates",
589 : Ops::Base::ToString(updateOriginShape_),
590 : "dim[axis] of updates must be less than dim[axis] of var"),
591 : return ge::GRAPH_FAILED);
592 :
593 : for (int64_t i = 1; i < static_cast<int64_t>(updateOriginShape_.GetDimNum()); i++) {
594 : if (i == absAxis_) {
595 : continue;
596 : }
597 : if (updateOriginShape_[i] != varOriginShape_[i]) {
598 : OP_LOGE_FOR_INVALID_SHAPES_WITH_REASON(
599 : context_->GetNodeName(), "updates, var",
600 : Ops::Base::ToString(updateOriginShape_) + ", " + Ops::Base::ToString(varOriginShape_),
601 : "dim[" + std::to_string(i) + "] of updates must be equal to dim[" + std::to_string(i) + "] of var");
602 : return ge::GRAPH_FAILED;
603 : }
604 : }
605 :
606 : if (indicesShapeRank_ != ONE_INDICES && indicesShapeRank_ != TWO_INDICES) {
607 : OP_LOGE_FOR_INVALID_SHAPEDIM_WITH_REASON(context_->GetNodeName(), "indices", std::to_string(indicesShapeRank_),
608 : "The shape dim of indices must be 1 or 2");
609 : return ge::GRAPH_FAILED;
610 : }
611 :
612 : if (indicesShapeRank_ == TWO_INDICES) {
613 : OP_CHECK_IF(indicesOriginShape_[1] != 2,
614 : OP_LOGE_FOR_INVALID_SHAPEDIM(context_->GetNodeName(), "indices",
615 : std::to_string(indicesOriginShape_[1]), "2"),
616 : return ge::GRAPH_FAILED);
617 : }
618 :
619 : return ge::GRAPH_SUCCESS;
620 : }
621 :
622 : void QuantUpdateScatterRegbaseTiling::PrintDebugInfo()
623 : {
624 : OP_LOGD(
625 : context_->GetNodeName(),
626 : "[QuantUpdateScatter]coreNum: %ld, eachCoreBsNum: %ld, lastCoreBsNum: %ld, srcBsStride: %ld, "
627 : "dstBsStride: %ld, indexElements: %ld, varDim1: %ld, varDim2: %ld, varDim3: %ld, "
628 : "innerLoopEle: %ld, innerLoopTimes: %ld, innerLoopTail: %ld, indicesShapeRank: %ld, quantScalesElements: %ld, "
629 : "quantZeroPointsElements: %ld, innerLoopTimesLastCore: %ld, innerLoopTailLastCore: %ld, "
630 : "innerLoopFullRpt: %ld, innerLoopFullRptLastCore: %ld, innerLoopTailRpt: %ld, innerLoopTailRptLastCore: %ld, "
631 : "srcFirBsStride: %ld, "
632 : "dstFirSecBsStride: %ld, updateDim0: %ld, updateDim1: %ld, updateDim2: %ld, updateDim3: %ld, "
633 : "updateOriLastDim: %ld, updateOriLastDimAlign: %ld",
634 : tilingData_.get_coreNum(), tilingData_.get_eachCoreBsNum(), tilingData_.get_lastCoreBsNum(),
635 : tilingData_.get_srcBsStride(), tilingData_.get_dstBsStride(), tilingData_.get_indexElements(),
636 : tilingData_.get_varDim1(), tilingData_.get_varDim2(), tilingData_.get_varDim3(), tilingData_.get_innerLoopEle(),
637 : tilingData_.get_innerLoopTimes(), tilingData_.get_innerLoopTail(), tilingData_.get_indicesShapeRank(),
638 : tilingData_.get_quantScalesElements(), tilingData_.get_quantZeroPointsElements(),
639 : tilingData_.get_innerLoopTimesLastCore(), tilingData_.get_innerLoopTailLastCore(),
640 : tilingData_.get_innerLoopFullRpt(), tilingData_.get_innerLoopFullRptLastCore(),
641 : tilingData_.get_innerLoopTailRpt(), tilingData_.get_innerLoopTailRptLastCore(),
642 : tilingData_.get_srcFirBsStride(), tilingData_.get_dstFirSecBsStride(), tilingData_.get_updateDim0(),
643 : tilingData_.get_updateDim1(), tilingData_.get_updateDim2(), tilingData_.get_updateDim3(),
644 : tilingData_.get_updateOriLastDim(), tilingData_.get_updateOriLastDimAlign());
645 : OP_LOGD(context_->GetNodeName(),
646 : "[QuantUpdateScatter]tilingKey: %lu, splitMode_: %lu, zeroPointsType_: %lu, divMode_: %lu, roundMode_: %lu",
647 : tilingKey_, splitMode_, zeroPointsType_, divMode_, castRoundMode_);
648 : }
649 :
650 : ge::graphStatus QuantUpdateScatterRegbaseTiling::DoTiling()
651 : {
652 : OP_CHECK_IF(PrepareTilingParams() != ge::GRAPH_SUCCESS,
653 : OP_LOGE(context_->GetNodeName(), "PrepareTilingParams failed!"), return ge::GRAPH_FAILED);
654 : OP_CHECK_IF(VerifyNullTenosr() != ge::GRAPH_SUCCESS, OP_LOGE(context_->GetNodeName(), "VerifyNullTenosr failed!"),
655 : return ge::GRAPH_FAILED);
656 : OP_CHECK_IF(VerifyParamsDtype() != ge::GRAPH_SUCCESS,
657 : OP_LOGE(context_->GetNodeName(), "VerifyParamsDtype return failed."), return ge::GRAPH_FAILED);
658 : OP_CHECK_IF(VerifyTilingQuantParams() != ge::GRAPH_SUCCESS,
659 : OP_LOGE(context_->GetNodeName(), "VerifyTilingQuantParams return failed."), return ge::GRAPH_FAILED);
660 : OP_CHECK_IF(VerifyTilingParams() != ge::GRAPH_SUCCESS,
661 : OP_LOGE(context_->GetNodeName(), "VerifyTilingParams failed!"), return ge::GRAPH_FAILED);
662 : OP_CHECK_IF(MergeDims() != ge::GRAPH_SUCCESS, OP_LOGE(context_->GetNodeName(), "MergeDims failed!"),
663 : return ge::GRAPH_FAILED);
664 : OP_CHECK_IF(GetTilingParam() != ge::GRAPH_SUCCESS, OP_LOGE(context_->GetNodeName(), "GetTilingParam failed!"),
665 : return ge::GRAPH_FAILED);
666 : tilingKey_ = GET_TPL_TILING_KEY(splitMode_, zeroPointsType_, divMode_, castRoundMode_);
667 : PrintDebugInfo();
668 :
669 : auto rawTilingData = context_->GetRawTilingData();
670 : OP_CHECK_NULL_WITH_CONTEXT(context_, rawTilingData);
671 : if (tilingData_.GetDataSize() > rawTilingData->GetCapacity()) {
672 : OP_LOGE_FOR_INVALID_VALUES_WITH_REASON(
673 : context_->GetNodeName(), "TilingDataSize,Capacity",
674 : std::to_string(tilingData_.GetDataSize()) + "," + std::to_string(rawTilingData->GetCapacity()),
675 : "The value of TilingDataSize must be greater than that of Capacity");
676 : return ge::GRAPH_FAILED;
677 : }
678 : tilingData_.SaveToBuffer(context_->GetRawTilingData()->GetData(), context_->GetRawTilingData()->GetCapacity());
679 : context_->GetRawTilingData()->SetDataSize(tilingData_.GetDataSize());
680 : context_->SetBlockDim(tilingData_.get_coreNum());
681 : context_->SetTilingKey(tilingKey_);
682 : size_t* currentWorkspace = context_->GetWorkspaceSizes(1);
683 : OP_CHECK_NULL_WITH_CONTEXT(context_, currentWorkspace);
684 : currentWorkspace[0] = SYNC_WORKSPACE_SIZE;
685 :
686 : return ge::GRAPH_SUCCESS;
687 : }
688 :
689 : static ge::graphStatus Tiling4QuantUpdateScatter(gert::TilingContext* context)
690 : {
691 : auto compileInfo = context->GetCompileInfo<QuantUpdateScatterCompileInfo>();
692 : OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
693 : QuantUpdateScatterRegbaseTiling tiling(context);
694 : return tiling.DoTiling();
695 : }
696 :
697 : static ge::graphStatus TilingPrepare4QuantUpdateScatter([[maybe_unused]] gert::TilingParseContext* context)
698 : {
699 : return ge::GRAPH_SUCCESS;
700 : }
701 :
702 : IMPL_OP_OPTILING(QuantUpdateScatter)
703 : .Tiling(Tiling4QuantUpdateScatter)
704 : .TilingParse<QuantUpdateScatterCompileInfo>(TilingPrepare4QuantUpdateScatter);
705 : } // namespace optiling
|
