Star1242
2035
代码
代码
Issues551
Pull Requests559
Actions
项目讨论
Wiki
分析
Star1242
2035
已合并
batch_norm_v3/rms_norm_grad : 还原 infer/dgamma 计算的内联写法 #8107
rk创建于 7月30日
batch_norm_v3/rms_norm_grad : 还原 infer/dgamma 计算的内联写法 #8107
已合并
rk创建于 7月30日
从wangpengbo26/ops-nncleancode9合入到cann/ops-nnmaster
共 5 个文件变更+100-58
Mnorm/batch_norm_v3/op_kernel/arch35/batch_norm_v3_infer_last_channel_continuous_a.h+19-4
@@ -156,8 +156,25 @@ private:
156 156 
157 __aicore__ inline void Compute(int64_t curTileBLen)157 __aicore__ inline void Compute(int64_t curTileBLen)
158 {158 {
159- InferComputeImpl<BatchNormV3InferLastChannelContinuousA, T>(159+ LocalTensor<T> x = xQueue_.DeQue<T>();
160- *this, xQueue_, yQueue_, betaFp32Buf_, gammaFp32Buf_, meanFp32Buf_, rstdFp32Buf_, curTileBLen);160+ LocalTensor<T> y = yQueue_.AllocTensor<T>();
161+ LocalTensor<float> betaFp32 = betaFp32Buf_.Get<float>();
162+ LocalTensor<float> gammaFp32 = gammaFp32Buf_.Get<float>();
163+ LocalTensor<float> meanFp32 = meanFp32Buf_.Get<float>();
164+ LocalTensor<float> rstdFp32 = rstdFp32Buf_.Get<float>();
165+ 
166+ __local_mem__ T* xLocal = (__local_mem__ T*)x.GetPhyAddr();
167+ __local_mem__ T* yLocal = (__local_mem__ T*)y.GetPhyAddr();
168+ __local_mem__ float* betaFp32Local = (__local_mem__ float*)betaFp32.GetPhyAddr();
169+ __local_mem__ float* gammaFp32Local = (__local_mem__ float*)gammaFp32.GetPhyAddr();
170+ __local_mem__ float* meanFp32Local = (__local_mem__ float*)meanFp32.GetPhyAddr();
171+ __local_mem__ float* rstdFp32Local = (__local_mem__ float*)rstdFp32.GetPhyAddr();
172+ 
173+ VFNormalize(xLocal, gammaFp32Local, betaFp32Local, meanFp32Local, rstdFp32Local, yLocal, curTileBLen);
174+ 
175+ yQueue_.EnQue(y);
176+ 
177+ xQueue_.FreeTensor<T>(x);
161 }178 }
162 179 
163 __aicore__ inline void VFPrepareParamCache(__local_mem__ T_GAMMA* gammaLocal, __local_mem__ T_GAMMA* betaLocal,180 __aicore__ inline void VFPrepareParamCache(__local_mem__ T_GAMMA* gammaLocal, __local_mem__ T_GAMMA* betaLocal,
@@ -195,7 +212,6 @@ private:
195 }212 }
196 }213 }
197 214 
198-public:
199 __aicore__ inline void VFNormalize(__local_mem__ T* xLocal, __local_mem__ float* gammaFp32Local,215 __aicore__ inline void VFNormalize(__local_mem__ T* xLocal, __local_mem__ float* gammaFp32Local,
200 __local_mem__ float* betaFp32Local, __local_mem__ float* meanFp32Local,216 __local_mem__ float* betaFp32Local, __local_mem__ float* meanFp32Local,
201 __local_mem__ float* rstdFp32Local, __local_mem__ T* yLocal, int64_t curTileBLen)217 __local_mem__ float* rstdFp32Local, __local_mem__ T* yLocal, int64_t curTileBLen)
@@ -233,7 +249,6 @@ public:
233 }249 }
234 }250 }
235 251 
236-private:
237 template <typename T_SRC>252 template <typename T_SRC>
238 __aicore__ inline void LoadParamForDtypeT(__local_mem__ T_SRC* src, RegTensor<float>& dst, MaskReg& preg,253 __aicore__ inline void LoadParamForDtypeT(__local_mem__ T_SRC* src, RegTensor<float>& dst, MaskReg& preg,
239 uint32_t offset)254 uint32_t offset)
Mnorm/batch_norm_v3/op_kernel/arch35/batch_norm_v3_infer_last_channel_small_a.h+20-4
@@ -144,11 +144,28 @@ private:
144 144 
145 __aicore__ inline void Compute(int64_t curTileBLen, int64_t curTileALen)145 __aicore__ inline void Compute(int64_t curTileBLen, int64_t curTileALen)
146 {146 {
147- InferComputeImpl<BatchNormV3InferLastChannelSmallA, T>(*this, xQueue_, yQueue_, betaFp32Buf_, gammaFp32Buf_,147+ LocalTensor<T> x = xQueue_.DeQue<T>();
148- meanFp32Buf_, rstdFp32Buf_, curTileBLen * curTileALen);148+ LocalTensor<T> y = yQueue_.AllocTensor<T>();
149+ LocalTensor<float> betaFp32 = betaFp32Buf_.Get<float>();
150+ LocalTensor<float> gammaFp32 = gammaFp32Buf_.Get<float>();
151+ LocalTensor<float> meanFp32 = meanFp32Buf_.Get<float>();
152+ LocalTensor<float> rstdFp32 = rstdFp32Buf_.Get<float>();
153+ 
154+ __local_mem__ T* xLocal = (__local_mem__ T*)x.GetPhyAddr();
155+ __local_mem__ T* yLocal = (__local_mem__ T*)y.GetPhyAddr();
156+ __local_mem__ float* betaFp32Local = (__local_mem__ float*)betaFp32.GetPhyAddr();
157+ __local_mem__ float* gammaFp32Local = (__local_mem__ float*)gammaFp32.GetPhyAddr();
158+ __local_mem__ float* meanFp32Local = (__local_mem__ float*)meanFp32.GetPhyAddr();
159+ __local_mem__ float* rstdFp32Local = (__local_mem__ float*)rstdFp32.GetPhyAddr();
160+ 
161+ VFNormalize(xLocal, gammaFp32Local, betaFp32Local, meanFp32Local, rstdFp32Local, yLocal,
162+ curTileBLen * curTileALen);
163+ 
164+ yQueue_.EnQue(y);
165+ 
166+ xQueue_.FreeTensor<T>(x);
149 }167 }
150 168 
151-public:
152 __aicore__ inline void VFNormalize(__local_mem__ T* xLocal, __local_mem__ float* gammaFp32Local,169 __aicore__ inline void VFNormalize(__local_mem__ T* xLocal, __local_mem__ float* gammaFp32Local,
153 __local_mem__ float* betaFp32Local, __local_mem__ float* meanFp32Local,170 __local_mem__ float* betaFp32Local, __local_mem__ float* meanFp32Local,
154 __local_mem__ float* rstdFp32Local, __local_mem__ T* yLocal, uint32_t curElemLen)171 __local_mem__ float* rstdFp32Local, __local_mem__ T* yLocal, uint32_t curElemLen)
@@ -188,7 +205,6 @@ public:
188 }205 }
189 }206 }
190 207 
191-private:
192 __aicore__ inline void CopyOutY(int64_t yGmOffset, int64_t curTileBLen, int64_t curTileALen)208 __aicore__ inline void CopyOutY(int64_t yGmOffset, int64_t curTileBLen, int64_t curTileALen)
193 {209 {
194 LocalTensor<T> y = yQueue_.DeQue<T>();210 LocalTensor<T> y = yQueue_.DeQue<T>();
Mnorm/batch_norm_v3/op_kernel/arch35/batch_norm_v3_infer_small_ab1.h+19-4
@@ -134,11 +134,27 @@ private:
134 134 
135 __aicore__ inline void Compute(int64_t curTileB0Len)135 __aicore__ inline void Compute(int64_t curTileB0Len)
136 {136 {
137- InferComputeImpl<BatchNormV3InferSmallAB1, T>(*this, xQueue_, yQueue_, betaFp32Buf_, gammaFp32Buf_,137+ LocalTensor<T> x = xQueue_.DeQue<T>();
138- meanFp32Buf_, rstdFp32Buf_, curTileB0Len);138+ LocalTensor<T> y = yQueue_.AllocTensor<T>();
139+ LocalTensor<float> betaFp32 = betaFp32Buf_.Get<float>();
140+ LocalTensor<float> gammaFp32 = gammaFp32Buf_.Get<float>();
141+ LocalTensor<float> meanFp32 = meanFp32Buf_.Get<float>();
142+ LocalTensor<float> rstdFp32 = rstdFp32Buf_.Get<float>();
143+ 
144+ __local_mem__ T* xLocal = (__local_mem__ T*)x.GetPhyAddr();
145+ __local_mem__ T* yLocal = (__local_mem__ T*)y.GetPhyAddr();
146+ __local_mem__ float* betaFp32Local = (__local_mem__ float*)betaFp32.GetPhyAddr();
147+ __local_mem__ float* gammaFp32Local = (__local_mem__ float*)gammaFp32.GetPhyAddr();
148+ __local_mem__ float* meanFp32Local = (__local_mem__ float*)meanFp32.GetPhyAddr();
149+ __local_mem__ float* rstdFp32Local = (__local_mem__ float*)rstdFp32.GetPhyAddr();
150+ 
151+ VFNormalize(xLocal, gammaFp32Local, betaFp32Local, meanFp32Local, rstdFp32Local, yLocal, curTileB0Len);
152+ 
153+ yQueue_.EnQue(y);
154+ 
155+ xQueue_.FreeTensor<T>(x);
139 }156 }
140 157 
141-public:
142 __aicore__ inline void VFNormalize(__local_mem__ T* xLocal, __local_mem__ float* gammaFp32Local,158 __aicore__ inline void VFNormalize(__local_mem__ T* xLocal, __local_mem__ float* gammaFp32Local,
143 __local_mem__ float* betaFp32Local, __local_mem__ float* meanFp32Local,159 __local_mem__ float* betaFp32Local, __local_mem__ float* meanFp32Local,
144 __local_mem__ float* rstdFp32Local, __local_mem__ T* yLocal,160 __local_mem__ float* rstdFp32Local, __local_mem__ T* yLocal,
@@ -180,7 +196,6 @@ public:
180 }196 }
181 }197 }
182 198 
183-private:
184 __aicore__ inline uint32_t GetSmallAB1ParamCacheElemLen() const199 __aicore__ inline uint32_t GetSmallAB1ParamCacheElemLen() const
185 {200 {
186 uint32_t abLen = static_cast<uint32_t>(tilingData_->totalALen * tilingData_->totalB1Len);201 uint32_t abLen = static_cast<uint32_t>(tilingData_->totalALen * tilingData_->totalB1Len);
Mnorm/batch_norm_v3/op_kernel/arch35/batch_norm_v3_regbase_common.h+0-26
@@ -59,32 +59,6 @@ struct RLessThanParams {
59 uint32_t remainderTailOffset3;59 uint32_t remainderTailOffset3;
60};60};
61 61 
62-template <typename Self, typename T>
63-__aicore__ inline void InferComputeImpl(Self& self, TQue<QuePosition::VECIN, 1>& xQueue,
64- TQue<QuePosition::VECOUT, 1>& yQueue, TBuf<TPosition::VECCALC>& betaBuf,
65- TBuf<TPosition::VECCALC>& gammaBuf, TBuf<TPosition::VECCALC>& meanBuf,
66- TBuf<TPosition::VECCALC>& rstdBuf, int64_t vfLen)
67-{
68- LocalTensor<T> x = xQueue.template DeQue<T>();
69- LocalTensor<T> y = yQueue.template AllocTensor<T>();
70- LocalTensor<float> betaFp32 = betaBuf.template Get<float>();
71- LocalTensor<float> gammaFp32 = gammaBuf.template Get<float>();
72- LocalTensor<float> meanFp32 = meanBuf.template Get<float>();
73- LocalTensor<float> rstdFp32 = rstdBuf.template Get<float>();
74- 
75- __local_mem__ T* xLocal = (__local_mem__ T*)x.GetPhyAddr();
76- __local_mem__ T* yLocal = (__local_mem__ T*)y.GetPhyAddr();
77- __local_mem__ float* betaFp32Local = (__local_mem__ float*)betaFp32.GetPhyAddr();
78- __local_mem__ float* gammaFp32Local = (__local_mem__ float*)gammaFp32.GetPhyAddr();
79- __local_mem__ float* meanFp32Local = (__local_mem__ float*)meanFp32.GetPhyAddr();
80- __local_mem__ float* rstdFp32Local = (__local_mem__ float*)rstdFp32.GetPhyAddr();
81- 
82- self.VFNormalize(xLocal, gammaFp32Local, betaFp32Local, meanFp32Local, rstdFp32Local, yLocal, vfLen);
83- 
84- yQueue.EnQue(y);
85- xQueue.template FreeTensor<T>(x);
86-}
87- 
88__aicore__ inline RLessThanParams GetRLessThanParams(uint32_t scaleCoef, uint32_t currentANumAlign, uint32_t r1)62__aicore__ inline RLessThanParams GetRLessThanParams(uint32_t scaleCoef, uint32_t currentANumAlign, uint32_t r1)
89{63{
90 RLessThanParams params;64 RLessThanParams params;
Mnorm/rms_norm_grad/op_kernel/arch35/rms_norm_grad_regbase_dgamma.h+42-20
@@ -226,10 +226,14 @@ public:
226 }226 }
227 }227 }
228 228 
229- __aicore__ inline void ComputePreDgamma(LocalTensor<DY_TYPE>& dyLocal, LocalTensor<X_TYPE>& xLocal,229+ __aicore__ inline void CalcDgamma(uint32_t inputOffset, uint32_t currentCols, bool isWithPad)
230- LocalTensor<RSTD_TYPE>& rstdLocal, LocalTensor<float>& dgammaOutLocal,
231- uint32_t currentCols, __local_mem__ float*& dgammaOutAddr)
232 {230 {
231+ LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();
232+ LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();
233+ LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();
234+ LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();
235+ 
236+ CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, 0);
233 xQueue_.EnQue(xLocal);237 xQueue_.EnQue(xLocal);
234 rstdQueue_.EnQue(rstdLocal);238 rstdQueue_.EnQue(rstdLocal);
235 dyQueue_.EnQue(dyLocal);239 dyQueue_.EnQue(dyLocal);
@@ -241,24 +245,12 @@ public:
241 __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();245 __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();
242 __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();246 __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();
243 __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();247 __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
244- dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();248+ __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();
245 249 
246 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);250 VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
247 dyQueue_.FreeTensor(dyLocal);251 dyQueue_.FreeTensor(dyLocal);
248 xQueue_.FreeTensor(xLocal);252 xQueue_.FreeTensor(xLocal);
249 rstdQueue_.FreeTensor(rstdLocal);253 rstdQueue_.FreeTensor(rstdLocal);
250- }
251- 
252- __aicore__ inline void CalcDgamma(uint32_t inputOffset, uint32_t currentCols, bool isWithPad)
253- {
254- LocalTensor<RSTD_TYPE> rstdLocal = rstdQueue_.template AllocTensor<RSTD_TYPE>();
255- LocalTensor<DY_TYPE> dyLocal = dyQueue_.template AllocTensor<DY_TYPE>();
256- LocalTensor<X_TYPE> xLocal = xQueue_.template AllocTensor<X_TYPE>();
257- LocalTensor<float> dgammaOutLocal = dgammaQueue_.template AllocTensor<float>();
258- 
259- CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, 0);
260- __local_mem__ float* dgammaOutAddr;
261- ComputePreDgamma(dyLocal, xLocal, rstdLocal, dgammaOutLocal, currentCols, dgammaOutAddr);
262 if (isWithPad) {254 if (isWithPad) {
263 uint32_t mainRows = rows_ - tiling_->rowsTailDG;255 uint32_t mainRows = rows_ - tiling_->rowsTailDG;
264 VFDuplicateRows(dgammaOutAddr, vlFp32_, rows_ * vlFp32_);256 VFDuplicateRows(dgammaOutAddr, vlFp32_, rows_ * vlFp32_);
@@ -287,8 +279,23 @@ public:
287 uint32_t rstdOffset = i * rowsPerUB_;279 uint32_t rstdOffset = i * rowsPerUB_;
288 280 
289 CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);281 CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);
290- __local_mem__ float* dgammaOutAddr;282+ xQueue_.EnQue(xLocal);
291- ComputePreDgamma(dyLocal, xLocal, rstdLocal, dgammaOutLocal, currentCols, dgammaOutAddr);283+ rstdQueue_.EnQue(rstdLocal);
284+ dyQueue_.EnQue(dyLocal);
285+ 
286+ dyLocal = dyQueue_.template DeQue<DY_TYPE>();
287+ xLocal = xQueue_.template DeQue<X_TYPE>();
288+ rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
289+ 
290+ __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();
291+ __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();
292+ __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
293+ __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();
294+ 
295+ VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
296+ dyQueue_.FreeTensor(dyLocal);
297+ xQueue_.FreeTensor(xLocal);
298+ rstdQueue_.FreeTensor(rstdLocal);
292 299 
293 VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rowsPerUB_);300 VFBinaryReduceSumWithoutTail(dgammaOutAddr, currentCols, rowsPerUB_);
294 UpdateCache(binaryAddCacheLocal, dgammaOutAddr, cacheID, vlFp32_);301 UpdateCache(binaryAddCacheLocal, dgammaOutAddr, cacheID, vlFp32_);
@@ -307,8 +314,23 @@ public:
307 int64_t cacheID = GetCacheID(i);314 int64_t cacheID = GetCacheID(i);
308 315 
309 CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);316 CopyInputsToUB(dyLocal, xLocal, rstdLocal, inputOffset, currentCols, rowsPerUB_, rstdOffset);
310- __local_mem__ float* dgammaOutAddr;317+ xQueue_.EnQue(xLocal);
311- ComputePreDgamma(dyLocal, xLocal, rstdLocal, dgammaOutLocal, currentCols, dgammaOutAddr);318+ rstdQueue_.EnQue(rstdLocal);
319+ dyQueue_.EnQue(dyLocal);
320+ 
321+ dyLocal = dyQueue_.template DeQue<DY_TYPE>();
322+ xLocal = xQueue_.template DeQue<X_TYPE>();
323+ rstdLocal = rstdQueue_.template DeQue<RSTD_TYPE>();
324+ 
325+ __local_mem__ DY_TYPE* dyAddr = (__local_mem__ DY_TYPE*)dyLocal[0].GetPhyAddr();
326+ __local_mem__ X_TYPE* xAddr = (__local_mem__ X_TYPE*)xLocal[0].GetPhyAddr();
327+ __local_mem__ RSTD_TYPE* rstdAddr = (__local_mem__ RSTD_TYPE*)rstdLocal[0].GetPhyAddr();
328+ __local_mem__ float* dgammaOutAddr = (__local_mem__ float*)dgammaOutLocal[0].GetPhyAddr();
329+ 
330+ VFCalcPreDgamma(dyAddr, xAddr, rstdAddr, dgammaOutAddr, currentCols, rowsPerUB_);
331+ dyQueue_.FreeTensor(dyLocal);
332+ xQueue_.FreeTensor(xLocal);
333+ rstdQueue_.FreeTensor(rstdLocal);
312 334 
313 // 处理累加的尾块335 // 处理累加的尾块
314 LocalTensor<RSTD_TYPE> rstdLocal1 = rstdQueue_.template AllocTensor<RSTD_TYPE>();336 LocalTensor<RSTD_TYPE> rstdLocal1 = rstdQueue_.template AllocTensor<RSTD_TYPE>();