perf(logit_grad): logitgrad算子A5性能优化-ops-nn-AtomGit

已合并

perf(logit_grad): logitgrad算子A5性能优化 #8743

luoyufan7创建于 21 天前从luoyufan7/ops-nn： pr/logitgrad-final-v2 合入到 cann/ops-nn： master

perf(logit_grad): logitgrad算子A5性能优化 #8743

已合并

luoyufan7创建于 21 天前

从luoyufan7/ops-nn pr/logitgrad-final-v2合入到cann/ops-nn master

共 1 个文件变更+143-2

Mloss/logit_grad/op_kernel/logit_grad.h+143-2

	@@ -28,6 +28,15 @@ constexpr int64_t PP_ELEMENT_NUM = 8 * 1024;
28	constexpr int64_t ONE_REPEAT_ELE_NUM_FP32 = 64;	28	constexpr int64_t ONE_REPEAT_ELE_NUM_FP32 = 64;
29	constexpr int64_t ALIGN = 16;	29	constexpr int64_t ALIGN = 16;
30		30
		31	+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
		Q Q qianzehong21 天前 950判断用 NPU_ARCH == 3510
		32	+constexpr AscendC::MicroAPI::CastTrait G5_FP16_TO_FP32_CAST_TRAIT = {
		33	+ AscendC::MicroAPI::RegLayout::ZERO,
		34	+ AscendC::MicroAPI::SatMode::UNKNOWN,
		35	+ AscendC::MicroAPI::MaskMergeMode::ZEROING,
		36	+ AscendC::RoundMode::UNKNOWN,
		37	+};
		38	+#endif
		39	+
