Star1003
1932
代码
代码
Issues249
Pull Requests361
Actions
项目讨论
Wiki
分析
Star1003
1932
已合并
Fixed the bug that the padIdx is used for calculation but the none value is set in the Interface #1100
fengjiawei1创建于 1月28日
Fixed the bug that the padIdx is used for calculation but the none value is set in the Interface #1100
已合并
fengjiawei1创建于 1月28日
从fengjiawei1/ops-nn_5365embedding_bag_ca合入到cann/ops-nnmaster
共 12 个文件变更+294-222
Mindex/embedding_bag/op_host/config/ascend950/embedding_bag_binary.json+2-2
@@ -1142,7 +1142,7 @@
1142 ]1142 ]
1143 },1143 },
1144 {1144 {
1145- "bin_filename": "EmbeddingBag_fp16_case3",1145+ "bin_filename": "EmbeddingBag_fp16_case4",
1146 "inputs": [1146 "inputs": [
1147 {1147 {
1148 "name": "weight",1148 "name": "weight",
@@ -1256,7 +1256,7 @@
1256 ]1256 ]
1257 },1257 },
1258 {1258 {
1259- "bin_filename": "EmbeddingBag_fp32_case3",1259+ "bin_filename": "EmbeddingBag_fp32_case4",
1260 "inputs": [1260 "inputs": [
1261 {1261 {
1262 "name": "weight",1262 "name": "weight",
Mindex/embedding_bag/op_host/embedding_bag_regbase_tiling.cpp+9-5
@@ -144,9 +144,13 @@ ge::graphStatus EmbeddingBagRegBaseTiling::GetShapeAttrsInfo()
144 isNeedSampleWeight_ = 0;144 isNeedSampleWeight_ = 0;
145 }145 }
146 146
147- inclueLastOfst_ = *(attrs->GetAttrPointer<bool>)(ATTR_INCLUDE_LAST_OFFSET);147+ auto inclueLastOfstPtr = attrs->GetAttrPointer<bool>(ATTR_INCLUDE_LAST_OFFSET);
148- paddingIdx_ = *(attrs->GetAttrPointer<int64_t>)(ATTR_PADD_INDEX);148+ OP_CHECK_NULL_WITH_CONTEXT(context_, inclueLastOfstPtr);
149- paddingIdx_ = paddingIdx_ < 0 ? paddingIdx_ + numEmbeddings_ : paddingIdx_;149+ inclueLastOfst_ = *inclueLastOfstPtr;
150+ 
151+ auto paddingIdxPtr = attrs->GetAttrPointer<bool>(ATTR_INCLUDE_LAST_OFFSET);
152+ OP_CHECK_NULL_WITH_CONTEXT(context_, paddingIdxPtr);
153+ paddingIdx_ = *paddingIdxPtr;
150 154 
151 if (embeddingDim_ * weightTypeSize_ <= MAX_SIMT_EMBDDING_BYTES) {155 if (embeddingDim_ * weightTypeSize_ <= MAX_SIMT_EMBDDING_BYTES) {
152 usedCoreNum_ = totalCoreNum_;156 usedCoreNum_ = totalCoreNum_;
@@ -278,7 +282,7 @@ int64_t EmbeddingBagRegBaseTiling::GetWeightAlignSize1D(int64_t weightRowFactor,
278 282 
279void EmbeddingBagRegBaseTiling::Compute1DFactor()283void EmbeddingBagRegBaseTiling::Compute1DFactor()
280{284{
281- auto halfUbSize = ubSize_ / DOUBLE_BUF - UB_RESERVED_BUFF;285+ auto halfUbSize = ubSize_ - UB_RESERVED_BUFF;
282 int64_t ubBlock = static_cast<int64_t>(Ops::Base::GetUbBlockSize(context_));286 int64_t ubBlock = static_cast<int64_t>(Ops::Base::GetUbBlockSize(context_));
283 287 
284 /* indices 分配4K */288 /* indices 分配4K */
@@ -446,7 +450,7 @@ void EmbeddingBagRegBaseTiling::SetTilingData()
446 tilingData_.set_indicesNumel(indicesNumel_);450 tilingData_.set_indicesNumel(indicesNumel_);
447 tilingData_.set_indicesLimit(indicesLimit_);451 tilingData_.set_indicesLimit(indicesLimit_);
448 tilingData_.set_sampleWeightNum(sampleWeightNum_);452 tilingData_.set_sampleWeightNum(sampleWeightNum_);
449- 453+ tilingData_.set_inclueLastOfst(static_cast<int64_t>(inclueLastOfst_));
CANN-robot
CANN-robotCANN-robot1月28日

错误处理与异常安全: 第450行新增了 tilingData_.set_inclueLastOfst(static_cast<int64_t>(inclueLastOfst_));。inclueLastOfst_ 是布尔类型（从第148行 *(attrs->GetAttrPointer<bool>)(ATTR_INCLUDE_LAST_OFFSET) 获取），将其强制转换为 int64_t 是合理的。但需要确保 inclueLastOfst_ 已被正确初始化。在第148行，inclueLastOfst_ 直接从属性指针赋值，如果属性不存在或指针为空，可能会导致未定义行为。虽然第139-140行对 attrs 进行了判空，但第148行直接解引用 attrs->GetAttrPointer<bool> 返回的指针，未检查该指针是否为空。

问题类型: 错误处理与异常安全 文件路径: index/embedding_bag/op_host/embedding_bag_regbase_tiling.cpp 行号: 450 问题代码:

tilingData_.set_inclueLastOfst(static_cast<int64_t>(inclueLastOfst_));

修改建议:

建议在第148行解引用属性指针之前，添加对指针有效性的检查。例如：`auto attrPtr = attrs->GetAttrPointer<bool>(ATTR_INCLUDE_LAST_OFFSET); OP_CHECK_NULL_WITH_CONTEXT(context_, attrPtr); inclueLastOfst_ = *attrPtr;`。这样可以避免潜在的空指针解引用风险。

此评论由代码审查工具自动生成

likedislike
450 TilingDataPrint();454 TilingDataPrint();
451}455}
452 456 
Mindex/embedding_bag/op_host/embedding_bag_tiling.cpp+2-17
@@ -18,6 +18,7 @@
18#include "log/log.h"18#include "log/log.h"
19#include "platform/platform_info.h"19#include "platform/platform_info.h"
20#include "platform/platform_infos_def.h"20#include "platform/platform_infos_def.h"
21+#include "tiling_base/tiling_util.h"
21#include "embedding_bag_tiling.h"22#include "embedding_bag_tiling.h"
22#include "embedding_bag_regbase_tiling.h"23#include "embedding_bag_regbase_tiling.h"
23#include "tiling/platform/platform_ascendc.h"24#include "tiling/platform/platform_ascendc.h"
@@ -248,25 +249,9 @@ void EmbeddingBagTiling::TilingDataPrint() const
248 OP_LOGD(tilingContext_, "tilingKey_: %ld", tilingKey_);249 OP_LOGD(tilingContext_, "tilingKey_: %ld", tilingKey_);
249}250}
250 251 
251-static bool IsRegbaseSocVersion4EmBeddingBag(platform_ascendc::SocVersion version)
252-{
253- const static std::set<platform_ascendc::SocVersion> regbaseSocVersions = {
254- platform_ascendc::SocVersion::ASCEND950
255- };
256- 
257- return regbaseSocVersions.find(version) != regbaseSocVersions.end();
258-}
259- 
260-bool IsRegbaseSocVersion4EmBeddingBag(const gert::TilingContext* context)
261-{
262- auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
263- auto socVersion = ascendcPlatform.GetSocVersion();
264- return IsRegbaseSocVersion4EmBeddingBag(socVersion);
265-}
266- 
267ge::graphStatus TilingEmbeddingBag(gert::TilingContext* context)252ge::graphStatus TilingEmbeddingBag(gert::TilingContext* context)
268{253{
269- if (IsRegbaseSocVersion4EmBeddingBag(context)) {254+ if (Ops::NN::OpTiling::IsRegbaseSocVersion(context)) {
270 return EmbeddingBagTilingForRegBase(context);255 return EmbeddingBagTilingForRegBase(context);
271 }256 }
272 EmbeddingBagTiling tilingObject(context);257 EmbeddingBagTiling tilingObject(context);
Mindex/embedding_bag/op_host/embedding_bag_tiling.h+1-1
@@ -7,7 +7,6 @@
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- 
11/*!10/*!
12 * \file embedding_bag_tiling.h11 * \file embedding_bag_tiling.h
13 * \brief12 * \brief
@@ -47,6 +46,7 @@ TILING_DATA_FIELD_DEF(int64_t, isNeedSampleWeight);
47TILING_DATA_FIELD_DEF(int64_t, indicesNumel);46TILING_DATA_FIELD_DEF(int64_t, indicesNumel);
48TILING_DATA_FIELD_DEF(int64_t, indicesLimit);47TILING_DATA_FIELD_DEF(int64_t, indicesLimit);
49TILING_DATA_FIELD_DEF(int64_t, sampleWeightNum);48TILING_DATA_FIELD_DEF(int64_t, sampleWeightNum);
49+TILING_DATA_FIELD_DEF(int64_t, inclueLastOfst);
50 50 
51END_TILING_DATA_DEF;51END_TILING_DATA_DEF;
52REGISTER_TILING_DATA_CLASS(EmbeddingBag, EmbeddingBagTilingData)52REGISTER_TILING_DATA_CLASS(EmbeddingBag, EmbeddingBagTilingData)
Mindex/embedding_bag/op_host/op_api/aclnn_embedding_bag.cpp+69-48
@@ -27,6 +27,8 @@
27#include "opdev/op_log.h"27#include "opdev/op_log.h"
28#include "opdev/shape_utils.h"28#include "opdev/shape_utils.h"
29#include "opdev/tensor_view_utils.h"29#include "opdev/tensor_view_utils.h"
30+#include "opdev/platform.h"
31+#include "op_api/aclnn_util.h"
30 32 
31using namespace op;33using namespace op;
32#ifdef __cplusplus34#ifdef __cplusplus
@@ -46,8 +48,9 @@ static const std::initializer_list<op::DataType> WEIGHT_DTYPE_SUPPORT_LIST_910 =
46 48 
47static inline const std::initializer_list<DataType>& GetDtypeSupportList()49static inline const std::initializer_list<DataType>& GetDtypeSupportList()
48{50{
49- if (GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND310P ||51+ auto curArch = GetCurrentPlatformInfo().GetCurNpuArch();
50- GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910) {52+ if (curArch == NpuArch::DAV_2002 ||
53+ curArch == NpuArch::DAV_1001) {
51 return WEIGHT_DTYPE_SUPPORT_LIST_910;54 return WEIGHT_DTYPE_SUPPORT_LIST_910;
52 } else {55 } else {
53 return WEIGHT_DTYPE_SUPPORT_LIST_910B;56 return WEIGHT_DTYPE_SUPPORT_LIST_910B;
@@ -145,7 +148,11 @@ static bool CheckDims(
145 }148 }
146 149 
147 if (indices != nullptr) {150 if (indices != nullptr) {
148- OP_CHECK_MAX_DIM(indices, MIN_SUPPORT_DIM, return false);151+ if (Ops::NN::AclnnUtil::IsRegbase()) {
152+ OP_CHECK_MAX_DIM(indices, MAX_SUPPORT_DIM, return false);
153+ } else {
154+ OP_CHECK_MAX_DIM(indices, MIN_SUPPORT_DIM, return false);
155+ }
149 }156 }
150 157 
151 if (offsets != nullptr) {158 if (offsets != nullptr) {
@@ -153,8 +160,13 @@ static bool CheckDims(
153 }160 }
154 161 
155 if (perSampleWeights != nullptr) {162 if (perSampleWeights != nullptr) {
156- OP_CHECK_MAX_DIM(perSampleWeights, MIN_SUPPORT_DIM, return false);163+ if (Ops::NN::AclnnUtil::IsRegbase()){
157- OP_CHECK_MIN_DIM(perSampleWeights, MIN_SUPPORT_DIM, return false);164+ OP_CHECK_MAX_DIM(perSampleWeights, MAX_SUPPORT_DIM, return false);
165+ OP_CHECK_MIN_DIM(perSampleWeights, MIN_SUPPORT_DIM, return false);
166+ } else {
167+ OP_CHECK_MAX_DIM(perSampleWeights, MIN_SUPPORT_DIM, return false);
168+ OP_CHECK_MIN_DIM(perSampleWeights, MIN_SUPPORT_DIM, return false);
169+ }
158 }170 }
159 171 
160 if (output != nullptr) {172 if (output != nullptr) {
@@ -171,11 +183,15 @@ static bool CheckDims(
171 }183 }
172 184 
173 if (maxIndices != nullptr) {185 if (maxIndices != nullptr) {
174- if (modeStr == "max") {186+ if (Ops::NN::AclnnUtil::IsRegbase()){
175 OP_CHECK_MAX_DIM(maxIndices, MAX_SUPPORT_DIM, return false);187 OP_CHECK_MAX_DIM(maxIndices, MAX_SUPPORT_DIM, return false);
176- OP_CHECK_MIN_DIM(maxIndices, MAX_SUPPORT_DIM, return false);
177 } else {188 } else {
178- OP_CHECK_MAX_DIM(maxIndices, MIN_SUPPORT_DIM, return false);189+ if (modeStr == "max") {
190+ OP_CHECK_MAX_DIM(maxIndices, MAX_SUPPORT_DIM, return false);
191+ OP_CHECK_MIN_DIM(maxIndices, MAX_SUPPORT_DIM, return false);
192+ } else {
193+ OP_CHECK_MAX_DIM(maxIndices, MIN_SUPPORT_DIM, return false);
194+ }
179 }195 }
180 }196 }
181 return true;197 return true;
@@ -198,12 +214,14 @@ static bool CheckShape(
198 const aclTensor* weight, const aclTensor* indices, const aclTensor* offsets, const aclTensor* perSampleWeights,214 const aclTensor* weight, const aclTensor* indices, const aclTensor* offsets, const aclTensor* perSampleWeights,
199 const std::string& modeStr, bool includeLastOffset, const aclTensor* output, const aclTensor* offset2bag,215 const std::string& modeStr, bool includeLastOffset, const aclTensor* output, const aclTensor* offset2bag,
200 const aclTensor* bagSize, const aclTensor* maxIndices)216 const aclTensor* bagSize, const aclTensor* maxIndices)
201-{217+{
202- if (modeStr != "sum" && perSampleWeights != nullptr) {218+ if (!Ops::NN::AclnnUtil::IsRegbase()) {
203- OP_LOGE(ACLNN_ERR_PARAM_INVALID, "per_sample_weights only supported with mode='sum'");219+ if (modeStr != "sum" && perSampleWeights != nullptr) {
204- return false;220+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "per_sample_weights only supported with mode='sum'");
221+ return false;
222+ }
205 }223 }
206- 224+
207 if (perSampleWeights != nullptr &&225 if (perSampleWeights != nullptr &&
208 indices->GetViewShape().GetShapeSize() != perSampleWeights->GetViewShape().GetShapeSize()) {226 indices->GetViewShape().GetShapeSize() != perSampleWeights->GetViewShape().GetShapeSize()) {
209 OP_LOGE(227 OP_LOGE(
@@ -217,50 +235,51 @@ static bool CheckShape(
217 OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputShape, return false);235 OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(output, outputShape, return false);
218 236 
219 if (offset2bag->GetViewShape().GetShapeSize() != 0 &&237 if (offset2bag->GetViewShape().GetShapeSize() != 0 &&
220- offset2bag->GetViewShape().GetShapeSize() != indices->GetViewShape().GetDim(0)) {238+ offset2bag->GetViewShape().GetShapeSize() != indices->GetViewShape().GetShapeSize()) {
221 OP_LOGE(239 OP_LOGE(
222 ACLNN_ERR_PARAM_INVALID, "offset2bag shape size should be %ld,but got %ld.",240 ACLNN_ERR_PARAM_INVALID, "offset2bag shape size should be %ld,but got %ld.",
223- indices->GetViewShape().GetDim(0), offset2bag->GetViewShape().GetShapeSize());241+ indices->GetViewShape().GetShapeSize(), offset2bag->GetViewShape().GetShapeSize());
224 return false;242 return false;
225- }243+ }
226 244 
227- if (includeLastOffset) {245+ if (Ops::NN::AclnnUtil::IsRegbase()){
228- if (bagSize->GetViewShape().GetShapeSize() != offsets->GetViewShape().GetDim(0) - 1) {
229- OP_LOGE(
230- ACLNN_ERR_PARAM_INVALID, "bagSize shape size should be %ld,but got %ld.",
231- offsets->GetViewShape().GetShapeSize() - 1, bagSize->GetViewShape().GetShapeSize());
232- return false;
233- }
234- } else {
235 if (bagSize->GetViewShape().GetShapeSize() != offsets->GetViewShape().GetDim(0)) {246 if (bagSize->GetViewShape().GetShapeSize() != offsets->GetViewShape().GetDim(0)) {
236 OP_LOGE(247 OP_LOGE(
237- ACLNN_ERR_PARAM_INVALID, "bagSize shape size should be %ld, but got %ld.",248+ ACLNN_ERR_PARAM_INVALID, "bagSize shape size should be %ld,but got %ld.",
238 offsets->GetViewShape().GetShapeSize(), bagSize->GetViewShape().GetShapeSize());249 offsets->GetViewShape().GetShapeSize(), bagSize->GetViewShape().GetShapeSize());
239 return false;250 return false;
251+ }
252+ } else {
253+ if (includeLastOffset) {
254+ if (bagSize->GetViewShape().GetShapeSize() != offsets->GetViewShape().GetDim(0) - 1) {
255+ OP_LOGE(
256+ ACLNN_ERR_PARAM_INVALID, "bagSize shape size should be %ld,but got %ld.",
257+ offsets->GetViewShape().GetShapeSize() - 1, bagSize->GetViewShape().GetShapeSize());
258+ return false;
259+ }
260+ } else {
261+ if (bagSize->GetViewShape().GetShapeSize() != offsets->GetViewShape().GetDim(0)) {
262+ OP_LOGE(
263+ ACLNN_ERR_PARAM_INVALID, "bagSize shape size should be %ld, but got %ld.",
264+ offsets->GetViewShape().GetShapeSize(), bagSize->GetViewShape().GetShapeSize());
265+ return false;
266+ }
240 }267 }
241 }268 }
242- 269+
243 if (modeStr == "max") {270 if (modeStr == "max") {
244 auto maxIndicesShape = GetOutPutShape(weight, offsets, includeLastOffset);271 auto maxIndicesShape = GetOutPutShape(weight, offsets, includeLastOffset);
245 OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(maxIndices, maxIndicesShape, return false);272 OP_CHECK_SHAPE_NOT_EQUAL_WITH_EXPECTED_SIZE(maxIndices, maxIndicesShape, return false);
246 } else {273 } else {
247- if (includeLastOffset) {274+ if (!Ops::NN::AclnnUtil::IsRegbase()) {
248- if (maxIndices->GetViewShape().GetShapeSize() != offsets->GetViewShape().GetDim(0) - 1) {275+ int64_t expectedSize = includeLastOffset ? offsets->GetViewShape().GetDim(0) - 1 : offsets->GetViewShape().GetDim(0);
249- OP_LOGE(276+ if (maxIndices->GetViewShape().GetShapeSize() != expectedSize) {
250- ACLNN_ERR_PARAM_INVALID, "maxIndices shape size should be %ld, but got %ld.",277+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "maxIndices shape size should be %ld, but got %ld.",expectedSize,
251- offsets->GetViewShape().GetShapeSize() - 1, maxIndices->GetViewShape().GetShapeSize());278+ maxIndices->GetViewShape().GetShapeSize());
252- return false;
253- }
254- } else {
255- if (maxIndices->GetViewShape().GetShapeSize() != offsets->GetViewShape().GetDim(0)) {
256- OP_LOGE(
257- ACLNN_ERR_PARAM_INVALID, "maxIndices shape size should be %ld,but got %ld.",
258- offsets->GetViewShape().GetShapeSize(), maxIndices->GetViewShape().GetShapeSize());
259 return false;279 return false;
260 }280 }
261 }281 }
262 }282 }
263- 
264 return true;283 return true;
265}284}
266 285 
@@ -374,15 +393,17 @@ aclnnStatus aclnnEmbeddingBagGetWorkspaceSize(
374 auto bagSizeL0Cast = l0op::Cast(std::get<2>(result), bagSize->GetDataType(), uniqueExecutor.get());393 auto bagSizeL0Cast = l0op::Cast(std::get<2>(result), bagSize->GetDataType(), uniqueExecutor.get());
375 CHECK_RET(bagSizeL0Cast != nullptr, ACLNN_ERR_INNER_NULLPTR);394 CHECK_RET(bagSizeL0Cast != nullptr, ACLNN_ERR_INNER_NULLPTR);
376 auto viewCopyBagSizeL0CastResult = l0op::ViewCopy(bagSizeL0Cast, bagSize, uniqueExecutor.get());395 auto viewCopyBagSizeL0CastResult = l0op::ViewCopy(bagSizeL0Cast, bagSize, uniqueExecutor.get());
377- 396+ auto maxIndicesShapeSize = std::get<3>(result)->GetViewShape().GetShapeSize();
378- auto maxIndicesL0Cast = l0op::Cast(std::get<3>(result), maxIndices->GetDataType(), uniqueExecutor.get());397+ if (maxIndicesShapeSize != 0){
379- CHECK_RET(maxIndicesL0Cast != nullptr, ACLNN_ERR_INNER_NULLPTR);398+ auto maxIndicesL0Cast = l0op::Cast(std::get<3>(result), maxIndices->GetDataType(), uniqueExecutor.get());
380- auto viewCopyMaxIndicesL0CastResult = l0op::ViewCopy(maxIndicesL0Cast, maxIndices, uniqueExecutor.get());399+ CHECK_RET(maxIndicesL0Cast != nullptr, ACLNN_ERR_INNER_NULLPTR);
381- CHECK_RET(400+ auto viewCopyMaxIndicesL0CastResult = l0op::ViewCopy(maxIndicesL0Cast, maxIndices, uniqueExecutor.get());
382- viewCopyOutputResult != nullptr && viewCopyMaxIndicesL0CastResult != nullptr &&401+ CHECK_RET(
383- viewCopyBagSizeL0CastResult != nullptr,402+ viewCopyOutputResult != nullptr && viewCopyMaxIndicesL0CastResult != nullptr &&
384- ACLNN_ERR_INNER_NULLPTR);403+ viewCopyBagSizeL0CastResult != nullptr,
385- 404+ ACLNN_ERR_INNER_NULLPTR);
405+ }
406+
386 // 固定写法，获取计算过程中需要使用的workspace大小407 // 固定写法，获取计算过程中需要使用的workspace大小
387 *workspaceSize = uniqueExecutor->GetWorkspaceSize();408 *workspaceSize = uniqueExecutor->GetWorkspaceSize();
388 uniqueExecutor.ReleaseTo(executor); // 需要把 uniqueExecutor持有executor转移给executor409 uniqueExecutor.ReleaseTo(executor); // 需要把 uniqueExecutor持有executor转移给executor
Mindex/embedding_bag/op_host/op_api/embedding_bag.cpp+74-34
@@ -11,14 +11,18 @@
11 * \file aclnn_embedding_bag.cpp11 * \file aclnn_embedding_bag.cpp
12 * \brief12 * \brief
13 */13 */
14- 14+
15+#include "opdev/aicpu/aicpu_task.h"
16+#include "aclnn_kernels/common/op_error_check.h"
15#include "embedding_bag.h"17#include "embedding_bag.h"
16#include "opdev/format_utils.h"18#include "opdev/format_utils.h"
17#include "opdev/op_dfx.h"19#include "opdev/op_dfx.h"
20+#include "opdev/platform.h"
18#include "opdev/make_op_executor.h"21#include "opdev/make_op_executor.h"
19#include "opdev/data_type_utils.h"22#include "opdev/data_type_utils.h"
20-#include "opdev/platform.h"
21#include "opdev/shape_utils.h"23#include "opdev/shape_utils.h"
24+#include "opdev/op_def.h"
25+#include "op_api/aclnn_util.h"
22 26 
23namespace l0op {27namespace l0op {
24 28 
@@ -54,52 +58,88 @@ static op::Shape GetOutPutShape(
54 return outputShape;58 return outputShape;
55}59}
56 60 
61+static aclTensor* AllocTensorForEmbeddingBag(
62+ const op::Shape& shape910,
63+ op::DataType dtype910,
64+ const op::Shape& shapeOther,
65+ op::DataType dtypeOther,
66+ op::SocVersion socVersion,
67+ aclOpExecutor* executor)
68+{
69+ if (Ops::NN::AclnnUtil::IsRegbase()) {
70+ return executor->AllocTensor(shape910, dtype910, op::Format::FORMAT_ND);
71+ } else {
72+ return executor->AllocTensor(shapeOther, dtypeOther, op::Format::FORMAT_ND);
73+ }
74+}
75+ 
76+static std::pair<op::Shape, op::Shape> GetMaxIndicesShapes(
77+ const std::string& modeStr,
78+ bool includeLastOffset,
79+ const op::Shape& offsetsShape,
80+ const op::Shape& outputShape)
81+{
82+ op::Shape shape910;
83+ op::Shape shapeOther;
84+ 
85+ if (modeStr == "max") {
86+ shape910 = outputShape;
87+ shapeOther = outputShape;
88+ } else {
89+ shape910.AppendDim(0);
90+ if (includeLastOffset) {
91+ shapeOther.AppendDim(offsetsShape.GetDim(0) - 1);
92+ } else {
93+ shapeOther = offsetsShape;
94+ }
95+ }
96+ return std::make_pair(shape910, shapeOther);
97+}
98+ 
99+static std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*> MakeNullEmbeddingBagResult() {
100+ return std::make_tuple(
101+ static_cast<aclTensor*>(nullptr),
102+ static_cast<aclTensor*>(nullptr),
103+ static_cast<aclTensor*>(nullptr),
104+ static_cast<aclTensor*>(nullptr));
105+}
106+ 
57const std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*> EmbeddingBag(107const std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*> EmbeddingBag(
58 const aclTensor* weight, const aclTensor* indices, const aclTensor* offsets, bool scaleGradByFreq,108 const aclTensor* weight, const aclTensor* indices, const aclTensor* offsets, bool scaleGradByFreq,
59 const std::string& modeStr, bool sparse, const aclTensor* perSampleWeights, bool includeLastOffset,109 const std::string& modeStr, bool sparse, const aclTensor* perSampleWeights, bool includeLastOffset,
60 int64_t paddingIdx, aclOpExecutor* executor)110 int64_t paddingIdx, aclOpExecutor* executor)
61{111{
62 if (!IsAiCoreSupport(weight, indices, offsets, perSampleWeights)) {112 if (!IsAiCoreSupport(weight, indices, offsets, perSampleWeights)) {
63- return std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*>(113+ return MakeNullEmbeddingBagResult();
64- static_cast<aclTensor*>(nullptr), static_cast<aclTensor*>(nullptr),
65- static_cast<aclTensor*>(nullptr), static_cast<aclTensor*>(nullptr));
66 }114 }
115+ auto socVersion = op::GetCurrentPlatformInfo().GetSocVersion();
116+ op::DataType indicesDtype = indices->GetDataType();
117+ op::DataType offsetsDtype = offsets->GetDataType();
118+ op::DataType indexPromoteDtype = (indicesDtype == offsetsDtype) ? indicesDtype : op::DataType::DT_INT64;
67 //申请output_tensor的Tensor119 //申请output_tensor的Tensor
68 auto outputShape = GetOutPutShape(weight, offsets, includeLastOffset);120 auto outputShape = GetOutPutShape(weight, offsets, includeLastOffset);
69 auto outputTensor = executor->AllocTensor(outputShape, weight->GetDataType());121 auto outputTensor = executor->AllocTensor(outputShape, weight->GetDataType());
70- 122+ aclTensor* offset2bag = nullptr;
71 // 申请offset2bag的Tensor123 // 申请offset2bag的Tensor
72- op::Shape offset2bagShape;124+ op::Shape offset2bagShape910, offset2bagShapeOther;
73- offset2bagShape.AppendDim(indices->GetViewShape().GetDim(0));125+ offset2bagShape910.AppendDim(indices->GetViewShape().GetShapeSize());
74- auto offset2bag = executor->AllocTensor(offset2bagShape, indices->GetDataType(), op::Format::FORMAT_ND);126+ offset2bagShapeOther.AppendDim(indices->GetViewShape().GetShapeSize());
75- 127+ offset2bag = AllocTensorForEmbeddingBag(offset2bagShape910, indexPromoteDtype, offset2bagShapeOther, indicesDtype, socVersion, executor);
76- // 申请bagSize的Tensor128+ aclTensor* bagSize = nullptr;
77- auto bagSize = executor->AllocTensor(offsets->GetViewShape(), offsets->GetDataType(), op::Format::FORMAT_ND);129+ op::Shape bagSizeShape910, bagSizeShapeOther;
130+ bagSizeShape910.AppendDim(offsets->GetViewShape().GetDim(0));
78 if (includeLastOffset) {131 if (includeLastOffset) {
79- op::Shape bagSizeShape;132+ bagSizeShapeOther.AppendDim(offsets->GetViewShape().GetDim(0) - 1);
80- bagSizeShape.AppendDim(offsets->GetViewShape().GetDim(0) - 1);
81- bagSize = executor->AllocTensor(bagSizeShape, offsets->GetDataType(), op::Format::FORMAT_ND);
82- }
83- 
84- // 申请maxIndices的Tensor
85- aclTensor* maxIndices;
86- if (modeStr == "max") {
87- auto maxIndicesShape = GetOutPutShape(weight, offsets, includeLastOffset);
88- maxIndices = executor->AllocTensor(maxIndicesShape, offsets->GetDataType(), op::Format::FORMAT_ND);
89 } else {133 } else {
90- maxIndices = executor->AllocTensor(offsets->GetViewShape(), offsets->GetDataType(), op::Format::FORMAT_ND);134+ bagSizeShapeOther = offsets->GetViewShape();
91- if (includeLastOffset) {
92- op::Shape maxIndicesShape;
93- maxIndicesShape.AppendDim(offsets->GetViewShape().GetDim(0) - 1);
94- maxIndices = executor->AllocTensor(maxIndicesShape, offsets->GetDataType(), op::Format::FORMAT_ND);
95- }
96 }135 }
136+ bagSize = AllocTensorForEmbeddingBag(bagSizeShape910, indexPromoteDtype, bagSizeShapeOther, offsetsDtype, socVersion, executor);
137+ aclTensor* maxIndices = nullptr;
138+ auto maxIndicesShapes = GetMaxIndicesShapes(modeStr, includeLastOffset, offsets->GetViewShape(), outputShape);
139+ maxIndices = AllocTensorForEmbeddingBag(maxIndicesShapes.first, indexPromoteDtype, maxIndicesShapes.second, offsetsDtype, socVersion, executor);
97 if (outputTensor == nullptr || offset2bag == nullptr || bagSize == nullptr || maxIndices == nullptr) {140 if (outputTensor == nullptr || offset2bag == nullptr || bagSize == nullptr || maxIndices == nullptr) {
98- return std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*>(141+ return MakeNullEmbeddingBagResult();
99- static_cast<aclTensor*>(nullptr), static_cast<aclTensor*>(nullptr),
100- static_cast<aclTensor*>(nullptr), static_cast<aclTensor*>(nullptr));
101 }142 }
102- 
103 L0_DFX(143 L0_DFX(
104 EmbeddingBag, weight, indices, offsets, scaleGradByFreq, modeStr, sparse, perSampleWeights, includeLastOffset);144 EmbeddingBag, weight, indices, offsets, scaleGradByFreq, modeStr, sparse, perSampleWeights, includeLastOffset);
105 145 
@@ -109,8 +149,8 @@ const std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*> EmbeddingBag(
109 OP_ATTR(modeStr, scaleGradByFreq, sparse, includeLastOffset, paddingIdx));149 OP_ATTR(modeStr, scaleGradByFreq, sparse, includeLastOffset, paddingIdx));
110 if (ret != ACL_SUCCESS) {150 if (ret != ACL_SUCCESS) {
111 OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "EmbeddingBagAiCore ADD_TO_LAUNCHER_LIST_AICORE failed.");151 OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "EmbeddingBagAiCore ADD_TO_LAUNCHER_LIST_AICORE failed.");
112- return std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*>(nullptr, nullptr, nullptr, nullptr);152+ return MakeNullEmbeddingBagResult();
113 }153 }
114 return std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*>(outputTensor, offset2bag, bagSize, maxIndices);154 return std::tuple<aclTensor*, aclTensor*, aclTensor*, aclTensor*>(outputTensor, offset2bag, bagSize, maxIndices);
115}155}
116-} // namespace l0op156+} // namespace l0op
Mindex/embedding_bag/op_kernel/arch35/embedding_bag_regbase_1d_max.h+25-14
@@ -47,6 +47,15 @@ public:
47 this->offset2bagGm_.SetGlobalBuffer((__gm__ I*)(gmParam[OFFSET2BAG_OUTPUT_IDX]));47 this->offset2bagGm_.SetGlobalBuffer((__gm__ I*)(gmParam[OFFSET2BAG_OUTPUT_IDX]));
48 this->maxIndicesGm_.SetGlobalBuffer((__gm__ I*)(gmParam[MAXINDICES_OUTPUT_IDX]));48 this->maxIndicesGm_.SetGlobalBuffer((__gm__ I*)(gmParam[MAXINDICES_OUTPUT_IDX]));
49 49
50+ 
51+ if (GetBlockIdx() == 0){
52+ InitGlobalMemory(this->yGm_, tiling_.embeddingDim * tiling_.nBags, (T)(0));
53+ int32_t eventIDMTE3ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
54+ SetFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
55+ WaitFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
56+ InitGlobalMemory(this->maxIndicesGm_, tiling_.embeddingDim * tiling_.nBags, (I)(-1));
57+ }
58+
50 this->weightCoreOfset_ = (GetBlockIdx() % tiling_.colTileNum) * tiling_.colNormalNum;59 this->weightCoreOfset_ = (GetBlockIdx() % tiling_.colTileNum) * tiling_.colNormalNum;
51 this->offsetStart_ = GetBlockIdx() / tiling_.colTileNum * tiling_.rowNormalNum;60 this->offsetStart_ = GetBlockIdx() / tiling_.colTileNum * tiling_.rowNormalNum;
52 this->curCoreBag_ =61 this->curCoreBag_ =
@@ -61,10 +70,11 @@ public:
61 pipe_.InitBuffer(this->inQueueIndices_, 1, tiling_.indicesFactor * sizeof(U));70 pipe_.InitBuffer(this->inQueueIndices_, 1, tiling_.indicesFactor * sizeof(U));
62 pipe_.InitBuffer(this->inQueueOffsets_, 1, (tiling_.offsetsFactor + 1) * sizeof(E));71 pipe_.InitBuffer(this->inQueueOffsets_, 1, (tiling_.offsetsFactor + 1) * sizeof(E));
63 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));72 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));
64- pipe_.InitBuffer(this->outQueueOffset2bag_, 1, tiling_.weightRowFactor * sizeof(I));73+ pipe_.InitBuffer(this->outQueueOffset2bag_, tiling_.weightRowFactor * sizeof(I));
65 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));74 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));
66- pipe_.InitBuffer(this->outQueueMaxIndices_, 1, (tiling_.weightDimFactor) * sizeof(I));75+ pipe_.InitBuffer(this->outQueueMaxIndices_, (tiling_.weightDimFactor) * sizeof(I));
67 pipe_.InitBuffer(this->maxIndicesCalcBuf_, (tiling_.weightRowFactor) * sizeof(I));76 pipe_.InitBuffer(this->maxIndicesCalcBuf_, (tiling_.weightRowFactor) * sizeof(I));
77+ SyncAll();
68 }78 }
69 79 
70 __aicore__ inline void HandleBagMaxNoPerSample(80 __aicore__ inline void HandleBagMaxNoPerSample(
@@ -86,7 +96,7 @@ public:
86 WaitFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);96 WaitFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
87 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();97 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();
88 DuplicateNegInf<T>(outYLocal_, static_cast<int32_t>(tiling_.weightDimFactor));98 DuplicateNegInf<T>(outYLocal_, static_cast<int32_t>(tiling_.weightDimFactor));
89- LocalTensor<I> maxIndicesLocal_ = this->outQueueMaxIndices_.template AllocTensor<I>();99+ LocalTensor<I> maxIndicesLocal_ = this->outQueueMaxIndices_.template Get<I>();
90 Duplicate(maxIndicesLocal_, static_cast<I>(-1), static_cast<int32_t>(tiling_.weightDimFactor));100 Duplicate(maxIndicesLocal_, static_cast<I>(-1), static_cast<int32_t>(tiling_.weightDimFactor));
91 int64_t curWeightNumber =101 int64_t curWeightNumber =
92 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;102 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;
@@ -98,6 +108,7 @@ public:
98 for (uint64_t indicesLoopIdx = 0; indicesLoopIdx < indicesLoop; ++indicesLoopIdx) {108 for (uint64_t indicesLoopIdx = 0; indicesLoopIdx < indicesLoop; ++indicesLoopIdx) {
99 int64_t curIndicesNumber =109 int64_t curIndicesNumber =
100 indicesLoopIdx == indicesLoop - 1 ? tailLoopIndicesNumber : tiling_.weightRowFactor;110 indicesLoopIdx == indicesLoop - 1 ? tailLoopIndicesNumber : tiling_.weightRowFactor;
111+ int64_t validRowNumber = curIndicesNumber;
101 // Offset2bagweight112 // Offset2bagweight
102 if (weightLoopIdx == 0) {113 if (weightLoopIdx == 0) {
103 this->CopyOffset2bagToGm(indiceStart, curIndicesNumber, curOffsetStart);114 this->CopyOffset2bagToGm(indiceStart, curIndicesNumber, curOffsetStart);
@@ -112,6 +123,7 @@ public:
112 if (weightIndex == tiling_.paddingIdx) {123 if (weightIndex == tiling_.paddingIdx) {
113 validIndicesFactorNumber = validIndicesFactorNumber - 1;124 validIndicesFactorNumber = validIndicesFactorNumber - 1;
114 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;125 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;
126+ validRowNumber = validRowNumber - 1;
115 continue;127 continue;
116 }128 }
117 maxIndicesCalcLocal_(indicesIndex) = static_cast<I>(weightIndex);129 maxIndicesCalcLocal_(indicesIndex) = static_cast<I>(weightIndex);
@@ -122,7 +134,7 @@ public:
122 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;134 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;
123 }135 }
124 this->ComputeMax(136 this->ComputeMax(
125- curWeightNumber, (uint16_t)validIndicesFactorNumber, outYLocal_, weightLocal_, maxIndicesLocal_);137+ curWeightNumber, (uint16_t)validRowNumber, outYLocal_, weightLocal_, maxIndicesLocal_);
126 this->inQueueWeight_.template FreeTensor(weightLocal_);138 this->inQueueWeight_.template FreeTensor(weightLocal_);
127 indiceStart = indiceStart + curIndicesNumber;139 indiceStart = indiceStart + curIndicesNumber;
128 }140 }
@@ -137,17 +149,19 @@ public:
137 this->CopyYToGm(outYOfset, curWeightNumber, outYLocal_);149 this->CopyYToGm(outYOfset, curWeightNumber, outYLocal_);
138 this->outQueueY_.template FreeTensor(outYLocal_);150 this->outQueueY_.template FreeTensor(outYLocal_);
139 this->CopyMaxIndicesToGm(outYOfset, curWeightNumber, maxIndicesLocal_);151 this->CopyMaxIndicesToGm(outYOfset, curWeightNumber, maxIndicesLocal_);
140- this->outQueueMaxIndices_.template FreeTensor(maxIndicesLocal_);
141 weightOfset = weightOfset + curWeightNumber;152 weightOfset = weightOfset + curWeightNumber;
142 }153 }
143 bagSizeLocal_(bagIdx) = bagSizeLocal_(bagIdx) - paddingCount;154 bagSizeLocal_(bagIdx) = bagSizeLocal_(bagIdx) - paddingCount;
155+ this->inQueueIndices_.template EnQue(indicesLocal_);
156+ int32_t eventIDSToV= static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
157+ SetFlag<HardEvent::S_V>(eventIDSToV);
158+ WaitFlag<HardEvent::S_V>(eventIDSToV);
144 }159 }
145 160 
146 __aicore__ inline void HandleBigBagMax(161 __aicore__ inline void HandleBigBagMax(
147 int64_t curBagIndiceNumber, int64_t bagIndiceStart, int64_t curOffsetStart, int64_t bagIdx,162 int64_t curBagIndiceNumber, int64_t bagIndiceStart, int64_t curOffsetStart, int64_t bagIdx,
148 LocalTensor<I> bagSizeLocal_)163 LocalTensor<I> bagSizeLocal_)
149 {164 {
150- LocalTensor<U> indicesLocal_ = this->inQueueIndices_.template DeQue<U>();
151 int32_t eventIDMTE2ToS1 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));165 int32_t eventIDMTE2ToS1 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
152 SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);166 SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);
153 WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);167 WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);
@@ -159,7 +173,7 @@ public:
159 int32_t eventIDMTE3ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));173 int32_t eventIDMTE3ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
160 SetFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);174 SetFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
161 WaitFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);175 WaitFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
162- LocalTensor<I> maxIndicesLocal_ = this->outQueueMaxIndices_.template AllocTensor<I>();176+ LocalTensor<I> maxIndicesLocal_ = this->outQueueMaxIndices_.template Get<I>();
163 Duplicate(maxIndicesLocal_, static_cast<I>(-1), static_cast<int32_t>(tiling_.weightDimFactor));177 Duplicate(maxIndicesLocal_, static_cast<I>(-1), static_cast<int32_t>(tiling_.weightDimFactor));
164 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();178 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();
165 DuplicateNegInf<T>(outYLocal_, static_cast<int32_t>(tiling_.weightDimFactor));179 DuplicateNegInf<T>(outYLocal_, static_cast<int32_t>(tiling_.weightDimFactor));
@@ -227,10 +241,12 @@ public:
227 this->CopyYToGm(outYOfset, curWeightNumber, outYLocal_);241 this->CopyYToGm(outYOfset, curWeightNumber, outYLocal_);
228 this->CopyMaxIndicesToGm(outYOfset, curWeightNumber, maxIndicesLocal_);242 this->CopyMaxIndicesToGm(outYOfset, curWeightNumber, maxIndicesLocal_);
229 this->outQueueY_.template FreeTensor(outYLocal_);243 this->outQueueY_.template FreeTensor(outYLocal_);
230- this->outQueueMaxIndices_.template FreeTensor(maxIndicesLocal_);
231 weightOfset = weightOfset + curWeightNumber;244 weightOfset = weightOfset + curWeightNumber;
232 }245 }
233 bagSizeLocal_(bagIdx) = bagSizeLocal_(bagIdx) - paddingCount;246 bagSizeLocal_(bagIdx) = bagSizeLocal_(bagIdx) - paddingCount;
247+ int32_t eventIDSToV= static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
248+ SetFlag<HardEvent::S_V>(eventIDSToV);
249+ WaitFlag<HardEvent::S_V>(eventIDSToV);
234 }250 }
235 251 
236 __aicore__ inline void ProcessNoPerSample()252 __aicore__ inline void ProcessNoPerSample()
@@ -263,12 +279,6 @@ public:
263 for (uint64_t bagIdx = 0; bagIdx < curLoopOffsetsNumber; ++bagIdx) {279 for (uint64_t bagIdx = 0; bagIdx < curLoopOffsetsNumber; ++bagIdx) {
264 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);280 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);
265 if (curBagIndiceNumber == 0) {281 if (curBagIndiceNumber == 0) {
266- uint64_t dimNumber = (uint64_t)this->curCoreEmbedDim_;
267- int64_t outYofset = curOffsetStart * tiling_.embeddingDim + this->weightCoreOfset_;
268- GlobalTensor<T> yOutGm_ = this->yGm_[outYofset];
269- InitGlobalMemory(yOutGm_, dimNumber, (T)(-1));
270- GlobalTensor<I> maxIndicesOutGm_ = this->maxIndicesGm_[outYofset];
271- InitGlobalMemory(maxIndicesOutGm_, dimNumber, (I)(-1));
272 curOffsetStart = curOffsetStart + 1;282 curOffsetStart = curOffsetStart + 1;
273 continue;283 continue;
274 }284 }
@@ -307,6 +317,7 @@ public:
307 return;317 return;
308 }318 }
309 ProcessNoPerSample();319 ProcessNoPerSample();
320+ this->DisposalBagSize(this->bagSizeGm_[tiling_.nBags], tiling_.inclueLastOfst);
310 }321 }
311 322 
312private:323private:
Mindex/embedding_bag/op_kernel/arch35/embedding_bag_regbase_1d_mean.h+11-6
@@ -42,6 +42,10 @@ public:
42 this->yGm_.SetGlobalBuffer((__gm__ T*)(gmParam[Y_OUTPUT_IDX]));42 this->yGm_.SetGlobalBuffer((__gm__ T*)(gmParam[Y_OUTPUT_IDX]));
43 this->bagSizeGm_.SetGlobalBuffer((__gm__ I*)(gmParam[BAGSIZE_OUTPUT_IDX]));43 this->bagSizeGm_.SetGlobalBuffer((__gm__ I*)(gmParam[BAGSIZE_OUTPUT_IDX]));
44 this->offset2bagGm_.SetGlobalBuffer((__gm__ I*)(gmParam[OFFSET2BAG_OUTPUT_IDX]));44 this->offset2bagGm_.SetGlobalBuffer((__gm__ I*)(gmParam[OFFSET2BAG_OUTPUT_IDX]));
45+ 
46+ if (GetBlockIdx() == 0){
47+ InitGlobalMemory(this->yGm_, tiling_.embeddingDim * tiling_.nBags, (T)(0));
48+ }
45 49
46 this->weightCoreOfset_ = (GetBlockIdx() % tiling_.colTileNum) * tiling_.colNormalNum;50 this->weightCoreOfset_ = (GetBlockIdx() % tiling_.colTileNum) * tiling_.colNormalNum;
47 this->offsetStart_ = GetBlockIdx() / tiling_.colTileNum * tiling_.rowNormalNum;51 this->offsetStart_ = GetBlockIdx() / tiling_.colTileNum * tiling_.rowNormalNum;
@@ -60,8 +64,9 @@ public:
60 pipe_.InitBuffer(this->inQueueIndices_, 1, tiling_.indicesFactor * sizeof(U));64 pipe_.InitBuffer(this->inQueueIndices_, 1, tiling_.indicesFactor * sizeof(U));
61 pipe_.InitBuffer(this->inQueueOffsets_, 1, (tiling_.offsetsFactor + 1) * sizeof(E));65 pipe_.InitBuffer(this->inQueueOffsets_, 1, (tiling_.offsetsFactor + 1) * sizeof(E));
62 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));66 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));
63- pipe_.InitBuffer(this->outQueueOffset2bag_, 1, tiling_.weightRowFactor * sizeof(I));67+ pipe_.InitBuffer(this->outQueueOffset2bag_, tiling_.weightRowFactor * sizeof(I));
64 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));68 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));
69+ SyncAll();
65 }70 }
66 71 
67 __aicore__ inline void HandleBigBagMean(72 __aicore__ inline void HandleBigBagMean(
@@ -175,19 +180,21 @@ public:
175 }180 }
176 int64_t weightLocalOffset = 0;181 int64_t weightLocalOffset = 0;
177 LocalTensor<T> weightLocal_ = this->inQueueWeight_.template AllocTensor<T>();182 LocalTensor<T> weightLocal_ = this->inQueueWeight_.template AllocTensor<T>();
183+ int64_t validRowNumber = curIndicesNumber;
178 for (uint64_t indicesIdx = 0; indicesIdx < curIndicesNumber; ++indicesIdx) {184 for (uint64_t indicesIdx = 0; indicesIdx < curIndicesNumber; ++indicesIdx) {
179 int64_t indiceLocalIdx = indiceStart + indicesIdx - this->copyIndicesStart_;185 int64_t indiceLocalIdx = indiceStart + indicesIdx - this->copyIndicesStart_;
180 U weightIndex = indicesLocal_(indiceLocalIdx);186 U weightIndex = indicesLocal_(indiceLocalIdx);
181 if (weightIndex == tiling_.paddingIdx) {187 if (weightIndex == tiling_.paddingIdx) {
182 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;188 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;
183 validIndicesFactorNumber = validIndicesFactorNumber - 1;189 validIndicesFactorNumber = validIndicesFactorNumber - 1;
190+ validRowNumber = validRowNumber - 1;
184 continue;191 continue;
185 }192 }
186 int64_t weightOffset = weightIndex * tiling_.embeddingDim + weightOfset;193 int64_t weightOffset = weightIndex * tiling_.embeddingDim + weightOfset;
187 this->CopyWeightNoPerSampleFromGm(weightOffset, curWeightNumber, weightLocal_[weightLocalOffset]);194 this->CopyWeightNoPerSampleFromGm(weightOffset, curWeightNumber, weightLocal_[weightLocalOffset]);
188 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;195 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;
189 }196 }
190- this->ComputeAdd(curWeightNumber, (uint16_t)validIndicesFactorNumber, outYLocal_, weightLocal_);197+ this->ComputeAdd(curWeightNumber, (uint16_t)validRowNumber, outYLocal_, weightLocal_);
191 this->inQueueWeight_.template FreeTensor(weightLocal_);198 this->inQueueWeight_.template FreeTensor(weightLocal_);
192 indiceStart = indiceStart + curIndicesNumber;199 indiceStart = indiceStart + curIndicesNumber;
193 }200 }
@@ -205,6 +212,7 @@ public:
205 weightOfset = weightOfset + curWeightNumber;212 weightOfset = weightOfset + curWeightNumber;
206 }213 }
207 bagSizeLocal(bagIdx) = bagSizeLocal(bagIdx) - paddingCount;214 bagSizeLocal(bagIdx) = bagSizeLocal(bagIdx) - paddingCount;
215+ this->inQueueIndices_.template EnQue(indicesLocal_);
208 }216 }
209 217 
210 __aicore__ inline void ProcessNoPerSample()218 __aicore__ inline void ProcessNoPerSample()
@@ -236,10 +244,6 @@ public:
236 for (uint64_t bagIdx = 0; bagIdx < curLoopOffsetsNumber; ++bagIdx) {244 for (uint64_t bagIdx = 0; bagIdx < curLoopOffsetsNumber; ++bagIdx) {
237 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);245 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);
238 if (curBagIndiceNumber == 0) {246 if (curBagIndiceNumber == 0) {
239- int64_t outYofset = curOffsetStart * tiling_.embeddingDim + this->weightCoreOfset_;
240- uint64_t dimNumber = (uint64_t)this->curCoreEmbedDim_;
241- GlobalTensor<T> yOutGm_ = this->yGm_[outYofset];
242- InitGlobalMemory(yOutGm_, dimNumber, (T)(-1));
243 curOffsetStart = curOffsetStart + 1;247 curOffsetStart = curOffsetStart + 1;
244 continue;248 continue;
245 }249 }
@@ -280,6 +284,7 @@ public:
280 return;284 return;
281 }285 }
282 ProcessNoPerSample();286 ProcessNoPerSample();
287+ this->DisposalBagSize(this->bagSizeGm_[tiling_.nBags], tiling_.inclueLastOfst);
283 }288 }
284 289 
285private:290private:
Mindex/embedding_bag/op_kernel/arch35/embedding_bag_regbase_1d_simt.h+4-1
@@ -233,7 +233,7 @@ __aicore__ inline void EmbeddingBagRegBaseSimt1D<W, I, O, P, COMP_T>::Process()
233 COMP_T shift = 0;233 COMP_T shift = 0;
234 COMP_T chunkPerBag = ops::CeilDiv(embeddingDimSize, static_cast<COMP_T>(BLOCK_DIM_0));234 COMP_T chunkPerBag = ops::CeilDiv(embeddingDimSize, static_cast<COMP_T>(BLOCK_DIM_0));
235 GetUintDivMagicAndShift(magic, shift, chunkPerBag);235 GetUintDivMagicAndShift(magic, shift, chunkPerBag);
236- 236+
237 if (mode == MODE_MAX) {237 if (mode == MODE_MAX) {
238 Simt::VF_CALL<SimtComputeMax1D<W, I, O, P, COMP_T>>(Simt::Dim3{BLOCK_DIM_0, BLOCK_DIM_1}, 238 Simt::VF_CALL<SimtComputeMax1D<W, I, O, P, COMP_T>>(Simt::Dim3{BLOCK_DIM_0, BLOCK_DIM_1},
239 (__gm__ W*)(weightGm_.GetPhyAddr()), 239 (__gm__ W*)(weightGm_.GetPhyAddr()),
@@ -277,6 +277,9 @@ __aicore__ inline void EmbeddingBagRegBaseSimt1D<W, I, O, P, COMP_T>::Process()
277 numBags, numIndices, chunkPerBag, magic, shift,277 numBags, numIndices, chunkPerBag, magic, shift,
278 embeddingDimSize, paddingIdx);278 embeddingDimSize, paddingIdx);
279 }279 }
280+ if (tilingData_.inclueLastOfst == 1) {
281+ bagSizeGm_(numBags) = 0;
282+ }
280}283}
281}284}
282 285 
Mindex/embedding_bag/op_kernel/arch35/embedding_bag_regbase_1d_sum.h+17-22
@@ -45,6 +45,10 @@ public:
45 if (GetBlockIdx() >= tiling_.usedCoreNum) {45 if (GetBlockIdx() >= tiling_.usedCoreNum) {
46 return;46 return;
47 }47 }
48+ if (GetBlockIdx() == 0){
49+ InitGlobalMemory(this->yGm_, tiling_.embeddingDim * tiling_.nBags, (T)(0));
50+ }
51+
48 this->weightCoreOfset_ = (GetBlockIdx() % tiling_.colTileNum) * tiling_.colNormalNum;52 this->weightCoreOfset_ = (GetBlockIdx() % tiling_.colTileNum) * tiling_.colNormalNum;
49 this->offsetStart_ = GetBlockIdx() / tiling_.colTileNum * tiling_.rowNormalNum;53 this->offsetStart_ = GetBlockIdx() / tiling_.colTileNum * tiling_.rowNormalNum;
50 this->curCoreBag_ =54 this->curCoreBag_ =
@@ -64,7 +68,8 @@ public:
64 pipe_.InitBuffer(this->inQueueIndices_, 1, tiling_.indicesFactor * sizeof(U));68 pipe_.InitBuffer(this->inQueueIndices_, 1, tiling_.indicesFactor * sizeof(U));
65 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));69 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));
66 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));70 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));
67- pipe_.InitBuffer(this->outQueueOffset2bag_, 1, tiling_.weightRowFactor * sizeof(I));71+ pipe_.InitBuffer(this->outQueueOffset2bag_, tiling_.weightRowFactor * sizeof(I));
72+ SyncAll();
68 }73 }
69 74 
70 __aicore__ inline void HandleBigBagSumNoPerSample(75 __aicore__ inline void HandleBigBagSumNoPerSample(
@@ -78,7 +83,6 @@ public:
78 int32_t eventIDMTE2ToS1 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));83 int32_t eventIDMTE2ToS1 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
79 SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);84 SetFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);
80 WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);85 WaitFlag<HardEvent::MTE2_S>(eventIDMTE2ToS1);
81- LocalTensor<U> indicesLocal_ = this->inQueueIndices_.template DeQue<U>();
82 for (uint64_t weightLoopIdx = 0; weightLoopIdx < this->weightLoop_; ++weightLoopIdx) {86 for (uint64_t weightLoopIdx = 0; weightLoopIdx < this->weightLoop_; ++weightLoopIdx) {
83 int32_t eventIDMTE3ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));87 int32_t eventIDMTE3ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
84 SetFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);88 SetFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
@@ -87,7 +91,6 @@ public:
87 SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);91 SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
88 WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);92 WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
89 int64_t indiceStart = bagIndiceStart;93 int64_t indiceStart = bagIndiceStart;
90- int64_t validIndicesFactorNumber = curBagIndiceNumber;
91 int64_t curWeightNumber =94 int64_t curWeightNumber =
92 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;95 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;
93 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();96 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();
@@ -118,7 +121,6 @@ public:
118 U weightIndex = indicesLocal_(indiceLocalIdx);121 U weightIndex = indicesLocal_(indiceLocalIdx);
119 if (weightIndex == tiling_.paddingIdx) {122 if (weightIndex == tiling_.paddingIdx) {
120 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;123 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;
121- validIndicesFactorNumber = validIndicesFactorNumber - 1;
122 validCopyIndices = validCopyIndices - 1;124 validCopyIndices = validCopyIndices - 1;
123 continue;125 continue;
124 }126 }
@@ -167,30 +169,30 @@ public:
167 int32_t eventIDMTE3ToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));169 int32_t eventIDMTE3ToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));
168 SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);170 SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
169 WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);171 WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
170- int64_t validIndicesFactorNumber = curBagIndiceNumber;
171 int64_t indiceStart = bagIndiceStart;172 int64_t indiceStart = bagIndiceStart;
172 for (uint64_t indicesLoopIdx = 0; indicesLoopIdx < indicesLoop; ++indicesLoopIdx) {173 for (uint64_t indicesLoopIdx = 0; indicesLoopIdx < indicesLoop; ++indicesLoopIdx) {
173 int64_t curIndicesNumber =174 int64_t curIndicesNumber =
174- indicesLoopIdx == indicesLoop - 1 ? tailLoopIndicesNumber : tiling_.indicesFactor;175+ indicesLoopIdx == indicesLoop - 1 ? tailLoopIndicesNumber : tiling_.weightRowFactor;
175 // Offset2bagweight176 // Offset2bagweight
176 if (weightLoopIdx == 0) {177 if (weightLoopIdx == 0) {
177 this->CopyOffset2bagToGm(indiceStart, curIndicesNumber, curOffsetStart);178 this->CopyOffset2bagToGm(indiceStart, curIndicesNumber, curOffsetStart);
178 }179 }
179 int64_t weightLocalOffset = 0;180 int64_t weightLocalOffset = 0;
181+ int64_t validRowNumber = curIndicesNumber;
180 LocalTensor<T> weightLocal_ = this->inQueueWeight_.template AllocTensor<T>();182 LocalTensor<T> weightLocal_ = this->inQueueWeight_.template AllocTensor<T>();
181 for (uint64_t indicesIdx = 0; indicesIdx < curIndicesNumber; ++indicesIdx) {183 for (uint64_t indicesIdx = 0; indicesIdx < curIndicesNumber; ++indicesIdx) {
182 int64_t indiceLocalIdx = indiceStart + indicesIdx - this->copyIndicesStart_;184 int64_t indiceLocalIdx = indiceStart + indicesIdx - this->copyIndicesStart_;
183 U weightIndex = indicesLocal_(indiceLocalIdx);185 U weightIndex = indicesLocal_(indiceLocalIdx);
184 if (weightIndex == tiling_.paddingIdx) {186 if (weightIndex == tiling_.paddingIdx) {
185 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;187 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;
186- validIndicesFactorNumber = validIndicesFactorNumber - 1;188+ validRowNumber = validRowNumber - 1;
187 continue;189 continue;
188 }190 }
189 int64_t weightOffset = weightIndex * tiling_.embeddingDim + weightOfset;191 int64_t weightOffset = weightIndex * tiling_.embeddingDim + weightOfset;
190 this->CopyWeightNoPerSampleFromGm(weightOffset, curWeightNumber, weightLocal_[weightLocalOffset]);192 this->CopyWeightNoPerSampleFromGm(weightOffset, curWeightNumber, weightLocal_[weightLocalOffset]);
191 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;193 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;
192 }194 }
193- this->ComputeAdd(curWeightNumber, (uint16_t)validIndicesFactorNumber, outYLocal_, weightLocal_);195+ this->ComputeAdd(curWeightNumber, (uint16_t)validRowNumber, outYLocal_, weightLocal_);
194 this->inQueueWeight_.template FreeTensor(weightLocal_);196 this->inQueueWeight_.template FreeTensor(weightLocal_);
195 indiceStart = indiceStart + curIndicesNumber;197 indiceStart = indiceStart + curIndicesNumber;
196 }198 }
@@ -203,6 +205,7 @@ public:
203 weightOfset = weightOfset + curWeightNumber;205 weightOfset = weightOfset + curWeightNumber;
204 }206 }
205 bagSizeLocal(bagIdx) = bagSizeLocal(bagIdx) - paddingCount;207 bagSizeLocal(bagIdx) = bagSizeLocal(bagIdx) - paddingCount;
208+ this->inQueueIndices_.template EnQue(indicesLocal_);
206 }209 }
207 210 
208 __aicore__ inline void ProcessNoPerSample()211 __aicore__ inline void ProcessNoPerSample()
@@ -233,10 +236,6 @@ public:
233 for (uint64_t bagIdx = 0; bagIdx < curLoopOffsetsNumber; ++bagIdx) {236 for (uint64_t bagIdx = 0; bagIdx < curLoopOffsetsNumber; ++bagIdx) {
234 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);237 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);
235 if (curBagIndiceNumber == 0) {238 if (curBagIndiceNumber == 0) {
236- uint64_t dimNumber = (uint64_t)this->curCoreEmbedDim_;
237- int64_t outYofset = curOffsetStart * tiling_.embeddingDim + this->weightCoreOfset_;
238- GlobalTensor<T> yOutGm_ = this->yGm_[outYofset];
239- InitGlobalMemory(yOutGm_, dimNumber, (T)(-1));
240 curOffsetStart = curOffsetStart + 1;239 curOffsetStart = curOffsetStart + 1;
241 continue;240 continue;
242 }241 }
@@ -288,18 +287,18 @@ public:
288 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;287 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;
289 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();288 LocalTensor<T> outYLocal_ = this->outQueueY_.template AllocTensor<T>();
290 Duplicate(outYLocal_, static_cast<T>(0), static_cast<int32_t>(tiling_.weightDimFactor));289 Duplicate(outYLocal_, static_cast<T>(0), static_cast<int32_t>(tiling_.weightDimFactor));
291- int64_t validIndicesFactorNumber = curBagIndiceNumber;
292 int64_t indiceStart = bagIndiceStart;290 int64_t indiceStart = bagIndiceStart;
293 int32_t eventIDMTE3ToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));291 int32_t eventIDMTE3ToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));
294 SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);292 SetFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
295 WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);293 WaitFlag<HardEvent::MTE3_S>(eventIDMTE3ToS);
296 for (uint64_t indicesLoopIdx = 0; indicesLoopIdx < indicesLoop; ++indicesLoopIdx) {294 for (uint64_t indicesLoopIdx = 0; indicesLoopIdx < indicesLoop; ++indicesLoopIdx) {
297 int64_t curIndicesNumber =295 int64_t curIndicesNumber =
298- indicesLoopIdx == indicesLoop - 1 ? tailLoopIndicesNumber : tiling_.indicesFactor;296+ indicesLoopIdx == indicesLoop - 1 ? tailLoopIndicesNumber : tiling_.weightRowFactor;
299 // Offset2bagweight297 // Offset2bagweight
300 if (weightLoopIdx == 0) {298 if (weightLoopIdx == 0) {
301 this->CopyOffset2bagToGm(indiceStart, curIndicesNumber, curOffsetStart);299 this->CopyOffset2bagToGm(indiceStart, curIndicesNumber, curOffsetStart);
302 }300 }
301+ int64_t validRowNumber = curIndicesNumber;
303 int64_t indicesIndex = 0;302 int64_t indicesIndex = 0;
304 int64_t weightLocalOffset = 0;303 int64_t weightLocalOffset = 0;
305 LocalTensor<T> weightLocal_ = this->inQueueWeight_.template AllocTensor<T>();304 LocalTensor<T> weightLocal_ = this->inQueueWeight_.template AllocTensor<T>();
@@ -309,7 +308,7 @@ public:
309 U weightIndex = indicesLocal_(indiceLocalIdx);308 U weightIndex = indicesLocal_(indiceLocalIdx);
310 if (weightIndex == tiling_.paddingIdx) {309 if (weightIndex == tiling_.paddingIdx) {
311 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;310 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;
312- validIndicesFactorNumber = validIndicesFactorNumber - 1;311+ validRowNumber = validRowNumber - 1;
313 continue;312 continue;
314 }313 }
315 if (indiceStart + indicesIdx <= tiling_.sampleWeightNum - 1) {314 if (indiceStart + indicesIdx <= tiling_.sampleWeightNum - 1) {
@@ -330,7 +329,7 @@ public:
330 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;329 weightLocalOffset = weightLocalOffset + tiling_.weightDimFactor;
331 }330 }
332 this->ComputeAddPerSample(331 this->ComputeAddPerSample(
333- curWeightNumber, (uint16_t)validIndicesFactorNumber, outYLocal_, weightLocal_,332+ curWeightNumber, (uint16_t)validRowNumber, outYLocal_, weightLocal_,
334 perSampleWeightCountLocal);333 perSampleWeightCountLocal);
335 this->inQueueWeight_.template FreeTensor(weightLocal_);334 this->inQueueWeight_.template FreeTensor(weightLocal_);
336 indiceStart = indiceStart + curIndicesNumber;335 indiceStart = indiceStart + curIndicesNumber;
@@ -344,6 +343,7 @@ public:
344 this->outQueueY_.template FreeTensor(outYLocal_);343 this->outQueueY_.template FreeTensor(outYLocal_);
345 }344 }
346 bagSizeLocal(bagIdx) = bagSizeLocal(bagIdx) - paddingCount;345 bagSizeLocal(bagIdx) = bagSizeLocal(bagIdx) - paddingCount;
346+ this->inQueueIndices_.template EnQue(indicesLocal_);
347 }347 }
348 348 
349 __aicore__ inline void HandleBigBagSumPerSample(349 __aicore__ inline void HandleBigBagSumPerSample(
@@ -361,7 +361,6 @@ public:
361 int64_t curWeightNumber =361 int64_t curWeightNumber =
362 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;362 weightLoopIdx == this->weightLoop_ - 1 ? this->tailLoopWeightNmber_ : tiling_.weightDimFactor;
363 int64_t indiceStart = bagIndiceStart;363 int64_t indiceStart = bagIndiceStart;
364- int64_t validIndicesFactorNumber = curBagIndiceNumber;
365 int32_t eventIDMTE3ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));364 int32_t eventIDMTE3ToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
366 SetFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);365 SetFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
367 WaitFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);366 WaitFlag<HardEvent::MTE3_V>(eventIDMTE3ToV);
@@ -400,7 +399,6 @@ public:
400 int64_t indiceLocalIdx = indiceStart + indicesIdx - this->copyIndicesStart_;399 int64_t indiceLocalIdx = indiceStart + indicesIdx - this->copyIndicesStart_;
401 U weightIndex = indicesLocal_(indiceLocalIdx);400 U weightIndex = indicesLocal_(indiceLocalIdx);
402 if (weightIndex == tiling_.paddingIdx) {401 if (weightIndex == tiling_.paddingIdx) {
403- validIndicesFactorNumber = validIndicesFactorNumber - 1;
404 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;402 paddingCount = weightLoopIdx == 0 ? paddingCount + 1 : paddingCount;
405 validCopyIndices = validCopyIndices - 1;403 validCopyIndices = validCopyIndices - 1;
406 continue;404 continue;
@@ -478,10 +476,6 @@ public:
478 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);476 int64_t curBagIndiceNumber = (int64_t)bagSizeLocal(bagIdx);
479 /* bag size is 0 */477 /* bag size is 0 */
480 if (curBagIndiceNumber == 0) {478 if (curBagIndiceNumber == 0) {
481- int64_t outYofset = curOffsetStart * tiling_.embeddingDim + this->weightCoreOfset_;
482- uint64_t dimNumbers = (uint64_t)this->curCoreEmbedDim_;
483- GlobalTensor<T> yOutGm_ = this->yGm_[outYofset];
484- InitGlobalMemory(yOutGm_, dimNumbers, (T)(-1));
485 curOffsetStart = curOffsetStart + 1;479 curOffsetStart = curOffsetStart + 1;
486 continue;480 continue;
487 }481 }
@@ -529,6 +523,7 @@ public:
529 } else {523 } else {
530 ProcessNoPerSample();524 ProcessNoPerSample();
531 }525 }
526+ this->DisposalBagSize(this->bagSizeGm_[tiling_.nBags], tiling_.inclueLastOfst);
532 }527 }
533 528 
534private:529private:
Mindex/embedding_bag/op_kernel/arch35/embedding_bag_regbase_common.h+25-30
@@ -74,7 +74,8 @@ __aicore__ inline void ComputeSumWithWeight(LocalTensor<T> weightLocal, LocalTe
74 uint32_t vfLen = platform::GetVRegSize() / sizeof(float);74 uint32_t vfLen = platform::GetVRegSize() / sizeof(float);
75 uint16_t loopCnt = (uint16_t)((colNum + vfLen) / vfLen);75 uint16_t loopCnt = (uint16_t)((colNum + vfLen) / vfLen);
76 uint32_t counter = colNum;76 uint32_t counter = colNum;
77- 77+ uint16_t numRows = static_cast<uint16_t>(rowNum);
78+
78 __VEC_SCOPE__79 __VEC_SCOPE__
79 {80 {
80 AscendC::MicroAPI::MaskReg maskRegUpdate;81 AscendC::MicroAPI::MaskReg maskRegUpdate;
@@ -84,7 +85,7 @@ __aicore__ inline void ComputeSumWithWeight(LocalTensor<T> weightLocal, LocalTe
84 uint32_t colOffset = i * vfLen; //85 uint32_t colOffset = i * vfLen; //
85 maskRegUpdate = AscendC::MicroAPI::UpdateMask<float>(counter);86 maskRegUpdate = AscendC::MicroAPI::UpdateMask<float>(counter);
86 ops::LoadOneTensorForDtypeT<T>(sumLocalAddr, sumCastReg, maskRegUpdate, colOffset);87 ops::LoadOneTensorForDtypeT<T>(sumLocalAddr, sumCastReg, maskRegUpdate, colOffset);
87- for (uint16_t j = 0; j < rowNum; ++j) {88+ for (uint16_t j = 0; j < numRows; ++j) {
88 uint32_t weightOffset = colOffset + j * colNumAlign;89 uint32_t weightOffset = colOffset + j * colNumAlign;
89 ops::LoadOneTensorForDtypeT<T>(weightLocalAddr, weightCastReg, maskRegUpdate, weightOffset);90 ops::LoadOneTensorForDtypeT<T>(weightLocalAddr, weightCastReg, maskRegUpdate, weightOffset);
90 AscendC::MicroAPI::Muls(weightCastReg, weightCastReg, val, maskRegUpdate);91 AscendC::MicroAPI::Muls(weightCastReg, weightCastReg, val, maskRegUpdate);
@@ -106,7 +107,7 @@ __aicore__ inline void ComputeSum(LocalTensor<T> weightLocal, LocalTensor<T> su
106 uint32_t vfLen = platform::GetVRegSize() / sizeof(float);107 uint32_t vfLen = platform::GetVRegSize() / sizeof(float);
107 uint16_t loopCnt = (uint16_t)((colNum + vfLen) / vfLen);108 uint16_t loopCnt = (uint16_t)((colNum + vfLen) / vfLen);
108 uint32_t counter = colNum;109 uint32_t counter = colNum;
109- 110+ uint16_t numRows = static_cast<uint16_t>(rowNum);
110 __VEC_SCOPE__111 __VEC_SCOPE__
111 {112 {
112 AscendC::MicroAPI::RegTensor<float> weightCastReg;113 AscendC::MicroAPI::RegTensor<float> weightCastReg;
@@ -117,7 +118,7 @@ __aicore__ inline void ComputeSum(LocalTensor<T> weightLocal, LocalTensor<T> su
117 maskRegUpdate = AscendC::MicroAPI::UpdateMask<float>(counter);118 maskRegUpdate = AscendC::MicroAPI::UpdateMask<float>(counter);
118 uint32_t colOffset = i * vfLen; //119 uint32_t colOffset = i * vfLen; //
119 ops::LoadOneTensorForDtypeT<T>(sumLocalAddr, sumCastReg, maskRegUpdate, colOffset);120 ops::LoadOneTensorForDtypeT<T>(sumLocalAddr, sumCastReg, maskRegUpdate, colOffset);
120- for (uint16_t j = 0; j < rowNum; ++j) {121+ for (uint16_t j = 0; j < numRows; ++j) {
121 uint32_t weightOffset = colOffset + j * colNumAlign;122 uint32_t weightOffset = colOffset + j * colNumAlign;
122 ops::LoadOneTensorForDtypeT<T>(weightLocalAddr, weightCastReg, maskRegUpdate, weightOffset);123 ops::LoadOneTensorForDtypeT<T>(weightLocalAddr, weightCastReg, maskRegUpdate, weightOffset);
123 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(124 AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(
@@ -163,9 +164,9 @@ public:
163 TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueperSampleWeights_;164 TQue<QuePosition::VECIN, DOUBLE_BUFFER> inQueueperSampleWeights_;
164 165 
165 TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueY_;166 TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueY_;
166- TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueOffset2bag_;167+ TBuf<QuePosition::VECOUT> outQueueOffset2bag_;
167 TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueBagSize_;168 TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueBagSize_;
168- TQue<QuePosition::VECOUT, DOUBLE_BUFFER> outQueueMaxIndices_;169+ TBuf<QuePosition::VECOUT> outQueueMaxIndices_;
169 170 
170 TBuf<QuePosition::VECCALC> maxIndicesCalcBuf_;171 TBuf<QuePosition::VECCALC> maxIndicesCalcBuf_;
171 TBuf<QuePosition::VECCALC> compareMaskBuf_;172 TBuf<QuePosition::VECCALC> compareMaskBuf_;
@@ -300,14 +301,13 @@ public:
300 __aicore__ inline void CopyOffset2bagToGm(301 __aicore__ inline void CopyOffset2bagToGm(
301 int64_t offset2bagOffset, int64_t outOffset2bagNumber, int64_t curOffsetStart)302 int64_t offset2bagOffset, int64_t outOffset2bagNumber, int64_t curOffsetStart)
302 {303 {
303- LocalTensor<I> offset2bagLocal_ = outQueueOffset2bag_.AllocTensor<I>();304+ LocalTensor<I> offset2bagLocal_ = outQueueOffset2bag_.Get<I>();
304 Duplicate(offset2bagLocal_, static_cast<I>(curOffsetStart), static_cast<int32_t>(outOffset2bagNumber));305 Duplicate(offset2bagLocal_, static_cast<I>(curOffsetStart), static_cast<int32_t>(outOffset2bagNumber));
305 int32_t eventIDVToMTE3 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));306 int32_t eventIDVToMTE3 = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
306 SetFlag<HardEvent::V_MTE3>(eventIDVToMTE3);307 SetFlag<HardEvent::V_MTE3>(eventIDVToMTE3);
307 WaitFlag<HardEvent::V_MTE3>(eventIDVToMTE3);308 WaitFlag<HardEvent::V_MTE3>(eventIDVToMTE3);
308 DataCopyExtParams copyParam{1, static_cast<uint32_t>(outOffset2bagNumber * sizeof(I)), 0, 0, 0};309 DataCopyExtParams copyParam{1, static_cast<uint32_t>(outOffset2bagNumber * sizeof(I)), 0, 0, 0};
309 DataCopyPad(offset2bagGm_[offset2bagOffset], offset2bagLocal_, copyParam);310 DataCopyPad(offset2bagGm_[offset2bagOffset], offset2bagLocal_, copyParam);
310- outQueueOffset2bag_.FreeTensor(offset2bagLocal_);
311 }311 }
312 312 
313 __aicore__ inline void ComputeBagSize(int64_t number, LocalTensor<E> offsetsLocal)313 __aicore__ inline void ComputeBagSize(int64_t number, LocalTensor<E> offsetsLocal)
@@ -443,8 +443,8 @@ public:
443 }443 }
444 444 
445 __aicore__ inline void ComputeMax(445 __aicore__ inline void ComputeMax(
446- int64_t number, uint16_t indicesNumber, LocalTensor<T> outYLocal, LocalTensor<T> weightLocal,446+ int64_t number, uint16_t indicesNumber, LocalTensor<T>& outYLocal, LocalTensor<T>& weightLocal,
447- LocalTensor<I> maxIndicesLocal)447+ LocalTensor<I>& maxIndicesLocal)
448 {448 {
449 int32_t eventIDSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));449 int32_t eventIDSToV = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
450 SetFlag<HardEvent::S_V>(eventIDSToV);450 SetFlag<HardEvent::S_V>(eventIDSToV);
@@ -497,7 +497,7 @@ public:
497 for (uint16_t j = 0; j < indicesNumber; ++j) {497 for (uint16_t j = 0; j < indicesNumber; ++j) {
498 auto weightLocalAddrUpdate = weightLocalAddr + i * vfISizeNum_ + j * weightDimFactor_;498 auto weightLocalAddrUpdate = weightLocalAddr + i * vfISizeNum_ + j * weightDimFactor_;
499 AscendC::MicroAPI::DataCopy(weightReg, weightLocalAddrUpdate);499 AscendC::MicroAPI::DataCopy(weightReg, weightLocalAddrUpdate);
500- AscendC::MicroAPI::Duplicate(maxIndicesCalcReg, maxIndicesCalcLocalAddr[j]);500+ AscendC::MicroAPI::Duplicate(maxIndicesCalcReg, maxIndicesCalcLocalAddr[j]); // 当前处理行的索引
501 AscendC::MicroAPI::Compare<T, CMPMODE::GT>(cmpMask, weightReg, outYReg, tMaskRegUpdate);501 AscendC::MicroAPI::Compare<T, CMPMODE::GT>(cmpMask, weightReg, outYReg, tMaskRegUpdate);
502 if constexpr (sizeof(I) / sizeof(T) == DIGIT_1) {502 if constexpr (sizeof(I) / sizeof(T) == DIGIT_1) {
503 AscendC::MicroAPI::Select(maxIndicesReg, maxIndicesCalcReg, maxIndicesReg, cmpMask);503 AscendC::MicroAPI::Select(maxIndicesReg, maxIndicesCalcReg, maxIndicesReg, cmpMask);
@@ -517,6 +517,10 @@ public:
517 AscendC::MicroAPI::DataCopy(maxIndicesLocalAddrUpdate, maxIndicesReg, iMaskRegUpdate);517 AscendC::MicroAPI::DataCopy(maxIndicesLocalAddrUpdate, maxIndicesReg, iMaskRegUpdate);
518 }518 }
519 }519 }
520+ 
521+ int32_t eventIDVToS = static_cast<int32_t>(GetTPipePtr()->FetchEventID(HardEvent::V_S));
522+ SetFlag<HardEvent::V_S>(eventIDVToS);
523+ WaitFlag<HardEvent::V_S>(eventIDVToS);
520 }524 }
521 525 
522 __aicore__ inline void ComputeWeightAdd(526 __aicore__ inline void ComputeWeightAdd(
@@ -532,25 +536,6 @@ public:
532 AscendC::Add(outYLocal_, weightLocal_, outYLocal_, weightNumber);536 AscendC::Add(outYLocal_, weightLocal_, outYLocal_, weightNumber);
533 }537 }
534 538 
535- __aicore__ inline void ComputeWeightMax(int64_t weightNumber, int64_t weightIdx)
536- {
537- LocalTensor<T> outYLocal_ = outQueueY_.DeQue<T>();
538- LocalTensor<T> weightLocal_ = inQueueWeight_.DeQue<T>();
539- LocalTensor<U> maxIndicesLocal_ = outQueueMaxIndices_.DeQue<U>();
540- 
541- LocalTensor<uint8_t> compareMaskBufLocal_ = compareMaskBuf_.Get<uint8_t>();
542- LocalTensor<U> maxIndicesCalcLocal_ = maxIndicesCalcBuf_.Get<U>();
543- AscendC::Duplicate(maxIndicesCalcLocal_, static_cast<U>(weightIdx), static_cast<int32_t>(weightNumber));
544- AscendC::Compare(compareMaskBufLocal_, weightLocal_, outYLocal_, CMPMODE::GE, weightNumber);
545- AscendC::Select(
546- maxIndicesLocal_, compareMaskBufLocal_, maxIndicesCalcLocal_, maxIndicesLocal_,
547- SELMODE::VSEL_TENSOR_TENSOR_MODE, weightNumber);
548- AscendC::Max(outYLocal_, weightLocal_, outYLocal_, weightNumber);
549- outQueueY_.EnQue(outYLocal_);
550- inQueueWeight_.EnQue(weightLocal_);
551- outQueueMaxIndices_.EnQue(maxIndicesLocal_);
552- }
553- 
554 __aicore__ inline void ComputeWeightMean(int64_t weightNumber)539 __aicore__ inline void ComputeWeightMean(int64_t weightNumber)
555 {540 {
556 LocalTensor<T> weightLocal_ = inQueueWeight_.DeQue<T>();541 LocalTensor<T> weightLocal_ = inQueueWeight_.DeQue<T>();
@@ -573,6 +558,16 @@ public:
573 DataCopyExtParams copyParam{1, static_cast<uint32_t>(weightNumber * sizeof(U)), 0, 0, 0};558 DataCopyExtParams copyParam{1, static_cast<uint32_t>(weightNumber * sizeof(U)), 0, 0, 0};
574 DataCopyPad(maxIndicesGm_[weightOffset], maxIndicesLocal, copyParam);559 DataCopyPad(maxIndicesGm_[weightOffset], maxIndicesLocal, copyParam);
575 }560 }
561+ 
562+ __aicore__ inline void DisposalBagSize(GlobalTensor<I> globalTensor, int64_t inclueLastOfst)
563+ {
564+ if (inclueLastOfst){
565+ LocalTensor<I> dataDupLocal = maxIndicesCalcBuf_.Get<I>();
566+ Duplicate(dataDupLocal, (I)0, 1);
567+ DataCopyExtParams copyParam{1, static_cast<uint32_t>(1 * sizeof(I)), 0, 0, 0};
568+ DataCopyPad(globalTensor, dataDupLocal, copyParam);
569+ }
570+ }
576};571};
577 572 
578#endif // EMBEDDING_BAG_H_REGBASE_COMMON_H573#endif // EMBEDDING_BAG_H_REGBASE_COMMON_H
Mindex/embedding_bag/tests/ut/op_host/test_embedding_bag_regbase_tiling.cpp+55-42
@@ -62,6 +62,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_0)
62 map<string, string> aicore_spec;62 map<string, string> aicore_spec;
63 map<string, string> intrinsics;63 map<string, string> intrinsics;
64 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};64 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};
65+ map<string, string> npuarchs = {{"NpuArch", "3510"}};
65 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);66 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
66 67 
67 // platform info68 // platform info
@@ -87,6 +88,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_0)
87 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");88 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
88 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(89 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
89 "AICoreintrinsicDtypeMap", intrinsics);90 "AICoreintrinsicDtypeMap", intrinsics);
91+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
90 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);92 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
91 93 
92 // tilingFunc simulate94 // tilingFunc simulate
@@ -134,6 +136,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_0)
134 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");136 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
135 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);137 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
136 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);138 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
139+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
137 140 
138 // workspaces nullptr return failed141 // workspaces nullptr return failed
139 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);142 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
@@ -162,6 +165,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_1)
162 map<string, string> aicore_spec;165 map<string, string> aicore_spec;
163 map<string, string> intrinsics;166 map<string, string> intrinsics;
164 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};167 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};
168+ map<string, string> npuarchs = {{"NpuArch", "3510"}};
165 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);169 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
166 170 
167 // platform info171 // platform info
@@ -187,6 +191,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_1)
187 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");191 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
188 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(192 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
189 "AICoreintrinsicDtypeMap", intrinsics);193 "AICoreintrinsicDtypeMap", intrinsics);
194+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
190 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);195 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
191 196 
192 // tilingFunc simulate197 // tilingFunc simulate
@@ -234,6 +239,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_1)
234 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");239 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
235 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);240 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
236 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);241 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
242+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
237 243 
238 // workspaces nullptr return failed244 // workspaces nullptr return failed
239 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);245 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
@@ -262,6 +268,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_3)
262 map<string, string> aicore_spec;268 map<string, string> aicore_spec;
263 map<string, string> intrinsics;269 map<string, string> intrinsics;
264 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};270 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};
271+ map<string, string> npuarchs = {{"NpuArch", "3510"}};
265 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);272 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
266 273 
267 // platform info274 // platform info
@@ -287,6 +294,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_3)
287 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");294 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
288 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(295 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
289 "AICoreintrinsicDtypeMap", intrinsics);296 "AICoreintrinsicDtypeMap", intrinsics);
297+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
290 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);298 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
291 299 
292 // tilingFunc simulate300 // tilingFunc simulate
@@ -334,6 +342,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_3)
334 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");342 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
335 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);343 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
336 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);344 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
345+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
337 346 
338 // workspaces nullptr return failed347 // workspaces nullptr return failed
339 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);348 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
@@ -362,6 +371,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_300)
362 map<string, string> aicore_spec;371 map<string, string> aicore_spec;
363 map<string, string> intrinsics;372 map<string, string> intrinsics;
364 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};373 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};
374+ map<string, string> npuarchs = {{"NpuArch", "3510"}};
365 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);375 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
366 376 
367 fe::PlatFormInfos platform_info;377 fe::PlatFormInfos platform_info;
@@ -386,6 +396,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_300)
386 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");396 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
387 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(397 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
388 "AICoreintrinsicDtypeMap", intrinsics);398 "AICoreintrinsicDtypeMap", intrinsics);
399+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
389 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);400 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
390 401 
391 // tilingFunc simulate402 // tilingFunc simulate
@@ -393,9 +404,9 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_300)
393 ASSERT_NE(param, nullptr);404 ASSERT_NE(param, nullptr);
394 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);405 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
395 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());406 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
396- 407+ 
397 // embeddingDim=64, weightTypeSize=4(float) -> 64*4=256 <= 256408 // embeddingDim=64, weightTypeSize=4(float) -> 64*4=256 <= 256
398- gert::StorageShape input_0 = {{1024, 64}, {1024, 64}}; 409+ gert::StorageShape input_0 = {{1024, 64}, {1024, 64}};
399 gert::StorageShape input_1 = {{512}, {512}}; // indices410 gert::StorageShape input_1 = {{512}, {512}}; // indices
400 gert::StorageShape input_2 = {{32}, {32}}; // offsets411 gert::StorageShape input_2 = {{32}, {32}}; // offsets
401 gert::StorageShape output_shape0 = {{32, 64}, {32, 64}}; // [num_bags, embedding_dim]412 gert::StorageShape output_shape0 = {{32, 64}, {32, 64}}; // [num_bags, embedding_dim]
@@ -419,7 +430,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_300)
419 .PlatformInfo(reinterpret_cast<char*>(&platform_info))430 .PlatformInfo(reinterpret_cast<char*>(&platform_info))
420 .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // weight: float431 .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // weight: float
421 .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // indices: int32432 .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // indices: int32
422- .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // 433+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) //
423 .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // y: float434 .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // y: float
424 .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offset2bag: int32435 .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offset2bag: int32
425 .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // bag_size: int32436 .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // bag_size: int32
@@ -435,9 +446,10 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_300)
435 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");446 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
436 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);447 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
437 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);448 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
438- 449+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
450+ 
439 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);451 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
440- 452+ 
441 auto tiling_key = tiling_context->GetTilingKey();453 auto tiling_key = tiling_context->GetTilingKey();
442 ASSERT_EQ(tiling_key, 300);454 ASSERT_EQ(tiling_key, 300);
443}455}
@@ -462,18 +474,18 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_301)
462 map<string, string> aicore_spec;474 map<string, string> aicore_spec;
463 map<string, string> intrinsics;475 map<string, string> intrinsics;
464 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};476 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};
477+ map<string, string> npuarchs = {{"NpuArch", "3510"}};
465 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);478 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
466- 479+ 
467 fe::PlatFormInfos platform_info;480 fe::PlatFormInfos platform_info;
468 platform_info.Init();481 platform_info.Init();
469-
470 struct EmbeddingBagCompileInfo {482 struct EmbeddingBagCompileInfo {
471 int32_t totalCoreNum = 0;483 int32_t totalCoreNum = 0;
472 int64_t sysWorkspaceSize = 0;484 int64_t sysWorkspaceSize = 0;
473 int64_t ubSizePlatForm = 0;485 int64_t ubSizePlatForm = 0;
474 bool isRegBase = true;486 bool isRegBase = true;
475 } compile_info;487 } compile_info;
476- 488+ 
477 auto kernel_holder =489 auto kernel_holder =
478 gert::KernelRunContextFaker()490 gert::KernelRunContextFaker()
479 .KernelIONum(2, 1)491 .KernelIONum(2, 1)
@@ -486,6 +498,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_301)
486 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");498 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
487 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(499 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
488 "AICoreintrinsicDtypeMap", intrinsics);500 "AICoreintrinsicDtypeMap", intrinsics);
501+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
489 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);502 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
490 503 
491 // tilingFunc simulate504 // tilingFunc simulate
@@ -493,15 +506,15 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_301)
493 ASSERT_NE(param, nullptr);506 ASSERT_NE(param, nullptr);
494 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);507 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
495 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());508 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
496- 509+ 
497 // embedding_dim=64, vocab_size=67108864 (64 * 67108864 = 4294967296 > UINT32_MAX=4294967295)510 // embedding_dim=64, vocab_size=67108864 (64 * 67108864 = 4294967296 > UINT32_MAX=4294967295)
498- gert::StorageShape input_0 = {{67108864, 64}, {67108864, 64}}; 511+ gert::StorageShape input_0 = {{67108864, 64}, {67108864, 64}};
499- gert::StorageShape input_1 = {{512}, {512}}; 512+ gert::StorageShape input_1 = {{512}, {512}};
500- gert::StorageShape input_2 = {{32}, {32}}; 513+ gert::StorageShape input_2 = {{32}, {32}};
501- gert::StorageShape output_shape0 = {{32, 64}, {32, 64}}; 514+ gert::StorageShape output_shape0 = {{32, 64}, {32, 64}};
502- gert::StorageShape output_shape1 = {{512}, {512}}; 515+ gert::StorageShape output_shape1 = {{512}, {512}};
503- gert::StorageShape output_shape2 = {{32}, {32}}; 516+ gert::StorageShape output_shape2 = {{32}, {32}};
504- gert::StorageShape output_shape3 = {{32, 64}, {32, 64}}; 517+ gert::StorageShape output_shape3 = {{32, 64}, {32, 64}};
505 518 
506 auto holder = gert::TilingContextFaker()519 auto holder = gert::TilingContextFaker()
507 .NodeIoNum(3, 4)520 .NodeIoNum(3, 4)
@@ -518,7 +531,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_301)
518 .PlatformInfo(reinterpret_cast<char*>(&platform_info))531 .PlatformInfo(reinterpret_cast<char*>(&platform_info))
519 .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // weight: float532 .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // weight: float
520 .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // indices: int32533 .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // indices: int32
521- .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offsets: int32 534+ .NodeInputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offsets: int32
522 .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // y: float535 .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // y: float
523 .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offset2bag: int32536 .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offset2bag: int32
524 .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // bag_size: int32537 .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // bag_size: int32
@@ -534,9 +547,10 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_301)
534 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");547 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
535 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);548 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
536 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);549 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
550+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
537 551 
538 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);552 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
539- 553+ 
540 auto tiling_key = tiling_context->GetTilingKey();554 auto tiling_key = tiling_context->GetTilingKey();
541 ASSERT_EQ(tiling_key, 301);555 ASSERT_EQ(tiling_key, 301);
542}556}
@@ -561,19 +575,17 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_302)
561 map<string, string> aicore_spec;575 map<string, string> aicore_spec;
562 map<string, string> intrinsics;576 map<string, string> intrinsics;
563 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};577 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};
578+ map<string, string> npuarchs = {{"NpuArch", "3510"}};
564 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);579 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
565- 
566-
567 fe::PlatFormInfos platform_info;580 fe::PlatFormInfos platform_info;
568 platform_info.Init();581 platform_info.Init();
569-
570 struct EmbeddingBagCompileInfo {582 struct EmbeddingBagCompileInfo {
571 int32_t totalCoreNum = 0;583 int32_t totalCoreNum = 0;
572 int64_t sysWorkspaceSize = 0;584 int64_t sysWorkspaceSize = 0;
573 int64_t ubSizePlatForm = 0;585 int64_t ubSizePlatForm = 0;
574 bool isRegBase = true;586 bool isRegBase = true;
575 } compile_info;587 } compile_info;
576- 588+ 
577 auto kernel_holder =589 auto kernel_holder =
578 gert::KernelRunContextFaker()590 gert::KernelRunContextFaker()
579 .KernelIONum(2, 1)591 .KernelIONum(2, 1)
@@ -586,6 +598,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_302)
586 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");598 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
587 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(599 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
588 "AICoreintrinsicDtypeMap", intrinsics);600 "AICoreintrinsicDtypeMap", intrinsics);
601+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
589 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);602 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
590 603 
591 // tilingFunc simulate604 // tilingFunc simulate
@@ -593,8 +606,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_302)
593 ASSERT_NE(param, nullptr);606 ASSERT_NE(param, nullptr);
594 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);607 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
595 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());608 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
596- 609+ 
597-
598 gert::StorageShape input_0 = {{1024, 64}, {1024, 64}}; // 小vocab_size610 gert::StorageShape input_0 = {{1024, 64}, {1024, 64}}; // 小vocab_size
599 gert::StorageShape input_1 = {{512}, {512}}; // indices611 gert::StorageShape input_1 = {{512}, {512}}; // indices
600 gert::StorageShape input_2 = {{32}, {32}}; // offsets612 gert::StorageShape input_2 = {{32}, {32}}; // offsets
@@ -618,7 +630,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_302)
618 .PlatformInfo(reinterpret_cast<char*>(&platform_info))630 .PlatformInfo(reinterpret_cast<char*>(&platform_info))
619 .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // weight: float631 .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // weight: float
620 .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // indices: int32632 .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // indices: int32
621- .NodeInputTd(2, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND) // offsets: int64 633+ .NodeInputTd(2, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND) // offsets: int64
622 .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // y: float634 .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) // y: float
623 .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offset2bag: int32635 .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // offset2bag: int32
624 .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // bag_size: int32636 .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) // bag_size: int32
@@ -634,6 +646,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_302)
634 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");646 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
635 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);647 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
636 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);648 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
649+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
637 650 
638 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);651 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
639 652 
@@ -641,7 +654,6 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_302)
641 ASSERT_EQ(tiling_key, 302);654 ASSERT_EQ(tiling_key, 302);
642}655}
643 656 
644- 
645TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_303)657TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_303)
646{658{
647 std::string op_type("EmbeddingBag");659 std::string op_type("EmbeddingBag");
@@ -662,18 +674,19 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_303)
662 map<string, string> aicore_spec;674 map<string, string> aicore_spec;
663 map<string, string> intrinsics;675 map<string, string> intrinsics;
664 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};676 std::map<std::string, std::string> soc_version_infos = {{"Short_SoC_version", "Ascend950"}};
677+ map<string, string> npuarchs = {{"NpuArch", "3510"}};
665 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);678 GetPlatFormInfos(compile_info_string.c_str(), soc_infos, aicore_spec, intrinsics);
666 679 
667 fe::PlatFormInfos platform_info;680 fe::PlatFormInfos platform_info;
668 platform_info.Init();681 platform_info.Init();
669- 682+ 
670 struct EmbeddingBagCompileInfo {683 struct EmbeddingBagCompileInfo {
671 int32_t totalCoreNum = 0;684 int32_t totalCoreNum = 0;
672 int64_t sysWorkspaceSize = 0;685 int64_t sysWorkspaceSize = 0;
673 int64_t ubSizePlatForm = 0;686 int64_t ubSizePlatForm = 0;
674 bool isRegBase = true;687 bool isRegBase = true;
675 } compile_info;688 } compile_info;
676- 689+ 
677 auto kernel_holder =690 auto kernel_holder =
678 gert::KernelRunContextFaker()691 gert::KernelRunContextFaker()
679 .KernelIONum(2, 1)692 .KernelIONum(2, 1)
@@ -686,6 +699,7 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_303)
686 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");699 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
687 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(700 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes(
688 "AICoreintrinsicDtypeMap", intrinsics);701 "AICoreintrinsicDtypeMap", intrinsics);
702+ kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
689 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);703 ASSERT_EQ(tiling_parse_func(kernel_holder.GetContext<gert::KernelContext>()), ge::GRAPH_SUCCESS);
690 704 
691 // tilingFunc simulate705 // tilingFunc simulate
@@ -693,11 +707,10 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_303)
693 ASSERT_NE(param, nullptr);707 ASSERT_NE(param, nullptr);
694 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);708 auto workspace_size_holer = gert::ContinuousVector::Create<size_t>(4096);
695 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());709 auto ws_size = reinterpret_cast<gert::ContinuousVector*>(workspace_size_holer.get());
696- 710+ 
697- 711+ gert::StorageShape input_0 = {{67108864, 64}, {67108864, 64}};
698- gert::StorageShape input_0 = {{67108864, 64}, {67108864, 64}}; 712+ gert::StorageShape input_1 = {{512}, {512}};
699- gert::StorageShape input_1 = {{512}, {512}}; 713+ gert::StorageShape input_2 = {{32}, {32}};
700- gert::StorageShape input_2 = {{32}, {32}};
701 gert::StorageShape output_shape0 = {{32, 64}, {32, 64}};714 gert::StorageShape output_shape0 = {{32, 64}, {32, 64}};
702 gert::StorageShape output_shape1 = {{512}, {512}};715 gert::StorageShape output_shape1 = {{512}, {512}};
703 gert::StorageShape output_shape2 = {{32}, {32}};716 gert::StorageShape output_shape2 = {{32}, {32}};
@@ -716,13 +729,13 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_303)
716 {"include_last_offset", Ops::NN::AnyValue::CreateFrom<bool>(false)},729 {"include_last_offset", Ops::NN::AnyValue::CreateFrom<bool>(false)},
717 {"padding_idx", Ops::NN::AnyValue::CreateFrom<int64_t>(-1)}})730 {"padding_idx", Ops::NN::AnyValue::CreateFrom<int64_t>(-1)}})
718 .PlatformInfo(reinterpret_cast<char*>(&platform_info))731 .PlatformInfo(reinterpret_cast<char*>(&platform_info))
719- .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) 732+ .NodeInputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
720- .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) 733+ .NodeInputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
721- .NodeInputTd(2, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND) 734+ .NodeInputTd(2, ge::DT_INT64, ge::FORMAT_ND, ge::FORMAT_ND)
722- .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND) 735+ .NodeOutputTd(0, ge::DT_FLOAT, ge::FORMAT_ND, ge::FORMAT_ND)
723- .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) 736+ .NodeOutputTd(1, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
724- .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) 737+ .NodeOutputTd(2, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
725- .NodeOutputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND) 738+ .NodeOutputTd(3, ge::DT_INT32, ge::FORMAT_ND, ge::FORMAT_ND)
726 .TilingData(param.get())739 .TilingData(param.get())
727 .Workspace(ws_size)740 .Workspace(ws_size)
728 .Build();741 .Build();
@@ -734,10 +747,10 @@ TEST_F(EmbeddingBagRegbaseTiling, embedding_bag_regbase_tiling_simt_303)
734 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");747 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetCoreNumByCoreType("AICore");
735 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);748 holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("AICoreintrinsicDtypeMap", intrinsics);
736 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);749 kernel_holder.GetContext<gert::TilingParseContext>()->GetPlatformInfo()->SetPlatformRes("version", soc_version_infos);
750+ holder.GetContext<gert::TilingContext>()->GetPlatformInfo()->SetPlatformRes("version", npuarchs);
737 751 
738 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);752 EXPECT_EQ(tiling_func(tiling_context), ge::GRAPH_SUCCESS);
739- 753+ 
740-
741 auto tiling_key = tiling_context->GetTilingKey();754 auto tiling_key = tiling_context->GetTilingKey();
742 ASSERT_EQ(tiling_key, 303);755 ASSERT_EQ(tiling_key, 303);
743}756}