Line data Source code
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 softmax_grad_tiling_base.cc
13 : * \brief
14 : */
15 :
16 : #include "softmax_grad_ext_tiling.h"
17 : #include <nlohmann/json.hpp>
18 : #include "atvoss/broadcast/broadcast_tiling.h"
19 :
20 : using namespace AscendC;
21 : using namespace Ops::NN::OpTiling;
22 : using namespace ge;
23 :
24 : namespace optiling {
25 : inline std::unique_ptr<nlohmann::json> GetCompileInfoJson(gert::TilingParseContext* context)
26 : {
27 : auto json_str = context->GetCompiledJson();
28 : OP_CHECK_IF(json_str == nullptr, OP_LOGE(context->GetNodeName(), "json_str is nullptr!"), return nullptr);
29 : std::unique_ptr<nlohmann::json> parsed_object_cinfo = std::make_unique<nlohmann::json>(
30 : nlohmann::json::parse(json_str));
31 : return parsed_object_cinfo;
32 : }
33 :
34 : std::string SoftmaxGradExtTilingBase::VectorToString(const std::vector<int64_t>& s)
35 : {
36 : std::stringstream ss;
37 : for (auto iter = s.begin(); iter != s.end(); ++iter) {
38 : ss << *iter;
39 : if (iter != s.end() - CONST_ONE) {
40 : ss << ", ";
41 : }
42 : }
43 : return ss.str();
44 : }
45 :
46 : std::string SoftmaxGradExtTilingBase::VectorToString(const int64_t* s, int64_t size)
47 : {
48 : std::stringstream ss;
49 : for (int64_t i = 0; i < size; i++) {
50 : ss << s[i];
51 : if (i != size - CONST_ONE) {
52 : ss << ", ";
53 : }
54 : }
55 : return ss.str();
56 : }
57 :
58 : ge::graphStatus SoftmaxGradExtTilingBase::GetAndCheckDtypes()
59 : {
60 : auto attrs = context_->GetAttrs();
61 : OP_CHECK_NULL_WITH_CONTEXT(context_, attrs);
62 :
63 : auto xDesc = context_->GetInputDesc(CONST_ZERO);
64 : OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc);
65 : xDtype_ = xDesc->GetDataType();
66 :
67 : auto xDesc1 = context_->GetInputDesc(CONST_ONE);
68 : OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc1);
69 : ge::DataType xDtype1 = xDesc1->GetDataType();
70 :
71 : auto xDesc2 = context_->GetInputDesc(CONST_TWO);
72 : OP_CHECK_NULL_WITH_CONTEXT(context_, xDesc2);
73 : ge::DataType xDtype2 = xDesc2->GetDataType();
74 :
75 : auto yDesc = context_->GetOutputDesc(CONST_ZERO);
76 : OP_CHECK_NULL_WITH_CONTEXT(context_, yDesc);
77 : yDtype_ = yDesc->GetDataType();
78 :
79 8 : OP_TILING_CHECK(xDtype_ != yDtype_ || xDtype_ != xDtype1 || xDtype_ != xDtype2,
80 : VECTOR_INNER_ERR_REPORT_TILIING(
81 : context_->GetNodeName(),
82 : "Input0 dtype [%s], Input1 dtype [%s], Input2 dtype [%s] and Output dtype [%s] should be same.",
83 : ge::TypeUtils::DataTypeToSerialString(xDtype_).c_str(),
84 : ge::TypeUtils::DataTypeToSerialString(xDtype1).c_str(),
85 : ge::TypeUtils::DataTypeToSerialString(xDtype2).c_str(),
86 : ge::TypeUtils::DataTypeToSerialString(yDtype_).c_str()),
87 : return ge::GRAPH_FAILED);
88 : OP_TILING_CHECK(xDtype_ != ge::DT_FLOAT16 && xDtype_ != ge::DT_FLOAT && xDtype_ != ge::DT_BF16,
89 : VECTOR_INNER_ERR_REPORT_TILIING(
90 : context_->GetNodeName(),
91 : "Input dtype is [%s], only support dtype ge::DT_FLOAT16, ge::DT_FLOAT or ge::DT_BF16.",
92 : ge::TypeUtils::DataTypeToSerialString(xDtype_).c_str()),
93 : return ge::GRAPH_FAILED);
94 :
95 : if (xDtype_ == ge::DT_FLOAT) {
96 : xDtypeSize_ = FLOAT32_BYTES;
97 : } else if (xDtype_ == ge::DT_FLOAT16 || xDtype_ == ge::DT_BF16) {
98 : xDtypeSize_ = FLOAT16_BYTES;
99 : }
100 :
101 : yDtypeSize_ = xDtypeSize_;
102 :
103 : return ge::GRAPH_SUCCESS;
104 : }
105 :
106 : ge::graphStatus SoftmaxGradExtTilingBase::GetDimsAndCheckShapeValid()
107 : {
108 : auto xShape = context_->GetInputShape(CONST_ZERO);
109 : OP_CHECK_NULL_WITH_CONTEXT(context_, xShape);
110 : auto xStorageShape = Ops::Base::EnsureNotScalar(xShape->GetStorageShape());
111 : xShapeSize_ = xStorageShape.GetDimNum();
112 :
113 : auto xShape1 = context_->GetInputShape(CONST_ONE);
114 : OP_CHECK_NULL_WITH_CONTEXT(context_, xShape1);
115 : auto xStorageShape1 = Ops::Base::EnsureNotScalar(xShape1->GetStorageShape());
116 :
117 : auto xShape2 = context_->GetInputShape(CONST_TWO);
118 : OP_CHECK_NULL_WITH_CONTEXT(context_, xShape2);
119 : auto xStorageShape2 = Ops::Base::EnsureNotScalar(xShape2->GetStorageShape());
120 : xShapeSize2_ = xStorageShape2.GetDimNum();
121 :
122 : auto yShape = context_->GetOutputShape(CONST_ZERO);
123 : OP_CHECK_NULL_WITH_CONTEXT(context_, yShape);
124 : auto yStorageShape = Ops::Base::EnsureNotScalar(yShape->GetStorageShape());
125 :
126 : OP_TILING_CHECK(
127 : xShapeSize_ > MAX_DIMS, // 超过支持的最大维度数
128 : VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(), "Input dim size [%ld] is larger than 6.", xShapeSize_),
129 : return ge::GRAPH_FAILED);
130 : OP_TILING_CHECK(
131 : xShapeSize_ == CONST_ZERO, // 检查输入的维度是否为0
132 : VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(), "Input dim size is zero, not support empty tensor."),
133 : return ge::GRAPH_FAILED);
134 :
135 : xShape_.resize(xShapeSize_); // 将xshape_大小调整为xShapeSize
136 : if (xShapeSize2_ == 1 && xStorageShape2.GetDim(0) == 1) {
137 : for (int i = 0; i < xShapeSize_; i++) {
138 : OP_TILING_CHECK(xStorageShape.GetDim(i) != yStorageShape.GetDim(i) ||
139 : xStorageShape.GetDim(i) != xStorageShape1.GetDim(i),
140 : VECTOR_INNER_ERR_REPORT_TILIING(
141 : context_->GetNodeName(),
142 : "Input0 dim[%d]: %ld, Input1 dim[%d]: %ld and Output dim[%d]: %ld should be "
143 : "same. Input2 dim[0]: %ld. ",
144 : i, xStorageShape.GetDim(i), i, xStorageShape1.GetDim(i), i, yStorageShape.GetDim(i),
145 : xStorageShape2.GetDim(0)),
146 : return ge::GRAPH_FAILED);
147 : OP_TILING_CHECK(xStorageShape.GetDim(i) <= CONST_ZERO,
148 : VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(), "Not support input dim[%d]: %ld.",
149 : i, xStorageShape.GetDim(i)),
150 : return ge::GRAPH_FAILED);
151 : xShape_[i] = xStorageShape.GetDim(i);
152 : }
153 : } else {
154 : isX2Scalar_ = false;
155 : for (int i = 0; i < xShapeSize_; i++) {
156 : OP_TILING_CHECK(
157 : xStorageShape.GetDim(i) != yStorageShape.GetDim(i) ||
158 : xStorageShape.GetDim(i) != xStorageShape1.GetDim(i) ||
159 : xStorageShape.GetDim(i) != xStorageShape2.GetDim(i),
160 : VECTOR_INNER_ERR_REPORT_TILIING(
161 : context_->GetNodeName(),
162 : "Input0 dim[%d]: %ld, Input1 dim[%d]: %ld, Input2 dim[%d]: %ld and Output dim[%d]: %ld should be "
163 : "same.",
164 : i, xStorageShape.GetDim(i), i, xStorageShape1.GetDim(i), i, xStorageShape2.GetDim(i), i,
165 : yStorageShape.GetDim(i)),
166 : return ge::GRAPH_FAILED);
167 : OP_TILING_CHECK(xStorageShape.GetDim(i) <= CONST_ZERO,
168 : VECTOR_INNER_ERR_REPORT_TILIING(context_->GetNodeName(), "Not support input dim[%d]: %ld.",
169 : i, xStorageShape.GetDim(i)),
170 : return ge::GRAPH_FAILED);
171 : xShape_[i] = xStorageShape.GetDim(i);
172 : }
173 : }
174 : return ge::GRAPH_SUCCESS;
175 : }
176 :
177 : ge::graphStatus SoftmaxGradExtTilingBase::GetAndCheckAxes()
178 : {
179 : auto attrs = context_->GetAttrs();
180 : OP_CHECK_NULL_WITH_CONTEXT(context_, attrs);
181 : auto attrAxis = attrs->GetAttrPointer<int64_t>(CONST_ZERO); // 默认-1轴reduce
182 : OP_CHECK_NULL_WITH_CONTEXT(context_, attrAxis);
183 : reduceAxes_ = xShapeSize_ - CONST_ONE;
184 : OP_TILING_CHECK((*attrAxis != -1 && *attrAxis != reduceAxes_),
185 : VECTOR_INNER_ERR_REPORT_TILIING(
186 : context_->GetNodeName(), "Axes is: %ld, axes only support -1 or %ld", *attrAxis, reduceAxes_),
187 : return ge::GRAPH_FAILED);
188 : return ge::GRAPH_SUCCESS;
189 : }
190 :
191 : ge::graphStatus SoftmaxGradExtTilingBase::GetShapeAttrsInfo()
192 : {
193 : OP_TILING_CHECK(context_ == nullptr, OP_LOGE("SoftmaxGradExtTilingBase", "context is nullptr."),
194 : return ge::GRAPH_FAILED);
195 :
196 : OP_TILING_CHECK(GetAndCheckDtypes() != ge::GRAPH_SUCCESS, , return ge::GRAPH_FAILED);
197 : OP_TILING_CHECK(GetDimsAndCheckShapeValid() != ge::GRAPH_SUCCESS, , return ge::GRAPH_FAILED);
198 : OP_TILING_CHECK(GetAndCheckAxes() != ge::GRAPH_SUCCESS, , return ge::GRAPH_FAILED);
199 :
200 : // 合轴(a1_, r_, a0_)
201 : a1_ = DIM_NUM_ONE;
202 : r_ = xShape_[reduceAxes_];
203 : a0_ = DIM_NUM_ONE;
204 : for (int i = 0; i < xShapeSize_; i++) {
205 : if (i < reduceAxes_) {
206 : a1_ *= xShape_[i];
207 : } else if (i > reduceAxes_) {
208 : a0_ *= xShape_[i];
209 : }
210 : }
211 :
212 : OP_LOGD(context_->GetNodeName(), "inputs original shape is:(%s), axes is:%ld, fused shape is: (%ld, %ld, %ld)\n",
213 : VectorToString(xShape_).c_str(), reduceAxes_, a1_, r_, a0_);
214 :
215 : return ge::GRAPH_SUCCESS;
216 : }
217 :
218 : ge::graphStatus SoftmaxGradExtTilingBase::GetPlatformInfo()
219 : {
220 : auto compileInfo = reinterpret_cast<const SoftmaxGradExtCompileInfo*>(context_->GetCompileInfo());
221 : OP_CHECK_NULL_WITH_CONTEXT(context_, compileInfo);
222 : blockSize_ = static_cast<uint64_t>(compileInfo->blockSize);
223 : vlFp32_ = static_cast<uint64_t>(compileInfo->vlFp32);
224 : vlFp16_ = static_cast<uint64_t>(compileInfo->vlFp16);
225 :
226 : OP_LOGD(context_->GetNodeName(), "blockSize: %ld, vlFp32: %ld, vlFp16: %ld.", blockSize_, vlFp32_, vlFp16_);
227 :
228 : auto platformInfoPtr = context_->GetPlatformInfo();
229 : if (platformInfoPtr == nullptr) {
230 : OP_LOGD(context_->GetNodeName(), "Entering into get core num from compile info."); // 说明是离线编译
231 : aicoreParams_.blockDim = static_cast<int32_t>(compileInfo->coreNum);
232 : aicoreParams_.ubSize = static_cast<int64_t>(compileInfo->ubSize);
233 : } else {
234 : OP_LOGD(context_->GetNodeName(), "Entering into get core num from platform."); // 说明是在线运行环境
235 : auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr);
236 : aicoreParams_.blockDim = static_cast<int64_t>(ascendcPlatform.GetCoreNumAiv());
237 : uint64_t ubSizeTemp = CONST_ZERO;
238 : ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizeTemp);
239 : aicoreParams_.ubSize = static_cast<int64_t>(ubSizeTemp);
240 : }
241 : return ge::GRAPH_SUCCESS;
242 : }
243 :
244 : static ge::graphStatus SoftmaxGradExtDSLTiling(gert::TilingContext* context,
245 : const SoftmaxGradExtCompileInfo* compileInfo)
246 : {
247 : // get input
248 : uint64_t inputNums = context->GetComputeNodeInputNum(); // 获取输入数量
249 : OP_TILING_CHECK(
250 : (inputNums < CONST_ONE),
251 : VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "inputNums is: %ld, less than one", inputNums),
252 : return ge::GRAPH_FAILED);
253 : std::vector<gert::Shape> inputShapes(inputNums);
254 : for (uint64_t i = 0; i < inputNums; i++) {
255 : auto inputShape = context->GetInputShape(i);
256 : OP_CHECK_NULL_WITH_CONTEXT(context, inputShape);
257 : inputShapes[i] = Ops::Base::EnsureNotScalar(inputShape->GetStorageShape());
258 : }
259 : auto desc = context->GetInputDesc(CONST_ZERO);
260 : OP_CHECK_NULL_WITH_CONTEXT(context, desc);
261 : DataType dataType = desc->GetDataType();
262 :
263 : // get attr for reduce axis
264 : auto attrs = context->GetAttrs();
265 : OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
266 : std::vector<int64_t> axis;
267 : int64_t reduceAxis = CONST_ZERO;
268 : auto axisListPtr = attrs->GetListInt(CONST_ZERO);
269 : int64_t xShapeSize = inputShapes[CONST_ZERO].GetDimNum(); // 获取输入张量的维度数
270 : if (axisListPtr == nullptr || axisListPtr->GetSize() == CONST_ZERO) {
271 : reduceAxis = xShapeSize - CONST_ONE;
272 : axis.emplace_back(reduceAxis);
273 : } else {
274 : for (size_t i = 0; i < axisListPtr->GetSize(); i++) {
275 : reduceAxis = axisListPtr->GetData()[i];
276 : OP_TILING_CHECK(
277 : (reduceAxis < -xShapeSize || reduceAxis > xShapeSize - CONST_ONE),
278 : VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "Dimension is: %ld, out of range [-%ld, %ld]",
279 : reduceAxis, xShapeSize, xShapeSize - CONST_ONE),
280 : return ge::GRAPH_FAILED);
281 :
282 : reduceAxis = reduceAxis < CONST_ZERO ? reduceAxis + xShapeSize : reduceAxis;
283 : if (std::find(axis.begin(), axis.end(), reduceAxis) == axis.end()) {
284 : axis.emplace_back(reduceAxis);
285 : }
286 : }
287 : }
288 : axis.emplace_back(reduceAxis);
289 : return ge::GRAPH_SUCCESS;
290 : }
291 :
292 : ge::graphStatus TilingPrepareForSoftmaxGradExtAscendC(gert::TilingParseContext* context)
293 : {
294 : OP_LOGD(context->GetNodeName(), "TilingPrepareForSoftmaxGradExtAscendC enter.");
295 :
296 : auto compileInfoPtr = context->GetCompiledInfo<SoftmaxGradExtCompileInfo>(); // 创建指针获取编译信息
297 : OP_CHECK_NULL_WITH_CONTEXT(context, compileInfoPtr);
298 :
299 : compileInfoPtr->blockSize = Ops::Base::GetUbBlockSize(
300 : context); // 设置blocksize为平台提供的通用缓冲区对齐大小=32U(4字节)
301 : compileInfoPtr->vlFp32 = Ops::Base::GetVRegSize(context) /
302 : FLOAT32_BYTES; // 设置vlfp32为平台提供的浮点数寄存器大小（以float32单位计算）=64
303 : compileInfoPtr->vlFp16 = Ops::Base::GetVRegSize(context) /
304 : FLOAT16_BYTES; // 设置vlfp16为平台提供的浮点数寄存器大小（以float16单位计算）=128
305 :
306 : fe::PlatFormInfos* platformInfoPtr = context->GetPlatformInfo();
307 : OP_LOGE_IF(platformInfoPtr == nullptr, ge::GRAPH_FAILED, context->GetNodeName(), "platformInfoPtr is null");
308 :
309 : auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfoPtr); // 创建platformAscendC对象,获取硬件信息
310 : compileInfoPtr->coreNum = ascendcPlatform.GetCoreNumAiv();
311 : OP_TILING_CHECK((compileInfoPtr->coreNum <= CONST_ZERO),
312 : VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "Get core num failed, core num: %u",
313 : static_cast<uint32_t>(compileInfoPtr->coreNum)),
314 : return ge::GRAPH_FAILED);
315 : uint64_t ubSizeTemp = CONST_ZERO; // 初始化临时变量
316 : ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSizeTemp); // 获取每个核心的通用缓冲区大小
317 : compileInfoPtr->ubSize = static_cast<int64_t>(ubSizeTemp);
318 : OP_TILING_CHECK((compileInfoPtr->ubSize <= CONST_ZERO),
319 : VECTOR_INNER_ERR_REPORT_TILIING(context->GetNodeName(), "Get ub size failed, ub size: %uB",
320 : static_cast<uint32_t>(compileInfoPtr->ubSize)),
321 : return ge::GRAPH_FAILED);
322 :
323 : return ge::GRAPH_SUCCESS;
324 : }
325 :
326 : ge::graphStatus TilingPrepareForSoftmaxGradExtDsl(gert::TilingParseContext* context)
327 : {
328 : auto compileInfoPtr = context->GetCompiledInfo<SoftmaxGradExtCompileInfo>();
329 : OP_LOGE_IF(compileInfoPtr == nullptr, ge::GRAPH_FAILED, context->GetNodeName(), "compileInfoPtr is null");
330 : std::unique_ptr<nlohmann::json> parsedObjectCInfo = GetCompileInfoJson(context); // 获取json格式的编译信息
331 : OP_CHECK_NULL_WITH_CONTEXT(context, parsedObjectCInfo);
332 : return ge::GRAPH_SUCCESS;
333 : }
334 :
335 : ge::graphStatus TilingForSoftmaxGradExt(gert::TilingContext* context)
336 : {
337 : if (context == nullptr) {
338 : OP_LOGE("SoftmaxGradExtTilingBase", "Tiling context is nullptr");
339 : return ge::GRAPH_FAILED;
340 : }
341 : OP_LOGD(context->GetNodeName(), "TilingForSoftmaxGradExt enter");
342 : auto compileInfo = reinterpret_cast<const SoftmaxGradExtCompileInfo*>(context->GetCompileInfo());
343 : OP_CHECK_NULL_WITH_CONTEXT(context, compileInfo);
344 : if (compileInfo->isAscendC) {
345 : OP_LOGD(context->GetNodeName(), "SoftmaxGradExtTilingBase Ascendc enter");
346 : return TilingRegistry::GetInstance().DoTilingImpl(context);
347 : }
348 : OP_LOGD(context->GetNodeName(), "SoftmaxGradExtDSLTiling enter");
349 : return SoftmaxGradExtDSLTiling(context, compileInfo);
350 : }
351 :
352 : ge::graphStatus TilingPrepareForSoftmaxGradExt(gert::TilingParseContext* context)
353 : {
354 : if (context == nullptr) {
355 : OP_LOGE("TilingPrepareForSoftmaxGradExt", "Tiling context is nullptr");
356 : return ge::GRAPH_FAILED;
357 : }
358 : OP_LOGD(context->GetNodeName(), "TilingPrepareForSoftmaxGradExt enter.");
359 :
360 : auto compileInfoPtr = context->GetCompiledInfo<SoftmaxGradExtCompileInfo>();
361 : OP_CHECK_IF((compileInfoPtr == nullptr), OP_LOGE(context->GetNodeName(), "compileInfoPtr is null"),
362 : return ge::GRAPH_FAILED);
363 : compileInfoPtr->isAscendC = IsRegbaseSocVersion(context);
364 : if (compileInfoPtr->isAscendC) {
365 : OP_LOGD(context, "TilingPrepareForSoftmaxGradExtAscendC enter");
366 : return TilingPrepareForSoftmaxGradExtAscendC(context);
367 : }
368 : OP_LOGD(context, "TilingPrepareForSoftmaxGradExtDsl enter");
369 : return TilingPrepareForSoftmaxGradExtDsl(context);
370 : }
371 :
372 : IMPL_OP_OPTILING(SoftmaxGradExt)
373 : .Tiling(TilingForSoftmaxGradExt)
374 : .TilingParse<SoftmaxGradExtCompileInfo>(TilingPrepareForSoftmaxGradExt);
375 :
376 : } // namespace optiling
|
