repeat_interleave_grad r=0场景修改-ops-nn-AtomGit

已合并

repeat_interleave_grad r=0场景修改 #2938

wuao创建于 3月19日从wow5523/ops-nn_9.0.0： 9.0.0 合入到 cann/ops-nn： 9.0.0

repeat_interleave_grad r=0场景修改 #2938

已合并

wuao创建于 3月19日

从wow5523/ops-nn_9.0.0 9.0.0合入到cann/ops-nn 9.0.0

共 2 个文件变更+54-0

Mindex/repeat_interleave_grad/op_kernel/v35/repeat_interleave_grad_david.h+34-0

	@@ -32,6 +32,7 @@ public:
32	constexpr static uint64_t MAX_NUM_PER_RES = 2 * 1024;	32	constexpr static uint64_t MAX_NUM_PER_RES = 2 * 1024;
33	constexpr static int32_t ELEMENT_ONE_REPEAT_ORI = platform::GetVRegSize() / sizeof(DataT);	33	constexpr static int32_t ELEMENT_ONE_REPEAT_ORI = platform::GetVRegSize() / sizeof(DataT);
34	constexpr static int32_t ELEMENT_ONE_REPEAT_COMPUTE = platform::GetVRegSize() / sizeof(PromoteDataT);	34	constexpr static int32_t ELEMENT_ONE_REPEAT_COMPUTE = platform::GetVRegSize() / sizeof(PromoteDataT);
		35	+ constexpr static float INIT_FLOAT_VALUE = 0.0;
35		36
36	__aicore__ inline RepeatInterleaveGradDavid(TPipe& pipe) : pipe_(pipe)	37	__aicore__ inline RepeatInterleaveGradDavid(TPipe& pipe) : pipe_(pipe)
37	{}	38	{}
	@@ -62,6 +63,8 @@ private:
62		63
63	__aicore__ inline void UpdateCacheAux(const int64_t cacheID, const int64_t stride, const int64_t count);	64	__aicore__ inline void UpdateCacheAux(const int64_t cacheID, const int64_t stride, const int64_t count);
64		65
		66	+ __aicore__ inline void ProcessZeroR(int64_t outputDataOffset, int32_t dimA);
		67	+
65	private:	68	private:
66	const RepeatInterleaveGradDavidTilingData* tiling_;	69	const RepeatInterleaveGradDavidTilingData* tiling_;
67	TPipe& pipe_;	70	TPipe& pipe_;
	@@ -169,6 +172,16 @@ __aicore__ inline void RepeatInterleaveGradDavid<DataT, PromoteDataT, IndexT>::P
169	}	172	}
170	}	173	}
171		174
		175	+template <typename DataT, typename PromoteDataT, typename IndexT>
		176	+__aicore__ inline void RepeatInterleaveGradDavid<DataT, PromoteDataT, IndexT>::ProcessZeroR(
		177	+ int64_t outputDataOffset, int32_t dimA)
		178	+{
		179	+ Duplicate<PromoteDataT>(computeRes_, INIT_FLOAT_VALUE, dimA);
		180	+ CopyOut(yGm_[outputDataOffset], computeRes_, dimA);
		181	+ SetFlag<HardEvent::MTE3_V>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
		182	+ WaitFlag<HardEvent::MTE3_V>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
		183	+}
		184	+
172	template <typename DataT, typename PromoteDataT, typename IndexT>	185	template <typename DataT, typename PromoteDataT, typename IndexT>
173	__aicore__ inline void RepeatInterleaveGradDavid<DataT, PromoteDataT, IndexT>::ProcessRepeatBlock(	186	__aicore__ inline void RepeatInterleaveGradDavid<DataT, PromoteDataT, IndexT>::ProcessRepeatBlock(
174	int32_t repeatFactor, LocalTensor<IndexT>& curRepeat, LocalTensor<IndexT>& preRepeat)	187	int32_t repeatFactor, LocalTensor<IndexT>& curRepeat, LocalTensor<IndexT>& preRepeat)
	@@ -192,6 +205,27 @@ __aicore__ inline void RepeatInterleaveGradDavid<DataT, PromoteDataT, IndexT>::P