31	template <typename T>	40	template <typename T>
32	class LogitGradND {	41	class LogitGradND {
33	public:	42	public:
	@@ -41,6 +50,10 @@ private:
41	__aicore__ inline void CopyInAndCast(int64_t inputOffset, int64_t dataCount);	50	__aicore__ inline void CopyInAndCast(int64_t inputOffset, int64_t dataCount);
42	__aicore__ inline void ComputeStepOne(int64_t dataCount);	51	__aicore__ inline void ComputeStepOne(int64_t dataCount);
43	__aicore__ inline void ComputeStepTwo(int64_t dataCount);	52	__aicore__ inline void ComputeStepTwo(int64_t dataCount);
		53	+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
		54	+ __aicore__ inline void ComputeFusedFp16(int64_t dataCount);
		55	+ __aicore__ inline void ComputeFusedBf16(int64_t dataCount);
		56	+#endif
44	__aicore__ inline void CastAndCopyOut(int64_t outputOffset, int64_t dataCount);	57	__aicore__ inline void CastAndCopyOut(int64_t outputOffset, int64_t dataCount);
45		58
46	private:	59	private:
	@@ -143,9 +156,19 @@ __aicore__ inline void LogitGradND<T>::Process()
143	eventId = pingPongFlag ? EVENT_ID1 : EVENT_ID0;	156	eventId = pingPongFlag ? EVENT_ID1 : EVENT_ID0;
144	CopyInAndCast(eachCoreStartOffset + localOffset, calNum);	157	CopyInAndCast(eachCoreStartOffset + localOffset, calNum);
145		158
		159	+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
		160	+ if constexpr (std::is_same_v<T, half>) {
		161	+ ComputeFusedFp16(calNum);
		162	+ } else if constexpr (std::is_same_v<T, bfloat16_t>) {
		163	+ ComputeFusedBf16(calNum);
		164	+ } else {
		165	+ ComputeStepOne(calNum);
		166	+ ComputeStepTwo(calNum);
		167	+ }
		168	+#else
146	ComputeStepOne(calNum);	169	ComputeStepOne(calNum);
147	-
148	ComputeStepTwo(calNum);	170	ComputeStepTwo(calNum);
		171	+#endif
149		172
150	CastAndCopyOut(eachCoreStartOffset + localOffset, calNum);	173	CastAndCopyOut(eachCoreStartOffset + localOffset, calNum);
151		174
	@@ -185,6 +208,11 @@ __aicore__ inline void LogitGradND<T>::CopyInAndCast(int64_t inputOffset, int64_
185		208
186	x1TensorFp32 = x1Tensor.template ReinterpretCast<float>();	209	x1TensorFp32 = x1Tensor.template ReinterpretCast<float>();
187	x2TensorFp32 = x2Tensor.template ReinterpretCast<float>();	210	x2TensorFp32 = x2Tensor.template ReinterpretCast<float>();
		211	+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
		212	+ if constexpr (std::is_same_v<T, half>) {
		213	+ return;
		214	+ }
		215	+#endif
188	if (std::is_same_v<T, bfloat16_t> \|\| std::is_same_v<T, half>) {	216	if (std::is_same_v<T, bfloat16_t> \|\| std::is_same_v<T, half>) {
189	Cast(x1TensorFp32, x1Tmp, RoundMode::CAST_NONE, dataCount);	217	Cast(x1TensorFp32, x1Tmp, RoundMode::CAST_NONE, dataCount);
190	PipeBarrier<PIPE_V>();	218	PipeBarrier<PIPE_V>();
	@@ -212,6 +240,119 @@ __aicore__ inline void LogitGradND<T>::ComputeStepOne(int64_t dataCount)
212	PipeBarrier<PIPE_V>();	240	PipeBarrier<PIPE_V>();
213	}	241	}
214		242
		243	+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
		244	+template <typename T>
		245	+__aicore__ inline void LogitGradND<T>::ComputeFusedFp16(int64_t dataCount)
		246	+{
		247	+ float lo = epslion;
		248	+ float hi = static_cast<float>(1.0) - epslion;
		249	+
		250	+ __VEC_SCOPE__
		251	+ {
		252	+ AscendC::MicroAPI::RegTensor<half> regXHalf;
		253	+ AscendC::MicroAPI::RegTensor<half> regDyHalf;
		254	+ AscendC::MicroAPI::RegTensor<float> regX;
		255	+ AscendC::MicroAPI::RegTensor<float> regDy;
		256	+ AscendC::MicroAPI::RegTensor<float> regTmp;
		257	+ AscendC::MicroAPI::RegTensor<float> regOut;
		258	+ AscendC::MicroAPI::RegTensor<float> regLo;
		259	+ AscendC::MicroAPI::RegTensor<float> regHi;
		260	+ AscendC::MicroAPI::RegTensor<float> regInvalid;
		261	+ AscendC::MicroAPI::MaskReg preg0;
		262	+ AscendC::MicroAPI::MaskReg maskGE;
		263	+ AscendC::MicroAPI::MaskReg maskLE;
		264	+ AscendC::MicroAPI::MaskReg maskValid;
		265	+ constexpr uint32_t vfLen = AscendC::VECTOR_REG_WIDTH / sizeof(float);
		266	+ uint32_t count = static_cast<uint32_t>(dataCount);
		267	+ uint16_t vfLoopNum = static_cast<uint16_t>((count + vfLen - 1) / vfLen);
		268	+ __local_mem__ half* xAddr = (__local_mem__ half*)x1Tmp.GetPhyAddr();
		269	+ __local_mem__ half* dyAddr = (__local_mem__ half*)x2Tmp.GetPhyAddr();
		270	+ __local_mem__ float* outAddr = (__local_mem__ float*)x1TensorFp32.GetPhyAddr();
		271	+
		272	+ AscendC::MicroAPI::Duplicate<float>(regLo, lo);
		273	+ AscendC::MicroAPI::Duplicate<float>(regHi, hi);
		274	+ AscendC::MicroAPI::Duplicate<float>(regInvalid, selectValue);
		275	+
		276	+ for (uint16_t i = 0; i < vfLoopNum; i++) {
		277	+ uint32_t rem = count - static_cast<uint32_t>(i) * vfLen;
		278	+ preg0 = AscendC::MicroAPI::UpdateMask<float>(rem);
		279	+ AscendC::MicroAPI::DataCopy<half, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regXHalf,
		280	+ xAddr + i * vfLen);
		281	+ AscendC::MicroAPI::DataCopy<half, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(regDyHalf,
		Q Q qianzehong21 天前此条代码评论区间+279至+281 使用LoadAlign，DataCopy会逐渐日落
		282	+ dyAddr + i * vfLen);
		283	+ AscendC::MicroAPI::Cast<float, half, G5_FP16_TO_FP32_CAST_TRAIT>(regX, regXHalf, preg0);
		284	+ AscendC::MicroAPI::Cast<float, half, G5_FP16_TO_FP32_CAST_TRAIT>(regDy, regDyHalf, preg0);
		285	+ AscendC::MicroAPI::Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(
		286	+ regTmp, regX, static_cast<float>(-1.0), preg0);
		287	+ AscendC::MicroAPI::Adds<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(
		288	+ regTmp, regTmp, static_cast<float>(1.0), preg0);
		289	+ AscendC::MicroAPI::Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(regTmp, regX, regTmp, preg0);
		290	+ AscendC::MicroAPI::Div<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(regDy, regDy, regTmp, preg0);
		291	+ AscendC::MicroAPI::Compare<float, AscendC::CMPMODE::GE>(maskGE, regX, regLo, preg0);
		292	+ AscendC::MicroAPI::Compare<float, AscendC::CMPMODE::LE>(maskLE, regX, regHi, preg0);
		293	+ AscendC::MicroAPI::MaskAnd(maskValid, maskGE, maskLE, preg0);
		294	+ AscendC::MicroAPI::Select<float>(regOut, regDy, regInvalid, maskValid);
		295	+ AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_NORM_B32>(outAddr + i * vfLen, regOut,
		Q Q qianzehong21 天前使用StoreAlign
		296	+ preg0);
		297	+ }
		298	+ }
		299	+ PipeBarrier<PIPE_V>();
		300	+}
		301	+#endif
		302	+
		303	+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
		304	+template <typename T>
		305	+__aicore__ inline void LogitGradND<T>::ComputeFusedBf16(int64_t dataCount)
		306	+{
		307	+ float lo = epslion;
		308	+ float hi = static_cast<float>(1.0) - epslion;
		309	+
		310	+ __VEC_SCOPE__
		311	+ {
		312	+ AscendC::MicroAPI::RegTensor<float> regX;
		313	+ AscendC::MicroAPI::RegTensor<float> regDy;
		314	+ AscendC::MicroAPI::RegTensor<float> regTmp;
		315	+ AscendC::MicroAPI::RegTensor<float> regOut;
		316	+ AscendC::MicroAPI::RegTensor<float> regLo;
		317	+ AscendC::MicroAPI::RegTensor<float> regHi;
		318	+ AscendC::MicroAPI::RegTensor<float> regInvalid;
		319	+ AscendC::MicroAPI::MaskReg preg0;
		320	+ AscendC::MicroAPI::MaskReg maskGE;
		321	+ AscendC::MicroAPI::MaskReg maskLE;
		322	+ AscendC::MicroAPI::MaskReg maskValid;
		323	+ constexpr uint32_t vfLen = AscendC::VECTOR_REG_WIDTH / sizeof(float);
		324	+ uint32_t count = static_cast<uint32_t>(dataCount);
		325	+ uint16_t vfLoopNum = static_cast<uint16_t>((count + vfLen - 1) / vfLen);
		326	+ __local_mem__ float* xAddr = (__local_mem__ float*)x1TensorFp32.GetPhyAddr();
		327	+ __local_mem__ float* dyAddr = (__local_mem__ float*)x2TensorFp32.GetPhyAddr();
		328	+
		329	+ AscendC::MicroAPI::Duplicate<float>(regLo, lo);
		330	+ AscendC::MicroAPI::Duplicate<float>(regHi, hi);
		331	+ AscendC::MicroAPI::Duplicate<float>(regInvalid, selectValue);
		332	+
		333	+ for (uint16_t i = 0; i < vfLoopNum; i++) {
		334	+ uint32_t rem = count - static_cast<uint32_t>(i) * vfLen;
		335	+ preg0 = AscendC::MicroAPI::UpdateMask<float>(rem);
		336	+ AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(regX, xAddr + i * vfLen);
		Q Q qianzehong21 天前同上
		337	+ AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(regDy, dyAddr + i * vfLen);
		338	+ AscendC::MicroAPI::Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(
		339	+ regTmp, regX, static_cast<float>(-1.0), preg0);
		340	+ AscendC::MicroAPI::Adds<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(
		341	+ regTmp, regTmp, static_cast<float>(1.0), preg0);
		342	+ AscendC::MicroAPI::Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(regTmp, regX, regTmp, preg0);
		343	+ AscendC::MicroAPI::Div<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(regDy, regDy, regTmp, preg0);
		344	+ AscendC::MicroAPI::Compare<float, AscendC::CMPMODE::GE>(maskGE, regX, regLo, preg0);
		345	+ AscendC::MicroAPI::Compare<float, AscendC::CMPMODE::LE>(maskLE, regX, regHi, preg0);
		346	+ AscendC::MicroAPI::MaskAnd(maskValid, maskGE, maskLE, preg0);
		347	+ AscendC::MicroAPI::Select<float>(regOut, regDy, regInvalid, maskValid);
		348	+ AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_NORM_B32>(xAddr + i * vfLen, regOut,
		Q Q qianzehong21 天前同上
		349	+ preg0);
		350	+ }
		351	+ }
		352	+ PipeBarrier<PIPE_V>();
		353	+}
		354	+#endif
		355	+
215	template <typename T>	356	template <typename T>
216	__aicore__ inline void LogitGradND<T>::ComputeStepTwo(int64_t dataCount)	357	__aicore__ inline void LogitGradND<T>::ComputeStepTwo(int64_t dataCount)
217	{	358	{
	@@ -252,4 +393,4 @@ __aicore__ inline void LogitGradND<T>::CastAndCopyOut(int64_t outputOffset, int6
252	}	393	}
253		394
254	} // namespace LogitGrad	395	} // namespace LogitGrad
255	-#endif	396	+#endif