Star915
1833
代码
代码
Issues230
Pull Requests324
Actions
项目讨论
Wiki
分析
Star915
1833
已合并
embadding bag dispose core dump #1655
dong-yanrong创建于 2月6日
embadding bag dispose core dump #1655
已合并
dong-yanrong创建于 2月6日
从dong-yanrong/ops-nndev_bag_0206合入到cann/ops-nnmaster
共 4 个文件变更+141-11
Mindex/embedding_bag/op_host/embedding_bag_infershape.cpp+134-9
@@ -166,7 +166,130 @@ inline ge::graphStatus InferShape4MaxIndices(
166 return ge::GRAPH_SUCCESS;166 return ge::GRAPH_SUCCESS;
167}167}
168 168 
169-static ge::graphStatus InferShapeForEmbeddingBagSupport(gert::InferShapeContext* context)169+inline ge::graphStatus InferShape4OutputSupport(
170+ gert::InferShapeContext* context, int64_t batch, int64_t embedding_dim, int64_t dimension,
171+ bool is_unknown_weight, bool is_unknown_indices, bool is_unknown_offset)
172+{
173+ auto output_shape = context->GetOutputShape(Y_IDX);
174+ OP_CHECK_NULL_WITH_CONTEXT(context, output_shape);
175+ output_shape->SetDimNum(OUTPUT_DIMS);
176+
177+ if (dimension == DIM_NUM_ONE) {
178+ if (is_unknown_offset) {
179+ output_shape->SetDim(BATCH_DIM, -1);
180+ } else {
181+ output_shape->SetDim(BATCH_DIM, batch);
182+ }
183+ } else {
184+ if (is_unknown_indices) {
185+ output_shape->SetDim(BATCH_DIM, -1);
186+ } else {
187+ output_shape->SetDim(BATCH_DIM, batch);
188+ }
189+ }
190+ 
191+ if (is_unknown_weight) {
192+ output_shape->SetDim(EMBEDDING_DIM_IDX, -1);
193+ } else {
194+ output_shape->SetDim(EMBEDDING_DIM_IDX, embedding_dim);
195+ }
196+ return ge::GRAPH_SUCCESS;
197+}
198+ 
199+inline ge::graphStatus InferShape4Offset2BagSupport(
200+ gert::InferShapeContext* context, int64_t indices_num, bool is_unknown_indices)
201+{
202+ auto offset_2_bag_shape = context->GetOutputShape(OFFSET_2_BAG_IDX);
203+ OP_CHECK_NULL_WITH_CONTEXT(context, offset_2_bag_shape);
204+ 
205+ offset_2_bag_shape->SetDimNum(INDICES_ONE_DIM);
206+ if (is_unknown_indices) {
207+ offset_2_bag_shape->SetDim(INDICES_ZERO_DIM, -1);
208+ } else {
209+ offset_2_bag_shape->SetDim(INDICES_ZERO_DIM, indices_num);
210+ }
211+ return ge::GRAPH_SUCCESS;
212+}
213+ 
214+inline ge::graphStatus InferShape4BagSizeSupport(
215+ gert::InferShapeContext* context, int64_t offsets_lens, int64_t dimension, bool is_unknown_indices, bool is_unknown_offset)
216+{
217+ auto bag_size_shape = context->GetOutputShape(BAG_SIZE_IDX);
218+ OP_CHECK_NULL_WITH_CONTEXT(context, bag_size_shape);
219+ bag_size_shape->SetDimNum(BAG_SIZE_DIM);
220+ 
221+ if (dimension == DIM_NUM_ONE) {
222+ if (is_unknown_offset) {
223+ bag_size_shape->SetDim(BAG_SIZE_ZERO_DIM, -1);
224+ } else {
225+ bag_size_shape->SetDim(BAG_SIZE_ZERO_DIM, offsets_lens);
226+ }
227+ } else {
228+ if (is_unknown_indices) {
229+ bag_size_shape->SetDim(BAG_SIZE_ZERO_DIM, -1);
230+ } else {
231+ bag_size_shape->SetDim(BAG_SIZE_ZERO_DIM, offsets_lens);
232+ }
233+ }
234+
235+ return ge::GRAPH_SUCCESS;
236+}
237+ 
238+inline ge::graphStatus InferShape4MaxIndicesSupport(
239+ gert::InferShapeContext* context, int64_t batch, int64_t embedding_dim, int64_t dimension,
240+ bool is_unknown_weight, bool is_unknown_indices, bool is_unknown_offset)
241+{
242+ auto max_indices_shape = context->GetOutputShape(MAX_INDICES_IDX);
243+ if (max_indices_shape == nullptr) {
244+ return ge::GRAPH_FAILED;
245+ }
246+ 
247+ auto* attrs = context->GetAttrs();
248+ if (attrs == nullptr) {
249+ return ge::GRAPH_FAILED;
250+ }
251+ const char* mode = attrs->GetAttrPointer<char>(MODE_IDX);
252+ if (mode == nullptr) {
253+ return ge::GRAPH_FAILED;
254+ }
255+ 
256+ if (strcmp(mode, ATTR_MODE_MAX) == 0) {
257+ max_indices_shape->SetDimNum(MAX_INDICES_DIM_TWO);
258+ if (dimension == DIM_NUM_ONE) {
259+ if (is_unknown_offset) {
260+ max_indices_shape->SetDim(MAX_INDICES_ZERO_DIM, -1);
261+ } else {
262+ max_indices_shape->SetDim(MAX_INDICES_ZERO_DIM, batch);
263+ }
264+ } else {
265+ if (is_unknown_indices) {
266+ max_indices_shape->SetDim(MAX_INDICES_ZERO_DIM, -1);
267+ } else {
268+ max_indices_shape->SetDim(MAX_INDICES_ZERO_DIM, batch);
269+ }
270+ }
271+ if (is_unknown_weight) {
272+ max_indices_shape->SetDim(MAX_INDICES_ONE_DIM, -1);
273+ } else {
274+ max_indices_shape->SetDim(MAX_INDICES_ONE_DIM, embedding_dim);
275+ }
276+ } else {
277+ max_indices_shape->SetDimNum(MAX_INDICES_DIM);
278+ max_indices_shape->SetDim(MAX_INDICES_ZERO_DIM, 0);
279+ }
280+ return ge::GRAPH_SUCCESS;
281+}
282+ 
283+static ge::graphStatus CheckIsUnknown(const gert::Shape* check_shape){
284+ bool is_unknown_rank = Ops::Base::IsUnknownRank(*check_shape);
285+ bool is_unknown_shape = Ops::Base::IsUnknownShape(*check_shape);
286+ if (is_unknown_rank || is_unknown_shape) {
287+ return true;
288+ }
289+ return false;
290+}
291+ 
292+static bool InferShapeForEmbeddingBagSupport(gert::InferShapeContext* context)
170{293{
171 OP_LOGD(context->GetNodeName(), "runtime2.0 EmbeddingBagSupport infershape running.");294 OP_LOGD(context->GetNodeName(), "runtime2.0 EmbeddingBagSupport infershape running.");
172 auto const offsets_shape = context->GetInputShape(OFFSETS_IDX);295 auto const offsets_shape = context->GetInputShape(OFFSETS_IDX);
@@ -189,28 +312,30 @@ static ge::graphStatus InferShapeForEmbeddingBagSupport(gert::InferShapeContext*
189 int64_t batch = get_batch(*include_last_offset, offsets_lens);312 int64_t batch = get_batch(*include_last_offset, offsets_lens);
190 int64_t bag_size = offsets_lens;313 int64_t bag_size = offsets_lens;
191 314 
315+ int64_t dimension = DIM_NUM_ONE;
192 if (indices_shape->GetDimNum() >= DIM_NUM_TWO) {316 if (indices_shape->GetDimNum() >= DIM_NUM_TWO) {
193 batch = indices_shape->GetDim(INDICES_ZERO_DIM);317 batch = indices_shape->GetDim(INDICES_ZERO_DIM);
194 indices_num = indices_shape->GetDim(INDICES_ZERO_DIM) * indices_shape->GetDim(INDICES_ONE_DIM);318 indices_num = indices_shape->GetDim(INDICES_ZERO_DIM) * indices_shape->GetDim(INDICES_ONE_DIM);
195 bag_size = batch;319 bag_size = batch;
320+ dimension = DIM_NUM_TWO;
196 }321 }
197- 322+ 
198- int64_t weight_dim_num = weight_shape->GetDimNum();323+ bool is_unknown_weight = CheckIsUnknown(weight_shape);
199- bool is_unknown_rank = Ops::Base::IsUnknownRank(*weight_shape);324+ bool is_unknown_indices = CheckIsUnknown(indices_shape);
200- bool is_unknown_shape = Ops::Base::IsUnknownShape(*weight_shape);325+ bool is_unknown_offsets = CheckIsUnknown(offsets_shape);
201 326 
202 OP_CHECK_IF(327 OP_CHECK_IF(
203- InferShape4Output(context, batch, embedding_dim, weight_dim_num, is_unknown_rank, is_unknown_shape) !=328+ InferShape4OutputSupport(context, batch, embedding_dim, dimension, is_unknown_weight, is_unknown_indices, is_unknown_offsets) !=
204 ge::GRAPH_SUCCESS,329 ge::GRAPH_SUCCESS,
205 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for output."), return GRAPH_FAILED);330 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for output."), return GRAPH_FAILED);
206 OP_CHECK_IF(331 OP_CHECK_IF(
207- InferShape4Offset2Bag(context, indices_num, is_unknown_rank, is_unknown_shape) != ge::GRAPH_SUCCESS,332+ InferShape4Offset2BagSupport(context, indices_num, is_unknown_indices) != ge::GRAPH_SUCCESS,
208 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for offset2bag."), return GRAPH_FAILED);333 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for offset2bag."), return GRAPH_FAILED);
209 OP_CHECK_IF(334 OP_CHECK_IF(
210- InferShape4BagSize(context, bag_size, is_unknown_rank, is_unknown_shape) != ge::GRAPH_SUCCESS,335+ InferShape4BagSizeSupport(context, bag_size, dimension, is_unknown_indices, is_unknown_offsets) != ge::GRAPH_SUCCESS,
211 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for bag_size."), return GRAPH_FAILED);336 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for bag_size."), return GRAPH_FAILED);
212 OP_CHECK_IF(337 OP_CHECK_IF(
213- InferShape4MaxIndices(context, batch, embedding_dim, is_unknown_rank, is_unknown_shape, true) != ge::GRAPH_SUCCESS,338+ InferShape4MaxIndicesSupport(context, batch, embedding_dim, dimension, is_unknown_weight, is_unknown_indices, is_unknown_offsets) != ge::GRAPH_SUCCESS,
214 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for max_indices."), return GRAPH_FAILED);339 OP_LOGE(context, "EmbeddingBagSupport failed to infer shape for max_indices."), return GRAPH_FAILED);
215 OP_LOGD(context->GetNodeName(), "runtime2.0 EmbeddingBagSupport infershape running success.");340 OP_LOGD(context->GetNodeName(), "runtime2.0 EmbeddingBagSupport infershape running success.");
216 341
Mindex/embedding_bag/op_host/embedding_bag_regbase_tiling.cpp+5-2
@@ -149,10 +149,10 @@ ge::graphStatus EmbeddingBagRegBaseTiling::GetShapeAttrsInfo()
149 OP_CHECK_NULL_WITH_CONTEXT(context_, inclueLastOfstPtr);149 OP_CHECK_NULL_WITH_CONTEXT(context_, inclueLastOfstPtr);
150 inclueLastOfst_ = *inclueLastOfstPtr;150 inclueLastOfst_ = *inclueLastOfstPtr;
151 151 
152- auto paddingIdxPtr = attrs->GetAttrPointer<bool>(ATTR_INCLUDE_LAST_OFFSET);152+ auto paddingIdxPtr = attrs->GetAttrPointer<int64_t>(ATTR_PADD_INDEX);
153 OP_CHECK_NULL_WITH_CONTEXT(context_, paddingIdxPtr);153 OP_CHECK_NULL_WITH_CONTEXT(context_, paddingIdxPtr);
154 paddingIdx_ = *paddingIdxPtr;154 paddingIdx_ = *paddingIdxPtr;
155- 155+
156 if (embeddingDim_ * weightTypeSize_ <= MAX_SIMT_EMBDDING_BYTES) {156 if (embeddingDim_ * weightTypeSize_ <= MAX_SIMT_EMBDDING_BYTES) {
157 usedCoreNum_ = totalCoreNum_;157 usedCoreNum_ = totalCoreNum_;
158 isSimt_ = 1;158 isSimt_ = 1;
@@ -164,6 +164,9 @@ ge::graphStatus EmbeddingBagRegBaseTiling::GetShapeAttrsInfo()
164 numBags_ = indiceShape.GetDim(0);164 numBags_ = indiceShape.GetDim(0);
165 indiceSize_ = indiceShape.GetDim(1);165 indiceSize_ = indiceShape.GetDim(1);
166 }166 }
167+ if (paddingIdx_ < 0) {
168+ paddingIdx_ = paddingIdx_ + numEmbeddings_;
169+ }
167 return ge::GRAPH_SUCCESS;170 return ge::GRAPH_SUCCESS;
168}171}
169 172 
Mindex/embedding_bag/op_kernel/arch35/embedding_bag_regbase_1d_mean.h+1-0
@@ -66,6 +66,7 @@ public:
66 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));66 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));
67 pipe_.InitBuffer(this->outQueueOffset2bag_, tiling_.weightRowFactor * sizeof(I));67 pipe_.InitBuffer(this->outQueueOffset2bag_, tiling_.weightRowFactor * sizeof(I));
68 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));68 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));
69+ pipe_.InitBuffer(this->maxIndicesCalcBuf_, (tiling_.weightRowFactor) * sizeof(I));
69 SyncAll();70 SyncAll();
70 }71 }
71 72 
Mindex/embedding_bag/op_kernel/arch35/embedding_bag_regbase_1d_sum.h+1-0
@@ -69,6 +69,7 @@ public:
69 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));69 pipe_.InitBuffer(this->outQueueY_, 1, tiling_.weightDimFactor * sizeof(T));
70 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));70 pipe_.InitBuffer(this->outQueueBagSize_, 1, (tiling_.offsetsFactor + 1) * sizeof(I));
71 pipe_.InitBuffer(this->outQueueOffset2bag_, tiling_.weightRowFactor * sizeof(I));71 pipe_.InitBuffer(this->outQueueOffset2bag_, tiling_.weightRowFactor * sizeof(I));
72+ pipe_.InitBuffer(this->maxIndicesCalcBuf_, (tiling_.weightRowFactor) * sizeof(I));
72 SyncAll();73 SyncAll();
73 }74 }
74 75