192	int64_t pOffset = lastRepeatCumsumVal_ * tiling_->lenN;	205	int64_t pOffset = lastRepeatCumsumVal_ * tiling_->lenN;
193	int64_t rOffset = lastRepeatCount_ * tiling_->lenN;	206	int64_t rOffset = lastRepeatCount_ * tiling_->lenN;
194	r = curRepeat.GetValue(k);	207	r = curRepeat.GetValue(k);
		208	+ if (r == 0) {
		209	+ int32_t nBaseFactor = tiling_->basicBlockSize / BLOCK_SIZE_BYTE * BLOCK_SIZE_BYTE / sizeof(DataT);
		210	+ if (nBaseFactor > MAX_NUM_PER_RES) {
		211	+ nBaseFactor = MAX_NUM_PER_RES;
		212	+ }
		213	+ UbParaUnit newSplitN;
		214	+ __RIGUtil::DoUbSplit(tiling_->lenN, nBaseFactor, newSplitN);
		215	+ int32_t nFactor = newSplitN.ubFactor;
		216	+ int64_t outputDataOffset = 0;
		217	+ int64_t nOffset = 0;
		218	+ for (int32_t n = 0; n < newSplitN.ubCount - 1; n++) { // 重新切分的整块
		219	+ nOffset = n * nFactor;
		220	+ outputDataOffset = mOutputOffset_ + rOffset + nOffset;
		221	+ ProcessZeroR(outputDataOffset, nFactor);
		222	+ }
		223	+ // 重新切分的尾块
		224	+ nOffset = (newSplitN.ubCount - 1) * nFactor;
		225	+ outputDataOffset = mOutputOffset_ + rOffset + nOffset;
		226	+ ProcessZeroR(outputDataOffset, newSplitN.ubTailFactor);
		227	+ continue;
		228	+ }
195		229
196	// r * lenN_ <= tiling_->basicBlockSize -> 正常计算，实际调用的也是AscendC的接口，不作区分	230	// r * lenN_ <= tiling_->basicBlockSize -> 正常计算，实际调用的也是AscendC的接口，不作区分
197	if (r < tiling_->rFactor) {	231	if (r < tiling_->rFactor) {

Mindex/repeat_interleave_grad/op_kernel/v35/repeat_interleave_grad_n_1.h+20-0

	@@ -69,6 +69,7 @@ public:
69	constexpr static int32_t ELEMENT_ONE_REPEAT_ORI = VL_LENGTH_B / sizeof(DataT);	69	constexpr static int32_t ELEMENT_ONE_REPEAT_ORI = VL_LENGTH_B / sizeof(DataT);
70	constexpr static int32_t ELEMENT_ONE_REPEAT_COMPUTE = VL_LENGTH_B / sizeof(PromoteDataT);	70	constexpr static int32_t ELEMENT_ONE_REPEAT_COMPUTE = VL_LENGTH_B / sizeof(PromoteDataT);
71	constexpr static int32_t ELEMENT_ONE_BLOCK_SIZE = BLOCK_SIZE_BYTE / sizeof(PromoteDataT);	71	constexpr static int32_t ELEMENT_ONE_BLOCK_SIZE = BLOCK_SIZE_BYTE / sizeof(PromoteDataT);
		72	+ constexpr static float INIT_FLOAT_VALUE = 0.0;
72		73
73	public:	74	public:
74	__aicore__ inline RepeatInterleaveGradDimNOneDavid(TPipe& pipe) : pipe_(pipe)	75	__aicore__ inline RepeatInterleaveGradDimNOneDavid(TPipe& pipe) : pipe_(pipe)
	@@ -142,6 +143,14 @@ public:
142	}	143	}
143		144
144	private:	145	private:
		146	+ __aicore__ inline void ProcessZeroR(int64_t outputDataOffset, int32_t dimA)
		147	+ {
		148	+ Duplicate<PromoteDataT>(computeRes_, INIT_FLOAT_VALUE, dimA);
		149	+ CopyOut(yGm_[outputDataOffset], computeRes_, dimA);
		150	+ SetFlag<HardEvent::MTE3_V>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
		151	+ WaitFlag<HardEvent::MTE3_V>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_V));
		152	+ }
		153	+
145	__aicore__ inline void ProcessRepeatBlock(int32_t repeatFactor, LocalTensor<IndexT>& repeatTensor)	154	__aicore__ inline void ProcessRepeatBlock(int32_t repeatFactor, LocalTensor<IndexT>& repeatTensor)
146	{	155	{
147	int32_t mCount = 0;	156	int32_t mCount = 0;
	@@ -177,6 +186,17 @@ private:
177	int64_t inputDataOffset = 0;	186	int64_t inputDataOffset = 0;
178	int64_t outputDataOffset = 0;	187	int64_t outputDataOffset = 0;
179	r = repeatTensor.GetValue(k);	188	r = repeatTensor.GetValue(k);
		189	+ if (r == 0) {
		190	+ for (int32_t m = 0; m < mCount - 1; m++) {
		191	+ mOffset = m * mFactor;
		192	+ outputDataOffset = mOutputOffset_ + rOffset + mOffset * tiling_->lenR;
		193	+ ProcessZeroR(outputDataOffset, mFactor);
		194	+ }
		195	+ mOffset = (mCount - 1) * mFactor;
		196	+ outputDataOffset = mOutputOffset_ + rOffset + mOffset * tiling_->lenR;
		197	+ ProcessZeroR(outputDataOffset, mTailFactor);
		198	+ continue;
		199	+ }
180	if (r == 1) {	200	if (r == 1) {
181	for (int32_t m = 0; m < mCount - 1; m++) {	201	for (int32_t m = 0; m < mCount - 1; m++) {
182	mOffset = m * mFactor;	202	mOffset = m * mFactor;