已合并
refactor: rename APIs in ops-nn batch3 (10 ops) #8568
hahaha22创建于 13 天前
refactor: rename APIs in ops-nn batch3 (10 ops) #8568
已合并
共 47 个文件变更+2631-2854
| @@ -67,10 +67,10 @@ private: | |||
| 67 | 67 | ||
| 68 | __aicore__ inline void CopyOutY(const LocalTensor<T>& yOutUb, int64_t ubA, int64_t offset); | 68 | __aicore__ inline void CopyOutY(const LocalTensor<T>& yOutUb, int64_t ubA, int64_t offset); |
| 69 | 69 | ||
| 70 | - __aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T* dst, RegTensor<float>& src, MaskReg& preg, | 70 | + __aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T* dst, RegTensor<float>& src, MaskReg& preg, |
| 71 | uint32_t offset); | 71 | uint32_t offset); |
| 72 | 72 | ||
| 73 | - __aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg, | 73 | + __aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg, |
| 74 | uint32_t offset); | 74 | uint32_t offset); |
| 75 | 75 | ||
| 76 | private: | 76 | private: |
| @@ -196,28 +196,28 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormCompute(const LocalTensor<T>& ds | |||
| 196 | int64_t tailSrcBOffset = floorVLCount * VL_FP32; | 196 | int64_t tailSrcBOffset = floorVLCount * VL_FP32; |
| 197 | int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32; | 197 | int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32; |
| 198 | 198 | ||
| 199 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr(); | 199 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr(); |
| 200 | - __local_mem__ T* foldGradA = (__local_mem__ T*)gradTensor.GetPhyAddr(); | 200 | + __ubuf__ T* foldGradA = (__ubuf__ T*)gradTensor.GetPhyAddr(); |
| 201 | - __local_mem__ T* foldGradB = (__local_mem__ T*)gradTensor.GetPhyAddr() + foldSrcBOffset; | 201 | + __ubuf__ T* foldGradB = (__ubuf__ T*)gradTensor.GetPhyAddr() + foldSrcBOffset; |
| 202 | - __local_mem__ T* tailGradA = (__local_mem__ T*)gradTensor.GetPhyAddr() + tailSrcAOffset; | 202 | + __ubuf__ T* tailGradA = (__ubuf__ T*)gradTensor.GetPhyAddr() + tailSrcAOffset; |
| 203 | - __local_mem__ T* tailGradB = (__local_mem__ T*)gradTensor.GetPhyAddr() + tailSrcBOffset; | 203 | + __ubuf__ T* tailGradB = (__ubuf__ T*)gradTensor.GetPhyAddr() + tailSrcBOffset; |
| 204 | - __local_mem__ T* unFoldGrad = (__local_mem__ T*)gradTensor.GetPhyAddr() + unFoldSrcOffset; | 204 | + __ubuf__ T* unFoldGrad = (__ubuf__ T*)gradTensor.GetPhyAddr() + unFoldSrcOffset; |
| 205 | 205 | ||
| 206 | __VEC_SCOPE__ | 206 | __VEC_SCOPE__ |
| 207 | { | 207 | { |
| 208 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 208 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 209 | - AscendC::MicroAPI::UnalignReg UReg; | 209 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 210 | 210 | ||
| 211 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 211 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 212 | - dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; | 212 | + dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; |
| 213 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { | 213 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { |
| 214 | AscendC::MicroAPI::RegTensor<float> reg0, reg1, reg2; | 214 | AscendC::MicroAPI::RegTensor<float> reg0, reg1, reg2; |
| 215 | LoadTensorForDtypeTIn(foldGradA, reg0, pFull, i * outerLoopStride + j * innerLoopStride); | 215 | LoadTensorForDtypeTIn(foldGradA, reg0, pFull, i * outerLoopStride + j * innerLoopStride); |
| 216 | LoadTensorForDtypeTIn(foldGradB, reg1, pFull, i * outerLoopStride + j * innerLoopStride); | 216 | LoadTensorForDtypeTIn(foldGradB, reg1, pFull, i * outerLoopStride + j * innerLoopStride); |
| 217 | 217 | ||
| 218 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg0, reg0, reg1, pFull); | 218 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg0, reg0, reg1, pFull); |
| 219 | - ReduceSum(reg2, reg0, pFull); | 219 | + Reduce<ReduceType::SUM>(reg2, reg0, pFull); |
| 220 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, reg2, UReg, 1); | 220 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, reg2, UReg, 1); |
| 221 | } | 221 | } |
| 222 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { | 222 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { |
| 223 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); | 223 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); |
| @@ -228,18 +228,18 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormCompute(const LocalTensor<T>& ds | |||
| 228 | LoadTensorForDtypeTIn(tailGradB, reg1, pMask, i * outerLoopStride + j * innerLoopStride); | 228 | LoadTensorForDtypeTIn(tailGradB, reg1, pMask, i * outerLoopStride + j * innerLoopStride); |
| 229 | 229 | ||
| 230 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg1, reg0, reg1, pMask); | 230 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg1, reg0, reg1, pMask); |
| 231 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pMask); | 231 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pMask); |
| 232 | - ReduceSum(reg2, reg0, pFull); | 232 | + Reduce<ReduceType::SUM>(reg2, reg0, pFull); |
| 233 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, reg2, UReg, 1); | 233 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, reg2, UReg, 1); |
| 234 | } | 234 | } |
| 235 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { | 235 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { |
| 236 | AscendC::MicroAPI::RegTensor<float> reg0, reg1; | 236 | AscendC::MicroAPI::RegTensor<float> reg0, reg1; |
| 237 | LoadTensorForDtypeTIn(unFoldGrad, reg0, pFull, i * outerLoopStride + j * innerLoopStride); | 237 | LoadTensorForDtypeTIn(unFoldGrad, reg0, pFull, i * outerLoopStride + j * innerLoopStride); |
| 238 | 238 | ||
| 239 | - ReduceSum(reg1, reg0, pFull); | 239 | + Reduce<ReduceType::SUM>(reg1, reg0, pFull); |
| 240 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, reg1, UReg, 1); | 240 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, reg1, UReg, 1); |
| 241 | } | 241 | } |
| 242 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 242 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 243 | } | 243 | } |
| 244 | } | 244 | } |
| 245 | NormComputePost(dstTensor, gradTensor, xTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); | 245 | NormComputePost(dstTensor, gradTensor, xTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); |
| @@ -265,9 +265,9 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePostWithMul(const LocalTe | |||
| 265 | 265 | ||
| 266 | uint16_t loopTimes = aSize; | 266 | uint16_t loopTimes = aSize; |
| 267 | if (rSize <= VL_FP32) { | 267 | if (rSize <= VL_FP32) { |
| 268 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 268 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 269 | - __local_mem__ T* grad = (__local_mem__ T*)gradTensor.GetPhyAddr(); | 269 | + __ubuf__ T* grad = (__ubuf__ T*)gradTensor.GetPhyAddr(); |
| 270 | - __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr(); | 270 | + __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr(); |
| 271 | 271 | ||
| 272 | __VEC_SCOPE__ | 272 | __VEC_SCOPE__ |
| 273 | { | 273 | { |
| @@ -281,7 +281,7 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePostWithMul(const LocalTe | |||
| 281 | LoadTensorForDtypeTIn(grad, reg0, pMask, i * rAligned); | 281 | LoadTensorForDtypeTIn(grad, reg0, pMask, i * rAligned); |
| 282 | LoadTensorForDtypeTIn(x, reg1, pMask, i * rAligned); | 282 | LoadTensorForDtypeTIn(x, reg1, pMask, i * rAligned); |
| 283 | 283 | ||
| 284 | - ReduceSum(reg2, reg0, pMask); | 284 | + Reduce<ReduceType::SUM>(reg2, reg0, pMask); |
| 285 | Duplicate(reg2, reg2, pFull); | 285 | Duplicate(reg2, reg2, pFull); |
| 286 | 286 | ||
| 287 | Exp(reg1, reg1, pMask); | 287 | Exp(reg1, reg1, pMask); |
| @@ -292,11 +292,11 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePostWithMul(const LocalTe | |||
| 292 | } | 292 | } |
| 293 | } | 293 | } |
| 294 | } else { | 294 | } else { |
| 295 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 295 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 296 | - __local_mem__ T* grad = (__local_mem__ T*)gradTensor.GetPhyAddr(); | 296 | + __ubuf__ T* grad = (__ubuf__ T*)gradTensor.GetPhyAddr(); |
| 297 | - __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr(); | 297 | + __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr(); |
| 298 | - __local_mem__ T* grad_1 = (__local_mem__ T*)gradTensor.GetPhyAddr() + VL_FP32; | 298 | + __ubuf__ T* grad_1 = (__ubuf__ T*)gradTensor.GetPhyAddr() + VL_FP32; |
| 299 | - __local_mem__ T* x_1 = (__local_mem__ T*)xTensor.GetPhyAddr() + VL_FP32; | 299 | + __ubuf__ T* x_1 = (__ubuf__ T*)xTensor.GetPhyAddr() + VL_FP32; |
| 300 | 300 | ||
| 301 | __VEC_SCOPE__ | 301 | __VEC_SCOPE__ |
| 302 | { | 302 | { |
| @@ -312,10 +312,10 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePostWithMul(const LocalTe | |||
| 312 | LoadTensorForDtypeTIn(grad_1, reg0_1, pMask, i * rAligned); | 312 | LoadTensorForDtypeTIn(grad_1, reg0_1, pMask, i * rAligned); |
| 313 | LoadTensorForDtypeTIn(x_1, reg1_1, pMask, i * rAligned); | 313 | LoadTensorForDtypeTIn(x_1, reg1_1, pMask, i * rAligned); |
| 314 | 314 | ||
| 315 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg2_1, reg0, pFull); | 315 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg2_1, reg0, pFull); |
| 316 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2_2, reg0, reg0_1, pMask); | 316 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg2_2, reg0, reg0_1, pMask); |
| 317 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg2_2, pMask); | 317 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg2_2, pMask); |
| 318 | - ReduceSum(reg2, reg0, pFull); | 318 | + Reduce<ReduceType::SUM>(reg2, reg0, pFull); |
| 319 | Duplicate(reg2, reg2, pFull); | 319 | Duplicate(reg2, reg2, pFull); |
| 320 | 320 | ||
| 321 | Exp(regExp, reg1, pFull); | 321 | Exp(regExp, reg1, pFull); |
| @@ -353,10 +353,10 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePost( | |||
| 353 | uint16_t oriRAligned = tl_->rAligned; | 353 | uint16_t oriRAligned = tl_->rAligned; |
| 354 | 354 | ||
| 355 | if (rSize <= VL_FP32) { | 355 | if (rSize <= VL_FP32) { |
| 356 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 356 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 357 | - __local_mem__ T* grad = (__local_mem__ T*)gradTensor.GetPhyAddr(); | 357 | + __ubuf__ T* grad = (__ubuf__ T*)gradTensor.GetPhyAddr(); |
| 358 | - __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr(); | 358 | + __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr(); |
| 359 | - __local_mem__ float* sumTmp = (__local_mem__ float*)binAddTmpTensor.GetPhyAddr(); | 359 | + __ubuf__ float* sumTmp = (__ubuf__ float*)binAddTmpTensor.GetPhyAddr(); |
| 360 | 360 | ||
| 361 | __VEC_SCOPE__ | 361 | __VEC_SCOPE__ |
| 362 | { | 362 | { |
| @@ -367,8 +367,8 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePost( | |||
| 367 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 367 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 368 | AscendC::MicroAPI::MaskReg maskOri; | 368 | AscendC::MicroAPI::MaskReg maskOri; |
| 369 | for (uint16_t i = 0; i < loopTimes; ++i) { | 369 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 370 | - DataCopy(reg0, (__local_mem__ float*)sumTmp + i * static_cast<uint32_t>(stride)); | 370 | + LoadAlign(reg0, (__ubuf__ float*)sumTmp + i * static_cast<uint32_t>(stride)); |
| 371 | - ReduceSum(reg1, reg0, pMask); | 371 | + Reduce<ReduceType::SUM>(reg1, reg0, pMask); |
| 372 | Duplicate(reg2, reg1, pFull); | 372 | Duplicate(reg2, reg1, pFull); |
| 373 | 373 | ||
| 374 | uint32_t sreg0 = static_cast<uint32_t>(oriR); | 374 | uint32_t sreg0 = static_cast<uint32_t>(oriR); |
| @@ -385,12 +385,12 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePost( | |||
| 385 | } | 385 | } |
| 386 | } | 386 | } |
| 387 | } else { | 387 | } else { |
| 388 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 388 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 389 | - __local_mem__ float* sumTmpA = (__local_mem__ float*)binAddTmpTensor.GetPhyAddr(); | 389 | + __ubuf__ float* sumTmpA = (__ubuf__ float*)binAddTmpTensor.GetPhyAddr(); |
| 390 | - __local_mem__ float* sumTmpB = (__local_mem__ float*)binAddTmpTensor.GetPhyAddr() + VL_FP32; | 390 | + __ubuf__ float* sumTmpB = (__ubuf__ float*)binAddTmpTensor.GetPhyAddr() + VL_FP32; |
| 391 | 391 | ||
| 392 | - __local_mem__ T* grad = (__local_mem__ T*)gradTensor.GetPhyAddr(); | 392 | + __ubuf__ T* grad = (__ubuf__ T*)gradTensor.GetPhyAddr(); |
| 393 | - __local_mem__ T* x = (__local_mem__ T*)xTensor.GetPhyAddr(); | 393 | + __ubuf__ T* x = (__ubuf__ T*)xTensor.GetPhyAddr(); |
| 394 | 394 | ||
| 395 | __VEC_SCOPE__ | 395 | __VEC_SCOPE__ |
| 396 | { | 396 | { |
| @@ -401,11 +401,11 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePost( | |||
| 401 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 401 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 402 | AscendC::MicroAPI::MaskReg maskOri; | 402 | AscendC::MicroAPI::MaskReg maskOri; |
| 403 | for (uint16_t i = 0; i < loopTimes; ++i) { | 403 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 404 | - DataCopy(reg0, (__local_mem__ float*)sumTmpA + i * stride); | 404 | + LoadAlign(reg0, (__ubuf__ float*)sumTmpA + i * stride); |
| 405 | - DataCopy(reg1, (__local_mem__ float*)sumTmpB + i * stride); | 405 | + LoadAlign(reg1, (__ubuf__ float*)sumTmpB + i * stride); |
| 406 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg1, reg0, reg1, pMask); | 406 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(reg1, reg0, reg1, pMask); |
| 407 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pMask); | 407 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(reg0, reg1, pMask); |
| 408 | - ReduceSum(reg2, reg0, pFull); | 408 | + Reduce<ReduceType::SUM>(reg2, reg0, pFull); |
| 409 | Duplicate(reg2, reg2, pFull); | 409 | Duplicate(reg2, reg2, pFull); |
| 410 | uint32_t sreg0 = static_cast<uint32_t>(oriR); | 410 | uint32_t sreg0 = static_cast<uint32_t>(oriR); |
| 411 | for (uint16_t j = 0; j < rLoopCount; ++j) { | 411 | for (uint16_t j = 0; j < rLoopCount; ++j) { |
| @@ -426,14 +426,14 @@ __aicore__ inline void LogSoftmaxGradAR<T>::NormComputePost( | |||
| 426 | } | 426 | } |
| 427 | 427 | ||
| 428 | template <typename T> | 428 | template <typename T> |
| 429 | -__aicore__ inline void LogSoftmaxGradAR<T>::LoadTensorForDtypeTIn(__local_mem__ T* src, RegTensor<float>& dst, | 429 | +__aicore__ inline void LogSoftmaxGradAR<T>::LoadTensorForDtypeTIn(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg, |
| 430 | - MaskReg& preg, uint32_t offset) | 430 | + uint32_t offset) |
| 431 | { | 431 | { |
| 432 | if constexpr (IsSameType<T, float>::value) { | 432 | if constexpr (IsSameType<T, float>::value) { |
| 433 | - DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset); | 433 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset); |
| 434 | } else { // fp16、bf16 | 434 | } else { // fp16、bf16 |
| 435 | RegTensor<T> xFp16; | 435 | RegTensor<T> xFp16; |
| 436 | - DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset)); | 436 | + LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset)); |
| 437 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); | 437 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 438 | } | 438 | } |
| 439 | } | 439 | } |
| @@ -452,16 +452,16 @@ __aicore__ inline void LogSoftmaxGradAR<T>::CopyInX(const LocalTensor<T>& xInUb, | |||
| 452 | } | 452 | } |
| 453 | 453 | ||
| 454 | template <typename T> | 454 | template <typename T> |
| 455 | -__aicore__ inline void LogSoftmaxGradAR<T>::StoreTensorForDtypeTOut(__local_mem__ T* dst, | 455 | +__aicore__ inline void LogSoftmaxGradAR<T>::StoreTensorForDtypeTOut(__ubuf__ T* dst, |
| 456 | AscendC::MicroAPI::RegTensor<float>& src, | 456 | AscendC::MicroAPI::RegTensor<float>& src, |
| 457 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset) | 457 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset) |
| 458 | { | 458 | { |
| 459 | if constexpr (IsSameType<T, float>::value) { | 459 | if constexpr (IsSameType<T, float>::value) { |
| 460 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); | 460 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); |
| 461 | } else { | 461 | } else { |
| 462 | AscendC::MicroAPI::RegTensor<T> xFp16; | 462 | AscendC::MicroAPI::RegTensor<T> xFp16; |
| 463 | Cast<T, float, castTraitFp32ToFp16>(xFp16, src, preg); | 463 | Cast<T, float, castTraitFp32ToFp16>(xFp16, src, preg); |
| 464 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); | 464 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); |
| 465 | } | 465 | } |
| 466 | } | 466 | } |
| 467 | 467 | ||
| @@ -44,7 +44,7 @@ public: | |||
| 44 | 44 | ||
| 45 | private: | 45 | private: |
| 46 | __aicore__ inline void CalculateOutVF(const LocalTensor<T>& yLocal, const LocalTensor<T>& xLocal, | 46 | __aicore__ inline void CalculateOutVF(const LocalTensor<T>& yLocal, const LocalTensor<T>& xLocal, |
| 47 | - const LocalTensor<T>& gradLocal, __local_mem__ float*& gradSumPtr, uint32_t a, | 47 | + const LocalTensor<T>& gradLocal, __ubuf__ float*& gradSumPtr, uint32_t a, |
| 48 | uint32_t ubFactor); | 48 | uint32_t ubFactor); |
| 49 | __aicore__ inline void CastVF(const LocalTensor<float>& gradFp32Local, const LocalTensor<T>& gradLocal, uint32_t a, | 49 | __aicore__ inline void CastVF(const LocalTensor<float>& gradFp32Local, const LocalTensor<T>& gradLocal, uint32_t a, |
| 50 | uint32_t ubFactor); | 50 | uint32_t ubFactor); |
| @@ -196,7 +196,7 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::Process() | |||
| 196 | 196 | ||
| 197 | yQueue_.FreeTensor(tmpLocal); | 197 | yQueue_.FreeTensor(tmpLocal); |
| 198 | 198 | ||
| 199 | - __local_mem__ float* gradSumPtr = (__local_mem__ float*)totalSumLocal_.GetPhyAddr(); | 199 | + __ubuf__ float* gradSumPtr = (__ubuf__ float*)totalSumLocal_.GetPhyAddr(); |
| 200 | // step 3. 遍历UB块，计算除法 | 200 | // step 3. 遍历UB块，计算除法 |
| 201 | for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) { | 201 | for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) { |
| 202 | int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; | 202 | int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; |
| @@ -233,12 +233,12 @@ template <typename T> | |||
| 233 | __aicore__ inline void LogSoftmaxGradArRecompute<T>::CalculateOutVF(const LocalTensor<T>& yLocal, | 233 | __aicore__ inline void LogSoftmaxGradArRecompute<T>::CalculateOutVF(const LocalTensor<T>& yLocal, |
| 234 | const LocalTensor<T>& xLocal, | 234 | const LocalTensor<T>& xLocal, |
| 235 | const LocalTensor<T>& gradLocal, | 235 | const LocalTensor<T>& gradLocal, |
| 236 | - __local_mem__ float*& gradSumPtr, uint32_t a, | 236 | + __ubuf__ float*& gradSumPtr, uint32_t a, |
| 237 | uint32_t ubFactor) | 237 | uint32_t ubFactor) |
| 238 | { | 238 | { |
| 239 | - __local_mem__ T* yPtr = (__local_mem__ T*)yLocal.GetPhyAddr(); | 239 | + __ubuf__ T* yPtr = (__ubuf__ T*)yLocal.GetPhyAddr(); |
| 240 | - __local_mem__ T* xPtr = (__local_mem__ T*)xLocal.GetPhyAddr(); | 240 | + __ubuf__ T* xPtr = (__ubuf__ T*)xLocal.GetPhyAddr(); |
| 241 | - __local_mem__ T* gradPtr = (__local_mem__ T*)gradLocal.GetPhyAddr(); | 241 | + __ubuf__ T* gradPtr = (__ubuf__ T*)gradLocal.GetPhyAddr(); |
| 242 | 242 | ||
| 243 | __VEC_SCOPE__ | 243 | __VEC_SCOPE__ |
| 244 | { | 244 | { |
| @@ -250,7 +250,7 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::CalculateOutVF(const LocalT | |||
| 250 | uint32_t width = ubFactor; | 250 | uint32_t width = ubFactor; |
| 251 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); | 251 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); |
| 252 | 252 | ||
| 253 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, gradSumPtr); | 253 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, gradSumPtr); |
| 254 | 254 | ||
| 255 | for (uint16_t j = 0; j < repeatTimes; j++) { | 255 | for (uint16_t j = 0; j < repeatTimes; j++) { |
| 256 | mask = MicroAPI::UpdateMask<float>(width); | 256 | mask = MicroAPI::UpdateMask<float>(width); |
| @@ -259,13 +259,13 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::CalculateOutVF(const LocalT | |||
| 259 | auto yAddr = yPtr + j * VL_FP32; | 259 | auto yAddr = yPtr + j * VL_FP32; |
| 260 | 260 | ||
| 261 | if constexpr (xToFp32_) { | 261 | if constexpr (xToFp32_) { |
| 262 | - MicroAPI::DataCopy<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(xRegFp16, xAddr); | 262 | + MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(xRegFp16, xAddr); |
| 263 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(xRegFp32, xRegFp16, mask); | 263 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(xRegFp32, xRegFp16, mask); |
| 264 | - MicroAPI::DataCopy<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(gradRegFp16, gradAddr); | 264 | + MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(gradRegFp16, gradAddr); |
| 265 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(gradRegFp32, gradRegFp16, mask); | 265 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(gradRegFp32, gradRegFp16, mask); |
| 266 | } else { | 266 | } else { |
| 267 | - MicroAPI::DataCopy(xRegFp32, xAddr); | 267 | + MicroAPI::LoadAlign(xRegFp32, xAddr); |
| 268 | - MicroAPI::DataCopy(gradRegFp32, gradAddr); | 268 | + MicroAPI::LoadAlign(gradRegFp32, gradAddr); |
| 269 | } | 269 | } |
| 270 | 270 | ||
| 271 | MicroAPI::Exp(expReg, xRegFp32, mask); | 271 | MicroAPI::Exp(expReg, xRegFp32, mask); |
| @@ -273,10 +273,10 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::CalculateOutVF(const LocalT | |||
| 273 | MicroAPI::Sub(vreg1, gradRegFp32, vreg0, mask); | 273 | MicroAPI::Sub(vreg1, gradRegFp32, vreg0, mask); |
| 274 | 274 | ||
| 275 | if constexpr (yToFp32_) { | 275 | if constexpr (yToFp32_) { |
| 276 | - MicroAPI::DataCopy(yAddr, vreg1, mask); | 276 | + MicroAPI::StoreAlign(yAddr, vreg1, mask); |
| 277 | } else { | 277 | } else { |
| 278 | MicroAPI::Cast<T, float, castTraitFp32ToFp16>(vreg2, vreg1, mask); | 278 | MicroAPI::Cast<T, float, castTraitFp32ToFp16>(vreg2, vreg1, mask); |
| 279 | - MicroAPI::DataCopy<T, MicroAPI::StoreDist::DIST_PACK_B32>(yAddr, vreg2, mask); | 279 | + MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_PACK_B32>(yAddr, vreg2, mask); |
| 280 | } | 280 | } |
| 281 | } | 281 | } |
| 282 | } | 282 | } |
| @@ -287,8 +287,8 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::CastVF(const LocalTensor<fl | |||
| 287 | const LocalTensor<T>& gradLocal, uint32_t a, | 287 | const LocalTensor<T>& gradLocal, uint32_t a, |
| 288 | uint32_t ubFactor) | 288 | uint32_t ubFactor) |
| 289 | { | 289 | { |
| 290 | - __local_mem__ float* gradFp32Ptr = (__local_mem__ float*)gradFp32Local.GetPhyAddr(); | 290 | + __ubuf__ float* gradFp32Ptr = (__ubuf__ float*)gradFp32Local.GetPhyAddr(); |
| 291 | - __local_mem__ T* gradPtr = (__local_mem__ T*)gradLocal.GetPhyAddr(); | 291 | + __ubuf__ T* gradPtr = (__ubuf__ T*)gradLocal.GetPhyAddr(); |
| 292 | 292 | ||
| 293 | __VEC_SCOPE__ | 293 | __VEC_SCOPE__ |
| 294 | { | 294 | { |
| @@ -305,13 +305,13 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::CastVF(const LocalTensor<fl | |||
| 305 | auto gradFp32Addr = gradFp32Ptr + j * VL_FP32; | 305 | auto gradFp32Addr = gradFp32Ptr + j * VL_FP32; |
| 306 | 306 | ||
| 307 | if constexpr (xToFp32_) { | 307 | if constexpr (xToFp32_) { |
| 308 | - MicroAPI::DataCopy<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, gradAddr); | 308 | + MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, gradAddr); |
| 309 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(vreg1, vreg0, mask); | 309 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(vreg1, vreg0, mask); |
| 310 | } else { | 310 | } else { |
| 311 | - MicroAPI::DataCopy(vreg1, gradAddr); | 311 | + MicroAPI::LoadAlign(vreg1, gradAddr); |
| 312 | } | 312 | } |
| 313 | 313 | ||
| 314 | - MicroAPI::DataCopy(gradFp32Addr, vreg1, mask); | 314 | + MicroAPI::StoreAlign(gradFp32Addr, vreg1, mask); |
| 315 | } | 315 | } |
| 316 | } | 316 | } |
| 317 | } | 317 | } |
| @@ -321,8 +321,8 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::FoldBlockVF(const LocalTens | |||
| 321 | const LocalTensor<T>& grad2Local, uint32_t a, | 321 | const LocalTensor<T>& grad2Local, uint32_t a, |
| 322 | uint32_t ubFactor) | 322 | uint32_t ubFactor) |
| 323 | { | 323 | { |
| 324 | - __local_mem__ float* grad1Fp32Ptr = (__local_mem__ float*)grad1Fp32Local.GetPhyAddr(); | 324 | + __ubuf__ float* grad1Fp32Ptr = (__ubuf__ float*)grad1Fp32Local.GetPhyAddr(); |
| 325 | - __local_mem__ T* grad2Ptr = (__local_mem__ T*)grad2Local.GetPhyAddr(); | 325 | + __ubuf__ T* grad2Ptr = (__ubuf__ T*)grad2Local.GetPhyAddr(); |
| 326 | 326 | ||
| 327 | __VEC_SCOPE__ | 327 | __VEC_SCOPE__ |
| 328 | { | 328 | { |
| @@ -340,18 +340,18 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::FoldBlockVF(const LocalTens | |||
| 340 | auto grad2Addr = grad2Ptr + j * VL_FP32; | 340 | auto grad2Addr = grad2Ptr + j * VL_FP32; |
| 341 | 341 | ||
| 342 | if constexpr (xToFp32_) { | 342 | if constexpr (xToFp32_) { |
| 343 | - MicroAPI::DataCopy<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, grad2Addr); | 343 | + MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, grad2Addr); |
| 344 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(vreg2, vreg0, mask); | 344 | MicroAPI::Cast<float, T, castTraitFp16ToFp32>(vreg2, vreg0, mask); |
| 345 | } else { | 345 | } else { |
| 346 | - MicroAPI::DataCopy(vreg2, grad2Addr); | 346 | + MicroAPI::LoadAlign(vreg2, grad2Addr); |
| 347 | } | 347 | } |
| 348 | 348 | ||
| 349 | - MicroAPI::DataCopy(vreg1, grad1Addr); | 349 | + MicroAPI::LoadAlign(vreg1, grad1Addr); |
| 350 | 350 | ||
| 351 | MicroAPI::Add(vreg3, vreg1, vreg2, mask); | 351 | MicroAPI::Add(vreg3, vreg1, vreg2, mask); |
| 352 | - MicroAPI::Copy<float, MicroAPI::MaskMergeMode::MERGING>(vreg1, vreg3, mask); | 352 | + MicroAPI::Move<float, MicroAPI::MaskMergeMode::MERGING>(vreg1, vreg3, mask); |
| 353 | 353 | ||
| 354 | - MicroAPI::DataCopy(grad1Addr, vreg1, maskFull); | 354 | + MicroAPI::StoreAlign(grad1Addr, vreg1, maskFull); |
| 355 | } | 355 | } |
| 356 | } | 356 | } |
| 357 | } | 357 | } |
| @@ -368,9 +368,9 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::UpdateCache(const LocalTens | |||
| 368 | uint32_t outerLoopStride = VL_FP32; | 368 | uint32_t outerLoopStride = VL_FP32; |
| 369 | uint32_t innerLoopStride = stride; | 369 | uint32_t innerLoopStride = stride; |
| 370 | 370 | ||
| 371 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 371 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 372 | - __local_mem__ float* cache = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheId * stride; | 372 | + __ubuf__ float* cache = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheId * stride; |
| 373 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 373 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 374 | 374 | ||
| 375 | __VEC_SCOPE__ | 375 | __VEC_SCOPE__ |
| 376 | { | 376 | { |
| @@ -379,12 +379,12 @@ __aicore__ inline void LogSoftmaxGradArRecompute<T>::UpdateCache(const LocalTens | |||
| 379 | MicroAPI::MaskReg pMask; | 379 | MicroAPI::MaskReg pMask; |
| 380 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 380 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 381 | pMask = MicroAPI::UpdateMask<float>(sreg); | 381 | pMask = MicroAPI::UpdateMask<float>(sreg); |
| 382 | - MicroAPI::DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 382 | + MicroAPI::LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 383 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 383 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 384 | - MicroAPI::DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride); | 384 | + MicroAPI::LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride); |
| 385 | MicroAPI::Add<float, MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 385 | MicroAPI::Add<float, MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 386 | } | 386 | } |
| 387 | - MicroAPI::DataCopy((__local_mem__ float*)cache + i * outerLoopStride, aReg, pMask); | 387 | + MicroAPI::StoreAlign((__ubuf__ float*)cache + i * outerLoopStride, aReg, pMask); |
| 388 | } | 388 | } |
| 389 | } | 389 | } |
| 390 | } | 390 | } |
| @@ -92,8 +92,8 @@ public: | |||
| 92 | LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>(); | 92 | LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>(); |
| 93 | LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>(); | 93 | LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>(); |
| 94 | 94 | ||
| 95 | - __local_mem__ T* gradLocal = (__local_mem__ T*)x0Tensor.GetPhyAddr(); | 95 | + __ubuf__ T* gradLocal = (__ubuf__ T*)x0Tensor.GetPhyAddr(); |
| 96 | - __local_mem__ T* xLocal = (__local_mem__ T*)x1Tensor.GetPhyAddr(); | 96 | + __ubuf__ T* xLocal = (__ubuf__ T*)x1Tensor.GetPhyAddr(); |
| 97 | CalcReduceSum(gradLocal, curTileA0Len); | 97 | CalcReduceSum(gradLocal, curTileA0Len); |
| 98 | CopyInAndTransPose(xOffsetPreLoad, nextTileA0Len, tilingData_->totalRLen); | 98 | CopyInAndTransPose(xOffsetPreLoad, nextTileA0Len, tilingData_->totalRLen); |
| 99 | CalcOutput(gradLocal, xLocal, curTileA0Len); | 99 | CalcOutput(gradLocal, xLocal, curTileA0Len); |
| @@ -110,8 +110,8 @@ public: | |||
| 110 | xOffset = curIdx * tilingData_->tileA0Len * tilingData_->totalRLen; | 110 | xOffset = curIdx * tilingData_->tileA0Len * tilingData_->totalRLen; |
| 111 | LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>(); | 111 | LocalTensor<T> x0Tensor = x0Queue_.DeQue<T>(); |
| 112 | LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>(); | 112 | LocalTensor<T> x1Tensor = x1Queue_.DeQue<T>(); |
| 113 | - __local_mem__ T* gradLocal = (__local_mem__ T*)x0Tensor.GetPhyAddr(); | 113 | + __ubuf__ T* gradLocal = (__ubuf__ T*)x0Tensor.GetPhyAddr(); |
| 114 | - __local_mem__ T* xLocal = (__local_mem__ T*)x1Tensor.GetPhyAddr(); | 114 | + __ubuf__ T* xLocal = (__ubuf__ T*)x1Tensor.GetPhyAddr(); |
| 115 | CalcReduceSum(gradLocal, curTileA0Len); | 115 | CalcReduceSum(gradLocal, curTileA0Len); |
| 116 | CalcOutput(gradLocal, xLocal, curTileA0Len); | 116 | CalcOutput(gradLocal, xLocal, curTileA0Len); |
| 117 | CalcTranspose(curTileA0Len, tilingData_->rAligned); | 117 | CalcTranspose(curTileA0Len, tilingData_->rAligned); |
| @@ -121,11 +121,11 @@ public: | |||
| 121 | } | 121 | } |
| 122 | 122 | ||
| 123 | private: | 123 | private: |
| 124 | - __aicore__ inline void CalcReduceSum(const __local_mem__ T* gradLocal, uint32_t curTileA0Len) | 124 | + __aicore__ inline void CalcReduceSum(const __ubuf__ T* gradLocal, uint32_t curTileA0Len) |
| 125 | { | 125 | { |
| 126 | - __local_mem__ float* tmpAddr = (__local_mem__ float*)tmpLocal_.GetPhyAddr(); | 126 | + __ubuf__ float* tmpAddr = (__ubuf__ float*)tmpLocal_.GetPhyAddr(); |
| 127 | - __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned] | 127 | + __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned] |
| 128 | - .GetPhyAddr(); | 128 | + .GetPhyAddr(); |
| 129 | 129 | ||
| 130 | uint32_t tileA0Len = tilingData_->tileA0Len; | 130 | uint32_t tileA0Len = tilingData_->tileA0Len; |
| 131 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); | 131 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); |
| @@ -144,8 +144,8 @@ private: | |||
| 144 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; | 144 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; |
| 145 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); | 145 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); |
| 146 | 146 | ||
| 147 | - DataCopy(tmpAddr + xOffset, gradReg, pregMask); | 147 | + StoreAlign(tmpAddr + xOffset, gradReg, pregMask); |
| 148 | - DataCopy(tmpAddr2 + xOffset, gradReg, pregMask); | 148 | + StoreAlign(tmpAddr2 + xOffset, gradReg, pregMask); |
| 149 | } | 149 | } |
| 150 | } | 150 | } |
| 151 | } | 151 | } |
| @@ -155,14 +155,13 @@ private: | |||
| 155 | AscendC::ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(xSumTensor_, tmpLocal_, srcShape, false); | 155 | AscendC::ReduceSum<float, AscendC::Pattern::Reduce::RA, true>(xSumTensor_, tmpLocal_, srcShape, false); |
| 156 | } | 156 | } |
| 157 | 157 | ||
| 158 | - __aicore__ inline void CalcOutput(const __local_mem__ T* gradLocal, const __local_mem__ T* xLocal, | 158 | + __aicore__ inline void CalcOutput(const __ubuf__ T* gradLocal, const __ubuf__ T* xLocal, uint32_t curTileA0Len) |
| 159 | - uint32_t curTileA0Len) | ||
| 160 | { | 159 | { |
| 161 | - __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor_.GetPhyAddr(); | 160 | + __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor_.GetPhyAddr(); |
| 162 | - __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned] | 161 | + __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tilingData_->tileA0Len * tilingData_->rAligned] |
| 163 | - .GetPhyAddr(); | 162 | + .GetPhyAddr(); |
| 164 | tmpLocalTy_ = tmpLocal_.template ReinterpretCast<T>(); | 163 | tmpLocalTy_ = tmpLocal_.template ReinterpretCast<T>(); |
| 165 | - __local_mem__ T* tmpAddrTy = (__local_mem__ T*)tmpLocalTy_.GetPhyAddr(); | 164 | + __ubuf__ T* tmpAddrTy = (__ubuf__ T*)tmpLocalTy_.GetPhyAddr(); |
| 166 | 165 | ||
| 167 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); | 166 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); |
| 168 | uint16_t loopA0Num = static_cast<uint16_t>(ops::CeilDiv(curTileA0Len, VL_FP32)); | 167 | uint16_t loopA0Num = static_cast<uint16_t>(ops::CeilDiv(curTileA0Len, VL_FP32)); |
| @@ -179,7 +178,7 @@ private: | |||
| 179 | 178 | ||
| 180 | for (uint16_t k = 0; k < loopA0Num; k++) { | 179 | for (uint16_t k = 0; k < loopA0Num; k++) { |
| 181 | pregMask = UpdateMask<float>(sreg); | 180 | pregMask = UpdateMask<float>(sreg); |
| 182 | - DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xSumLocal + k * VL_FP32); | 181 | + LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xSumLocal + k * VL_FP32); |
| 183 | for (uint16_t i = 0; i < curTileRLenVl; i++) { | 182 | for (uint16_t i = 0; i < curTileRLenVl; i++) { |
| 184 | uint32_t xOffset = i * tileA0LenLocal + k * VL_FP32; | 183 | uint32_t xOffset = i * tileA0LenLocal + k * VL_FP32; |
| 185 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); | 184 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); |
| @@ -190,11 +189,12 @@ private: | |||
| 190 | Sub(gradReg, gradReg, xReg, pregMask); | 189 | Sub(gradReg, gradReg, xReg, pregMask); |
| 191 | 190 | ||
| 192 | if constexpr (xToFp32_) { | 191 | if constexpr (xToFp32_) { |
| 193 | - MicroAPI::DataCopy(tmpAddrTy + xOffset, gradReg, pregMask); | 192 | + MicroAPI::StoreAlign(tmpAddrTy + xOffset, gradReg, pregMask); |
| 194 | } else { // fp16、bf16 | 193 | } else { // fp16、bf16 |
| 195 | RegTensor<T> xFp16; | 194 | RegTensor<T> xFp16; |
| 196 | MicroAPI::Cast<T, float, castTraitFp32ToFp16>(xFp16, gradReg, pregMask); | 195 | MicroAPI::Cast<T, float, castTraitFp32ToFp16>(xFp16, gradReg, pregMask); |
| 197 | - MicroAPI::DataCopy<T, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + xOffset, xFp16, pregMask); | 196 | + MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + xOffset, xFp16, |
| 197 | + pregMask); | ||
| 198 | } | 198 | } |
| 199 | } | 199 | } |
| 200 | } | 200 | } |
| @@ -266,29 +266,29 @@ private: | |||
| 266 | yQueue_.EnQue(yLocal_); | 266 | yQueue_.EnQue(yLocal_); |
| 267 | } | 267 | } |
| 268 | 268 | ||
| 269 | - __aicore__ inline void LoadTensorForDtypeT(const __local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg, | 269 | + __aicore__ inline void LoadTensorForDtypeT(const __ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg, |
| 270 | uint32_t offset) | 270 | uint32_t offset) |
| 271 | { | 271 | { |
| 272 | if constexpr (xToFp32_) { | 272 | if constexpr (xToFp32_) { |
| 273 | - DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset); | 273 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset); |
| 274 | } else { // fp16、bf16 | 274 | } else { // fp16、bf16 |
| 275 | RegTensor<T> xFp16; | 275 | RegTensor<T> xFp16; |
| 276 | - DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset)); | 276 | + LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset)); |
| 277 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); | 277 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 278 | } | 278 | } |
| 279 | } | 279 | } |
| 280 | 280 | ||
| 281 | __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen) | 281 | __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen) |
| 282 | { | 282 | { |
| 283 | - static constexpr MultiCopyConfig config = {false}; | 283 | + static constexpr NdDmaConfig config = {false}; |
| 284 | - MultiCopyLoopInfo<CONST_TWO> copyLoopInfo; | 284 | + NdDmaLoopInfo<CONST_TWO> copyLoopInfo; |
| 285 | copyLoopInfo.loopSrcStride[0] = 1; | 285 | copyLoopInfo.loopSrcStride[0] = 1; |
| 286 | copyLoopInfo.loopSrcStride[1] = totalRLen; | 286 | copyLoopInfo.loopSrcStride[1] = totalRLen; |
| 287 | copyLoopInfo.loopDstStride[0] = tilingData_->tileA0Len; | 287 | copyLoopInfo.loopDstStride[0] = tilingData_->tileA0Len; |
| 288 | copyLoopInfo.loopDstStride[1] = 1; | 288 | copyLoopInfo.loopDstStride[1] = 1; |
| 289 | copyLoopInfo.loopSize[0] = totalRLen; | 289 | copyLoopInfo.loopSize[0] = totalRLen; |
| 290 | copyLoopInfo.loopSize[1] = curTileA0Len; | 290 | copyLoopInfo.loopSize[1] = curTileA0Len; |
| 291 | - MultiCopyParams<T, CONST_TWO> params = {copyLoopInfo, 0}; | 291 | + NdDmaParams<T, CONST_TWO> params = {copyLoopInfo, 0}; |
| 292 | 292 | ||
| 293 | LocalTensor<T> x0Local_ = x0Queue_.AllocTensor<T>(); | 293 | LocalTensor<T> x0Local_ = x0Queue_.AllocTensor<T>(); |
| 294 | DataCopy<T, CONST_TWO, config>(x0Local_, x0Gm_[xGmOffset], params); | 294 | DataCopy<T, CONST_TWO, config>(x0Local_, x0Gm_[xGmOffset], params); |
| @@ -358,4 +358,4 @@ private: | |||
| 358 | }; | 358 | }; |
| 359 | } // namespace LogSoftmaxGradOps | 359 | } // namespace LogSoftmaxGradOps |
| 360 | 360 | ||
| 361 | -#endif | 361 | +#endif |
| @@ -93,8 +93,8 @@ public: | |||
| 93 | LocalTensor<T> gradTensor = gradQueue_.DeQue<T>(); | 93 | LocalTensor<T> gradTensor = gradQueue_.DeQue<T>(); |
| 94 | LocalTensor<T> xTensor = xQueue_.DeQue<T>(); | 94 | LocalTensor<T> xTensor = xQueue_.DeQue<T>(); |
| 95 | 95 | ||
| 96 | - __local_mem__ T* gradLocal = (__local_mem__ T*)gradTensor.GetPhyAddr(); | 96 | + __ubuf__ T* gradLocal = (__ubuf__ T*)gradTensor.GetPhyAddr(); |
| 97 | - __local_mem__ T* xLocal = (__local_mem__ T*)xTensor.GetPhyAddr(); | 97 | + __ubuf__ T* xLocal = (__ubuf__ T*)xTensor.GetPhyAddr(); |
| 98 | 98 | ||
| 99 | yMain_ = yQueue_.AllocTensor<float>(); | 99 | yMain_ = yQueue_.AllocTensor<float>(); |
| 100 | 100 | ||
| @@ -114,9 +114,9 @@ public: | |||
| 114 | } | 114 | } |
| 115 | 115 | ||
| 116 | private: | 116 | private: |
| 117 | - __aicore__ inline void CalcReduceSum(const __local_mem__ T* gradLocal, uint32_t curTileA0Len, int64_t a0BlockOffset) | 117 | + __aicore__ inline void CalcReduceSum(const __ubuf__ T* gradLocal, uint32_t curTileA0Len, int64_t a0BlockOffset) |
| 118 | { | 118 | { |
| 119 | - __local_mem__ float* yLocal = (__local_mem__ float*)yMain_.GetPhyAddr() + a0BlockOffset; | 119 | + __ubuf__ float* yLocal = (__ubuf__ float*)yMain_.GetPhyAddr() + a0BlockOffset; |
| 120 | 120 | ||
| 121 | uint32_t tileA0Len = tilingData_->tileA0Len; | 121 | uint32_t tileA0Len = tilingData_->tileA0Len; |
| 122 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); | 122 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); |
| @@ -134,7 +134,7 @@ private: | |||
| 134 | for (uint16_t i = 0; i < curTileRLenVl; i++) { | 134 | for (uint16_t i = 0; i < curTileRLenVl; i++) { |
| 135 | uint32_t gradOffset = i * tileA0Len + k * VL_FP32; | 135 | uint32_t gradOffset = i * tileA0Len + k * VL_FP32; |
| 136 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, gradOffset); | 136 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, gradOffset); |
| 137 | - DataCopy(((__local_mem__ float*)yLocal) + gradOffset, gradReg, pregMask); | 137 | + StoreAlign(((__ubuf__ float*)yLocal) + gradOffset, gradReg, pregMask); |
| 138 | } | 138 | } |
| 139 | } | 139 | } |
| 140 | } | 140 | } |
| @@ -145,11 +145,11 @@ private: | |||
| 145 | false); | 145 | false); |
| 146 | } | 146 | } |
| 147 | 147 | ||
| 148 | - __aicore__ inline void CalcOutput(const __local_mem__ T* gradLocal, const __local_mem__ T* xLocal, | 148 | + __aicore__ inline void CalcOutput(const __ubuf__ T* gradLocal, const __ubuf__ T* xLocal, uint32_t curTileA0Len, |
| 149 | - uint32_t curTileA0Len, int64_t a0BlockOffset) | 149 | + int64_t a0BlockOffset) |
| 150 | { | 150 | { |
| 151 | - __local_mem__ T* yLocal = (__local_mem__ T*)yMain_.GetPhyAddr() + a0BlockOffset; | 151 | + __ubuf__ T* yLocal = (__ubuf__ T*)yMain_.GetPhyAddr() + a0BlockOffset; |
| 152 | - __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor_.GetPhyAddr(); | 152 | + __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor_.GetPhyAddr(); |
| 153 | 153 | ||
| 154 | uint32_t tileA0Len = tilingData_->tileA0Len; | 154 | uint32_t tileA0Len = tilingData_->tileA0Len; |
| 155 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); | 155 | uint16_t curTileRLenVl = static_cast<uint16_t>(tilingData_->totalRLen); |
| @@ -166,7 +166,7 @@ private: | |||
| 166 | 166 | ||
| 167 | for (uint16_t k = 0; k < loopA0Num; k++) { | 167 | for (uint16_t k = 0; k < loopA0Num; k++) { |
| 168 | pregMask = UpdateMask<float>(sreg); | 168 | pregMask = UpdateMask<float>(sreg); |
| 169 | - DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xSumLocal + k * VL_FP32); | 169 | + LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xSumLocal + k * VL_FP32); |
| 170 | for (uint16_t i = 0; i < curTileRLenVl; i++) { | 170 | for (uint16_t i = 0; i < curTileRLenVl; i++) { |
| 171 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; | 171 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; |
| 172 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); | 172 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); |
| @@ -177,25 +177,25 @@ private: | |||
| 177 | Sub(gradReg, gradReg, xReg, pregMask); | 177 | Sub(gradReg, gradReg, xReg, pregMask); |
| 178 | 178 | ||
| 179 | if constexpr (IsSameType<T, float>::value) { | 179 | if constexpr (IsSameType<T, float>::value) { |
| 180 | - DataCopy(((__local_mem__ float*)yLocal) + xOffset, gradReg, pregMask); | 180 | + StoreAlign(((__ubuf__ float*)yLocal) + xOffset, gradReg, pregMask); |
| 181 | } else { // fp16、bf16 | 181 | } else { // fp16、bf16 |
| 182 | RegTensor<T> xFp16; | 182 | RegTensor<T> xFp16; |
| 183 | Cast<T, float, castTraitFp32ToFp16>(xFp16, gradReg, pregMask); | 183 | Cast<T, float, castTraitFp32ToFp16>(xFp16, gradReg, pregMask); |
| 184 | - DataCopy<T, StoreDist::DIST_PACK_B32>(((__local_mem__ T*)yLocal) + xOffset, xFp16, pregMask); | 184 | + StoreAlign<T, StoreDist::DIST_PACK_B32>(((__ubuf__ T*)yLocal) + xOffset, xFp16, pregMask); |
| 185 | } | 185 | } |
| 186 | } | 186 | } |
| 187 | } | 187 | } |
| 188 | } | 188 | } |
| 189 | } | 189 | } |
| 190 | 190 | ||
| 191 | - __aicore__ inline void LoadTensorForDtypeT(const __local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg, | 191 | + __aicore__ inline void LoadTensorForDtypeT(const __ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg, |
| 192 | uint32_t offset) | 192 | uint32_t offset) |
| 193 | { | 193 | { |
| 194 | if constexpr (IsSameType<T, float>::value) { | 194 | if constexpr (IsSameType<T, float>::value) { |
| 195 | - DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset); | 195 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset); |
| 196 | } else { // fp16、bf16 | 196 | } else { // fp16、bf16 |
| 197 | RegTensor<T> xFp16; | 197 | RegTensor<T> xFp16; |
| 198 | - DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset)); | 198 | + LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset)); |
| 199 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); | 199 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 200 | } | 200 | } |
| 201 | } | 201 | } |
| @@ -151,8 +151,8 @@ private: | |||
| 151 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); | 151 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); |
| 152 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; | 152 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; |
| 153 | 153 | ||
| 154 | - __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr(); | 154 | + __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr(); |
| 155 | - __local_mem__ T* grad = (__local_mem__ T*)gradMain_.GetPhyAddr(); | 155 | + __ubuf__ T* grad = (__ubuf__ T*)gradMain_.GetPhyAddr(); |
| 156 | 156 | ||
| 157 | __VEC_SCOPE__ | 157 | __VEC_SCOPE__ |
| 158 | { | 158 | { |
| @@ -167,7 +167,7 @@ private: | |||
| 167 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 167 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 168 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; | 168 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; |
| 169 | LoadTensorForDtypeT(grad, gradReg, pregMask, xOffset); | 169 | LoadTensorForDtypeT(grad, gradReg, pregMask, xOffset); |
| 170 | - DataCopy((__local_mem__ float*)dst + xOffset, gradReg, pregMask); | 170 | + StoreAlign((__ubuf__ float*)dst + xOffset, gradReg, pregMask); |
| 171 | } | 171 | } |
| 172 | } | 172 | } |
| 173 | } | 173 | } |
| @@ -186,8 +186,8 @@ private: | |||
| 186 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); | 186 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); |
| 187 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; | 187 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; |
| 188 | 188 | ||
| 189 | - __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr(); | 189 | + __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr(); |
| 190 | - __local_mem__ T* grad = (__local_mem__ T*)gradFold.GetPhyAddr(); | 190 | + __ubuf__ T* grad = (__ubuf__ T*)gradFold.GetPhyAddr(); |
| 191 | 191 | ||
| 192 | __VEC_SCOPE__ | 192 | __VEC_SCOPE__ |
| 193 | { | 193 | { |
| @@ -202,9 +202,9 @@ private: | |||
| 202 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 202 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 203 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; | 203 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; |
| 204 | LoadTensorForDtypeT(grad, grad0Reg, pregMask, i * outerLoopSrcStride + j * VL_FP32); | 204 | LoadTensorForDtypeT(grad, grad0Reg, pregMask, i * outerLoopSrcStride + j * VL_FP32); |
| 205 | - DataCopy(grad1Reg, (__local_mem__ float*)dst + xOffset); | 205 | + LoadAlign(grad1Reg, (__ubuf__ float*)dst + xOffset); |
| 206 | Add(grad0Reg, grad1Reg, grad0Reg, pregMask); | 206 | Add(grad0Reg, grad1Reg, grad0Reg, pregMask); |
| 207 | - DataCopy((__local_mem__ float*)dst + xOffset, grad0Reg, pregMask); | 207 | + StoreAlign((__ubuf__ float*)dst + xOffset, grad0Reg, pregMask); |
| 208 | } | 208 | } |
| 209 | } | 209 | } |
| 210 | } | 210 | } |
| @@ -258,14 +258,14 @@ private: | |||
| 258 | __aicore__ inline void CalcOutput(int64_t curTileRLen, uint32_t curTileA0Len, uint16_t loopA0Num) | 258 | __aicore__ inline void CalcOutput(int64_t curTileRLen, uint32_t curTileA0Len, uint16_t loopA0Num) |
| 259 | { | 259 | { |
| 260 | LocalTensor<T> grad = gradQueue_.DeQue<T>(); | 260 | LocalTensor<T> grad = gradQueue_.DeQue<T>(); |
| 261 | - __local_mem__ T* gradLocal = (__local_mem__ T*)grad.GetPhyAddr(); | 261 | + __ubuf__ T* gradLocal = (__ubuf__ T*)grad.GetPhyAddr(); |
| 262 | LocalTensor<T> x = xQueue_.DeQue<T>(); | 262 | LocalTensor<T> x = xQueue_.DeQue<T>(); |
| 263 | - __local_mem__ T* xLocal = (__local_mem__ T*)x.GetPhyAddr(); | 263 | + __ubuf__ T* xLocal = (__ubuf__ T*)x.GetPhyAddr(); |
| 264 | 264 | ||
| 265 | LocalTensor<T> y = yQueue_.template AllocTensor<T>(); | 265 | LocalTensor<T> y = yQueue_.template AllocTensor<T>(); |
| 266 | - __local_mem__ T* yLocal = (__local_mem__ T*)y.GetPhyAddr(); | 266 | + __ubuf__ T* yLocal = (__ubuf__ T*)y.GetPhyAddr(); |
| 267 | 267 | ||
| 268 | - __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor_.GetPhyAddr(); | 268 | + __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor_.GetPhyAddr(); |
| 269 | 269 | ||
| 270 | uint32_t tileA0Len = tilingData_->tileA0Len; | 270 | uint32_t tileA0Len = tilingData_->tileA0Len; |
| 271 | uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen); | 271 | uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen); |
| @@ -280,7 +280,7 @@ private: | |||
| 280 | 280 | ||
| 281 | for (uint16_t k = 0; k < loopA0Num; k++) { | 281 | for (uint16_t k = 0; k < loopA0Num; k++) { |
| 282 | pregMask = UpdateMask<float>(sreg); | 282 | pregMask = UpdateMask<float>(sreg); |
| 283 | - DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xSumLocal + k * VL_FP32); | 283 | + LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xSumLocal + k * VL_FP32); |
| 284 | for (uint16_t i = 0; i < curTileRLenVl; i++) { | 284 | for (uint16_t i = 0; i < curTileRLenVl; i++) { |
| 285 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; | 285 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; |
| 286 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); | 286 | LoadTensorForDtypeT(gradLocal, gradReg, pregMask, xOffset); |
| @@ -292,11 +292,11 @@ private: | |||
| 292 | 292 | ||
| 293 | // copy out | 293 | // copy out |
| 294 | if constexpr (IsSameType<T, float>::value) { | 294 | if constexpr (IsSameType<T, float>::value) { |
| 295 | - DataCopy(((__local_mem__ float*)yLocal) + xOffset, gradReg, pregMask); | 295 | + StoreAlign(((__ubuf__ float*)yLocal) + xOffset, gradReg, pregMask); |
| 296 | } else { // fp16、bf16 | 296 | } else { // fp16、bf16 |
| 297 | RegTensor<T> xFp16; | 297 | RegTensor<T> xFp16; |
| 298 | Cast<T, float, castTraitFp32ToFp16>(xFp16, gradReg, pregMask); | 298 | Cast<T, float, castTraitFp32ToFp16>(xFp16, gradReg, pregMask); |
| 299 | - DataCopy<T, StoreDist::DIST_PACK_B32>(((__local_mem__ T*)yLocal) + xOffset, xFp16, pregMask); | 299 | + StoreAlign<T, StoreDist::DIST_PACK_B32>(((__ubuf__ T*)yLocal) + xOffset, xFp16, pregMask); |
| 300 | } | 300 | } |
| 301 | } | 301 | } |
| 302 | } | 302 | } |
| @@ -308,14 +308,13 @@ private: | |||
| 308 | xQueue_.FreeTensor<T>(x); | 308 | xQueue_.FreeTensor<T>(x); |
| 309 | } | 309 | } |
| 310 | 310 | ||
| 311 | - __aicore__ inline void LoadTensorForDtypeT(__local_mem__ T* src, RegTensor<float>& dst, MaskReg& preg, | 311 | + __aicore__ inline void LoadTensorForDtypeT(__ubuf__ T* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset) |
| 312 | - uint32_t offset) | ||
| 313 | { | 312 | { |
| 314 | if constexpr (IsSameType<T, float>::value) { | 313 | if constexpr (IsSameType<T, float>::value) { |
| 315 | - DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset); | 314 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset); |
| 316 | } else { // fp16、bf16 | 315 | } else { // fp16、bf16 |
| 317 | RegTensor<T> xFp16; | 316 | RegTensor<T> xFp16; |
| 318 | - DataCopy<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T*)src + offset)); | 317 | + LoadAlign<T, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T*)src + offset)); |
| 319 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); | 318 | Cast<float, T, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 320 | } | 319 | } |
| 321 | } | 320 | } |
| @@ -28,6 +28,10 @@ | |||
| 28 | 28 | ||
| 29 | namespace LogSoftmaxGradOps { | 29 | namespace LogSoftmaxGradOps { |
| 30 | using namespace AscendC; | 30 | using namespace AscendC; |
| 31 | +using AscendC::Reg::LoadAlign; | ||
| 32 | +using AscendC::Reg::Move; | ||
| 33 | +using AscendC::Reg::Reduce; | ||
| 34 | +using AscendC::Reg::StoreAlign; | ||
| 31 | 35 | ||
| 32 | constexpr static AscendC::MicroAPI::CastTrait castTraitFp16ToFp32 = { | 36 | constexpr static AscendC::MicroAPI::CastTrait castTraitFp16ToFp32 = { |
| 33 | AscendC::MicroAPI::RegLayout::ZERO, | 37 | AscendC::MicroAPI::RegLayout::ZERO, |
| @@ -180,8 +184,8 @@ __aicore__ inline void LogSoftmaxGradOpsBase::CastToFp32From(const LocalTensor<f | |||
| 180 | if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) { | 184 | if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) { |
| 181 | __VEC_SCOPE__ | 185 | __VEC_SCOPE__ |
| 182 | { | 186 | { |
| 183 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 187 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 184 | - __local_mem__ T* src = (__local_mem__ T*)srcTensor.GetPhyAddr(); | 188 | + __ubuf__ T* src = (__ubuf__ T*)srcTensor.GetPhyAddr(); |
| 185 | uint32_t count; | 189 | uint32_t count; |
| 186 | AscendC::MicroAPI::RegTensor<float> fp32Reg; | 190 | AscendC::MicroAPI::RegTensor<float> fp32Reg; |
| 187 | AscendC::MicroAPI::RegTensor<T> b16Reg; | 191 | AscendC::MicroAPI::RegTensor<T> b16Reg; |
| @@ -190,10 +194,10 @@ __aicore__ inline void LogSoftmaxGradOpsBase::CastToFp32From(const LocalTensor<f | |||
| 190 | count = static_cast<uint32_t>(colSize); | 194 | count = static_cast<uint32_t>(colSize); |
| 191 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 195 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 192 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 196 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 193 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( | 197 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( |
| 194 | - b16Reg, (__local_mem__ T*)src + i * outerLoopSrcStride + j * innerLoopStride); | 198 | + b16Reg, (__ubuf__ T*)src + i * outerLoopSrcStride + j * innerLoopStride); |
| 195 | Cast<float, T, castTraitFp16ToFp32>(fp32Reg, b16Reg, pMask); | 199 | Cast<float, T, castTraitFp16ToFp32>(fp32Reg, b16Reg, pMask); |
| 196 | - DataCopy((__local_mem__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask); | 200 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask); |
| 197 | } | 201 | } |
| 198 | } | 202 | } |
| 199 | } | 203 | } |
| @@ -223,8 +227,8 @@ __aicore__ inline void LogSoftmaxGradOpsBase::CastFromFp32To(const LocalTensor<T | |||
| 223 | if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) { | 227 | if constexpr (IsSameType<T, half>::value || IsSameType<T, bfloat16_t>::value) { |
| 224 | __VEC_SCOPE__ | 228 | __VEC_SCOPE__ |
| 225 | { | 229 | { |
| 226 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 230 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 227 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 231 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 228 | uint32_t count; | 232 | uint32_t count; |
| 229 | AscendC::MicroAPI::RegTensor<float> fp32Reg; | 233 | AscendC::MicroAPI::RegTensor<float> fp32Reg; |
| 230 | AscendC::MicroAPI::RegTensor<T> b16Reg; | 234 | AscendC::MicroAPI::RegTensor<T> b16Reg; |
| @@ -233,10 +237,10 @@ __aicore__ inline void LogSoftmaxGradOpsBase::CastFromFp32To(const LocalTensor<T | |||
| 233 | count = static_cast<uint32_t>(colSize); | 237 | count = static_cast<uint32_t>(colSize); |
| 234 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 238 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 235 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 239 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 236 | - DataCopy(fp32Reg, (__local_mem__ float*)src + i * outerLoopSrcStride + j * innerLoopStride); | 240 | + LoadAlign(fp32Reg, (__ubuf__ float*)src + i * outerLoopSrcStride + j * innerLoopStride); |
| 237 | Cast<T, float, castTraitFp32ToFp16>(b16Reg, fp32Reg, pMask); | 241 | Cast<T, float, castTraitFp32ToFp16>(b16Reg, fp32Reg, pMask); |
| 238 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>( | 242 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>( |
| 239 | - (__local_mem__ T*)dst + i * outerLoopDstStride + j * innerLoopStride, b16Reg, pMask); | 243 | + (__ubuf__ T*)dst + i * outerLoopDstStride + j * innerLoopStride, b16Reg, pMask); |
| 240 | } | 244 | } |
| 241 | } | 245 | } |
| 242 | } | 246 | } |
| @@ -322,19 +326,19 @@ __aicore__ inline void LogSoftmaxGradOpsBase::VectorAdd(const LocalTensor<float> | |||
| 322 | static_cast<int64_t>(platform::GetVRegSize())); | 326 | static_cast<int64_t>(platform::GetVRegSize())); |
| 323 | __VEC_SCOPE__ | 327 | __VEC_SCOPE__ |
| 324 | { | 328 | { |
| 325 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 329 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 326 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 330 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 327 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 331 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 328 | uint32_t sreg = static_cast<uint32_t>(count); | 332 | uint32_t sreg = static_cast<uint32_t>(count); |
| 329 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 333 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 330 | AscendC::MicroAPI::MaskReg pMask; | 334 | AscendC::MicroAPI::MaskReg pMask; |
| 331 | for (uint16_t i = 0; i < loopTimes; ++i) { | 335 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 332 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 336 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 333 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32); | 337 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32); |
| 334 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32); | 338 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32); |
| 335 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 339 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 336 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 340 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 337 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask); | 341 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask); |
| 338 | } | 342 | } |
| 339 | } | 343 | } |
| 340 | } | 344 | } |
| @@ -352,20 +356,20 @@ __aicore__ inline void LogSoftmaxGradOpsBase::VectorAdd(const LocalTensor<float> | |||
| 352 | uint32_t innerLoopStride = stride; | 356 | uint32_t innerLoopStride = stride; |
| 353 | __VEC_SCOPE__ | 357 | __VEC_SCOPE__ |
| 354 | { | 358 | { |
| 355 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 359 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 356 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 360 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 357 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 361 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 358 | uint32_t count = nSize; | 362 | uint32_t count = nSize; |
| 359 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 363 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 360 | AscendC::MicroAPI::MaskReg pMask; | 364 | AscendC::MicroAPI::MaskReg pMask; |
| 361 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 365 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 362 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 366 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 363 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 367 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 364 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride); | 368 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride); |
| 365 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride + j * innerLoopStride); | 369 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride + j * innerLoopStride); |
| 366 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 370 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 367 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 371 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 368 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask); | 372 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask); |
| 369 | } | 373 | } |
| 370 | } | 374 | } |
| 371 | } | 375 | } |
| @@ -383,19 +387,19 @@ __aicore__ inline void LogSoftmaxGradOpsBase::VectorMul(const LocalTensor<float> | |||
| 383 | static_cast<int64_t>(platform::GetVRegSize())); | 387 | static_cast<int64_t>(platform::GetVRegSize())); |
| 384 | __VEC_SCOPE__ | 388 | __VEC_SCOPE__ |
| 385 | { | 389 | { |
| 386 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 390 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 387 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 391 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 388 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 392 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 389 | uint32_t sreg = static_cast<uint32_t>(count); | 393 | uint32_t sreg = static_cast<uint32_t>(count); |
| 390 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 394 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 391 | AscendC::MicroAPI::MaskReg pMask; | 395 | AscendC::MicroAPI::MaskReg pMask; |
| 392 | 396 | ||
| 393 | for (uint16_t i = 0; i < loopTimes; ++i) { | 397 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 394 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 398 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 395 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32); | 399 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32); |
| 396 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32); | 400 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32); |
| 397 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 401 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 398 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, cReg, pMask); | 402 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, cReg, pMask); |
| 399 | } | 403 | } |
| 400 | } | 404 | } |
| 401 | } | 405 | } |
| @@ -419,19 +423,19 @@ __aicore__ inline void LogSoftmaxGradOpsBase::NlastBroadcastMul(const LocalTenso | |||
| 419 | uint32_t innerLoopStride = aSize; | 423 | uint32_t innerLoopStride = aSize; |
| 420 | __VEC_SCOPE__ | 424 | __VEC_SCOPE__ |
| 421 | { | 425 | { |
| 422 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 426 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 423 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 427 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 424 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 428 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 425 | uint32_t count = static_cast<uint32_t>(aSize); | 429 | uint32_t count = static_cast<uint32_t>(aSize); |
| 426 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 430 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 427 | AscendC::MicroAPI::MaskReg pMask; | 431 | AscendC::MicroAPI::MaskReg pMask; |
| 428 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 432 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 429 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 433 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 430 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride); | 434 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride); |
| 431 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 435 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 432 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride); | 436 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride); |
| 433 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 437 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 434 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); | 438 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); |
| 435 | } | 439 | } |
| 436 | } | 440 | } |
| 437 | } | 441 | } |
| @@ -457,40 +461,40 @@ __aicore__ inline void LogSoftmaxGradOpsBase::LastReduceSumSmallR(const LocalTen | |||
| 457 | if (rSize <= VL_FP32) { | 461 | if (rSize <= VL_FP32) { |
| 458 | __VEC_SCOPE__ | 462 | __VEC_SCOPE__ |
| 459 | { | 463 | { |
| 460 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 464 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 461 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 465 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 462 | uint32_t count = static_cast<uint32_t>(rSize); | 466 | uint32_t count = static_cast<uint32_t>(rSize); |
| 463 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 467 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 464 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 468 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 465 | - AscendC::MicroAPI::UnalignReg UReg; | 469 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 466 | for (uint16_t i = 0; i < loopTimes; ++i) { | 470 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 467 | - DataCopy(aReg, (__local_mem__ float*)src + i * stride); | 471 | + LoadAlign(aReg, (__ubuf__ float*)src + i * stride); |
| 468 | - ReduceSum(bReg, aReg, pMask); | 472 | + Reduce<ReduceType::SUM>(bReg, aReg, pMask); |
| 469 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 473 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 470 | } | 474 | } |
| 471 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 475 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 472 | } | 476 | } |
| 473 | } else { | 477 | } else { |
| 474 | __VEC_SCOPE__ | 478 | __VEC_SCOPE__ |
| 475 | { | 479 | { |
| 476 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 480 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 477 | - __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 481 | + __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 478 | - __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32; | 482 | + __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32; |
| 479 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); | 483 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); |
| 480 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 484 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 481 | - AscendC::MicroAPI::UnalignReg UReg; | 485 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 482 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 486 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 483 | AscendC::MicroAPI::MaskReg | 487 | AscendC::MicroAPI::MaskReg |
| 484 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 488 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 485 | for (uint16_t i = 0; i < loopTimes; ++i) { | 489 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 486 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * stride); | 490 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * stride); |
| 487 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * stride); | 491 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * stride); |
| 488 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 492 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 489 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 493 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 490 | - ReduceSum(bReg, aReg, pFull); | 494 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 491 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 495 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 492 | } | 496 | } |
| 493 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 497 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 494 | } | 498 | } |
| 495 | } | 499 | } |
| 496 | } | 500 | } |
| @@ -536,43 +540,43 @@ __aicore__ inline void LogSoftmaxGradOpsBase::LastReduceSum(const LocalTensor<fl | |||
| 536 | 540 | ||
| 537 | __VEC_SCOPE__ | 541 | __VEC_SCOPE__ |
| 538 | { | 542 | { |
| 539 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr(); | 543 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr(); |
| 540 | - __local_mem__ float* foldSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 544 | + __ubuf__ float* foldSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 541 | - __local_mem__ float* foldSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; | 545 | + __ubuf__ float* foldSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; |
| 542 | - __local_mem__ float* tailSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; | 546 | + __ubuf__ float* tailSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; |
| 543 | - __local_mem__ float* tailSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; | 547 | + __ubuf__ float* tailSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; |
| 544 | - __local_mem__ float* unFoldSrc = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; | 548 | + __ubuf__ float* unFoldSrc = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; |
| 545 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 549 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 546 | - AscendC::MicroAPI::UnalignReg UReg; | 550 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 547 | 551 | ||
| 548 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 552 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 549 | - dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; | 553 | + dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; |
| 550 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { | 554 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { |
| 551 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; | 555 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; |
| 552 | - DataCopy(aReg, (__local_mem__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); | 556 | + LoadAlign(aReg, (__ubuf__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); |
| 553 | - DataCopy(bReg, (__local_mem__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); | 557 | + LoadAlign(bReg, (__ubuf__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); |
| 554 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); | 558 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); |
| 555 | - ReduceSum(dReg, cReg, pFull); | 559 | + Reduce<ReduceType::SUM>(dReg, cReg, pFull); |
| 556 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, dReg, UReg, 1); | 560 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, dReg, UReg, 1); |
| 557 | } | 561 | } |
| 558 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { | 562 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { |
| 559 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); | 563 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); |
| 560 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 564 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 561 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 565 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 562 | - DataCopy(aReg, (__local_mem__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); | 566 | + LoadAlign(aReg, (__ubuf__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); |
| 563 | - DataCopy(bReg, (__local_mem__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); | 567 | + LoadAlign(bReg, (__ubuf__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); |
| 564 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 568 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 565 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 569 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 566 | - ReduceSum(bReg, aReg, pFull); | 570 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 567 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 571 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 568 | } | 572 | } |
| 569 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { | 573 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { |
| 570 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 574 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 571 | - DataCopy(aReg, (__local_mem__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); | 575 | + LoadAlign(aReg, (__ubuf__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); |
| 572 | - ReduceSum(bReg, aReg, pFull); | 576 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 573 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 577 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 574 | } | 578 | } |
| 575 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 579 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 576 | } | 580 | } |
| 577 | } | 581 | } |
| 578 | LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); | 582 | LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); |
| @@ -580,13 +584,13 @@ __aicore__ inline void LogSoftmaxGradOpsBase::LastReduceSum(const LocalTensor<fl | |||
| 580 | 584 | ||
| 581 | template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1> | 585 | template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1> |
| 582 | struct NlastDichotomyAdd { | 586 | struct NlastDichotomyAdd { |
| 583 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 587 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 584 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 588 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 585 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 589 | + uint32_t stride) |
| 586 | { | 590 | { |
| 587 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 591 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 588 | - __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO; | 592 | + __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO; |
| 589 | - __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO; | 593 | + __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO; |
| 590 | if constexpr (TailCount <= 0) { | 594 | if constexpr (TailCount <= 0) { |
| 591 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, | 595 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, |
| 592 | stride * CONST_TWO); | 596 | stride * CONST_TWO); |
| @@ -594,13 +598,13 @@ struct NlastDichotomyAdd { | |||
| 594 | } | 598 | } |
| 595 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 599 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 596 | } | 600 | } |
| 597 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 601 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 598 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 602 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 599 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 603 | + uint32_t stride, uint32_t offset) |
| 600 | { | 604 | { |
| 601 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 605 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 602 | - __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO; | 606 | + __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO; |
| 603 | - __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO; | 607 | + __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO; |
| 604 | if constexpr (TailCount <= 0) { | 608 | if constexpr (TailCount <= 0) { |
| 605 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, | 609 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, |
| 606 | stride * CONST_TWO, offset); | 610 | stride * CONST_TWO, offset); |
| @@ -618,49 +622,49 @@ struct NlastDichotomyAdd { | |||
| 618 | 622 | ||
| 619 | template <int32_t TailCount, int32_t Index, int32_t Depth> | 623 | template <int32_t TailCount, int32_t Index, int32_t Depth> |
| 620 | struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { | 624 | struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { |
| 621 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 625 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 622 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 626 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 623 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 627 | + uint32_t stride) |
| 624 | { | 628 | { |
| 625 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 629 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 626 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 630 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 627 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 631 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 628 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 632 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 629 | } | 633 | } |
| 630 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 634 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 631 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 635 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 632 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 636 | + uint32_t stride, uint32_t offset) |
| 633 | { | 637 | { |
| 634 | if constexpr (TailCount <= 0) { | 638 | if constexpr (TailCount <= 0) { |
| 635 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 639 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 636 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 640 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 637 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 641 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 638 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 642 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 639 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 643 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 640 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 644 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 641 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 645 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 642 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 646 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 643 | } else { | 647 | } else { |
| 644 | if constexpr (Index + Depth < TailCount) { | 648 | if constexpr (Index + Depth < TailCount) { |
| 645 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 649 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 646 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 650 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 647 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 651 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 648 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 652 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 649 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 653 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 650 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 654 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 651 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 655 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 652 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 656 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 653 | } else if constexpr (Index < TailCount) { | 657 | } else if constexpr (Index < TailCount) { |
| 654 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 658 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 655 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 659 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 656 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 660 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 657 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 661 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 658 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 662 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 659 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 663 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 660 | } else { | 664 | } else { |
| 661 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 665 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 662 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 666 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 663 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 667 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 664 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 668 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 665 | } | 669 | } |
| 666 | } | 670 | } |
| @@ -669,25 +673,25 @@ struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { | |||
| 669 | 673 | ||
| 670 | template <> | 674 | template <> |
| 671 | struct NlastDichotomyAdd<CONST_TWO> { | 675 | struct NlastDichotomyAdd<CONST_TWO> { |
| 672 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 676 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 673 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 677 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 674 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 678 | + uint32_t stride) |
| 675 | { | 679 | { |
| 676 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 680 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 677 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 681 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 678 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 682 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 679 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 683 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 680 | } | 684 | } |
| 681 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 685 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 682 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 686 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 683 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 687 | + uint32_t stride, uint32_t offset) |
| 684 | { | 688 | { |
| 685 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 689 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 686 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 690 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 687 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 691 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 688 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 692 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 689 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 693 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 690 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 694 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 691 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 695 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 692 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 696 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 693 | } | 697 | } |
| @@ -695,11 +699,11 @@ struct NlastDichotomyAdd<CONST_TWO> { | |||
| 695 | 699 | ||
| 696 | template <> | 700 | template <> |
| 697 | struct NlastDichotomyAdd<1> { | 701 | struct NlastDichotomyAdd<1> { |
| 698 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 702 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 699 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 703 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 700 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 704 | + uint32_t stride) |
| 701 | { | 705 | { |
| 702 | - DataCopy(acc, (__local_mem__ float*)srcA); | 706 | + LoadAlign(acc, (__ubuf__ float*)srcA); |
| 703 | } | 707 | } |
| 704 | }; | 708 | }; |
| 705 | 709 | ||
| @@ -714,32 +718,32 @@ __aicore__ inline void LogSoftmaxGradOpsBase::NlastReduceSumSmallR(const LocalTe | |||
| 714 | if constexpr (RSize == 1) { | 718 | if constexpr (RSize == 1) { |
| 715 | __VEC_SCOPE__ | 719 | __VEC_SCOPE__ |
| 716 | { | 720 | { |
| 717 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 721 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 718 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 722 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 719 | uint32_t count = static_cast<uint32_t>(aSize); | 723 | uint32_t count = static_cast<uint32_t>(aSize); |
| 720 | AscendC::MicroAPI::RegTensor<float> aReg; | 724 | AscendC::MicroAPI::RegTensor<float> aReg; |
| 721 | AscendC::MicroAPI::MaskReg pMask; | 725 | AscendC::MicroAPI::MaskReg pMask; |
| 722 | for (uint16_t i = 0; i < loopTimes; ++i) { | 726 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 723 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 727 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 724 | - DataCopy(aReg, (__local_mem__ float*)src + i * VL_FP32); | 728 | + LoadAlign(aReg, (__ubuf__ float*)src + i * VL_FP32); |
| 725 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask); | 729 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask); |
| 726 | } | 730 | } |
| 727 | } | 731 | } |
| 728 | } else { | 732 | } else { |
| 729 | __VEC_SCOPE__ | 733 | __VEC_SCOPE__ |
| 730 | { | 734 | { |
| 731 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 735 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 732 | - __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 736 | + __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 733 | - __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride; | 737 | + __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride; |
| 734 | uint32_t count = static_cast<uint32_t>(aSize); | 738 | uint32_t count = static_cast<uint32_t>(aSize); |
| 735 | AscendC::MicroAPI::RegTensor<float> aReg; | 739 | AscendC::MicroAPI::RegTensor<float> aReg; |
| 736 | AscendC::MicroAPI::MaskReg pMask; | 740 | AscendC::MicroAPI::MaskReg pMask; |
| 737 | for (uint16_t i = 0; i < loopTimes; ++i) { | 741 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 738 | - __local_mem__ float* curSrcA = srcA + i * VL_FP32; | 742 | + __ubuf__ float* curSrcA = srcA + i * VL_FP32; |
| 739 | - __local_mem__ float* curSrcB = srcB + i * VL_FP32; | 743 | + __ubuf__ float* curSrcB = srcB + i * VL_FP32; |
| 740 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 744 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 741 | NlastDichotomyAdd<RSize>::LoadAndAccumulate(aReg, curSrcA, curSrcB, pMask, stride); | 745 | NlastDichotomyAdd<RSize>::LoadAndAccumulate(aReg, curSrcA, curSrcB, pMask, stride); |
| 742 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask); | 746 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask); |
| 743 | } | 747 | } |
| 744 | } | 748 | } |
| 745 | } | 749 | } |
| @@ -812,35 +816,34 @@ __aicore__ inline void LogSoftmaxGradOpsBase::NlastReduceSumLargeR(const LocalTe | |||
| 812 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 816 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 813 | pMask = plt_b32(count, POST_UPDATE); | 817 | pMask = plt_b32(count, POST_UPDATE); |
| 814 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { | 818 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { |
| 815 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + | 819 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopStride + |
| 816 | - i * outerLoopStride + j * innerLoopStride; | 820 | + j * innerLoopStride; |
| 817 | - __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride + | 821 | + __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride + j * srcStride; |
| 818 | - j * srcStride; | 822 | + __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride + i * outerLoopStride + |
| 819 | - __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride + | 823 | + j * srcStride; |
| 820 | - i * outerLoopStride + j * srcStride; | ||
| 821 | NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, foldOffset); | 824 | NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, foldOffset); |
| 822 | - DataCopy((__local_mem__ float*)dst, aReg, pMask); | 825 | + StoreAlign((__ubuf__ float*)dst, aReg, pMask); |
| 823 | } | 826 | } |
| 824 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { | 827 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { |
| 825 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + | 828 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopStride + |
| 826 | - i * outerLoopStride + mainFoldLoopTimes * innerLoopStride; | 829 | + mainFoldLoopTimes * innerLoopStride; |
| 827 | - __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride + | 830 | + __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + i * outerLoopStride + |
| 828 | - mainFoldLoopTimes * srcStride; | 831 | + mainFoldLoopTimes * srcStride; |
| 829 | - __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + stride + | 832 | + __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + stride + i * outerLoopStride + |
| 830 | - i * outerLoopStride + mainFoldLoopTimes * srcStride; | 833 | + mainFoldLoopTimes * srcStride; |
| 831 | NlastDichotomyAdd<COMPRESSION, TailCount>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, | 834 | NlastDichotomyAdd<COMPRESSION, TailCount>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride, |
| 832 | foldOffset); | 835 | foldOffset); |
| 833 | - DataCopy((__local_mem__ float*)dst, aReg, pMask); | 836 | + StoreAlign((__ubuf__ float*)dst, aReg, pMask); |
| 834 | } | 837 | } |
| 835 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { | 838 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { |
| 836 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + unFoldDstOffset + | 839 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + unFoldDstOffset + |
| 837 | - i * outerLoopStride + j * innerLoopStride; | 840 | + i * outerLoopStride + j * innerLoopStride; |
| 838 | - __local_mem__ float* srcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + | 841 | + __ubuf__ float* srcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + i * outerLoopStride + |
| 839 | - i * outerLoopStride + j * srcStride; | 842 | + j * srcStride; |
| 840 | - __local_mem__ float* srcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + stride + | 843 | + __ubuf__ float* srcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset + stride + |
| 841 | - i * outerLoopStride + j * srcStride; | 844 | + i * outerLoopStride + j * srcStride; |
| 842 | NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride); | 845 | NlastDichotomyAdd<COMPRESSION>::LoadAndAccumulate(aReg, srcA, srcB, pMask, stride); |
| 843 | - DataCopy((__local_mem__ float*)dst, aReg, pMask); | 846 | + StoreAlign((__ubuf__ float*)dst, aReg, pMask); |
| 844 | } | 847 | } |
| 845 | } | 848 | } |
| 846 | } | 849 | } |
| @@ -901,20 +904,20 @@ __aicore__ inline void LogSoftmaxGradOpsBase::UpdateCache(const LocalTensor<floa | |||
| 901 | uint32_t innerLoopStride = stride; | 904 | uint32_t innerLoopStride = stride; |
| 902 | __VEC_SCOPE__ | 905 | __VEC_SCOPE__ |
| 903 | { | 906 | { |
| 904 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 907 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 905 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 908 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 906 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 909 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 907 | uint32_t sreg = static_cast<uint32_t>(count); | 910 | uint32_t sreg = static_cast<uint32_t>(count); |
| 908 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 911 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 909 | AscendC::MicroAPI::MaskReg pMask; | 912 | AscendC::MicroAPI::MaskReg pMask; |
| 910 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 913 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 911 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 914 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 912 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 915 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 913 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 916 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 914 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride); | 917 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride); |
| 915 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 918 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 916 | } | 919 | } |
| 917 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 920 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 918 | } | 921 | } |
| 919 | } | 922 | } |
| 920 | } | 923 | } |
| @@ -933,24 +936,24 @@ __aicore__ inline void LogSoftmaxGradOpsBase::Normalize(const LocalTensor<float> | |||
| 933 | uint32_t innerLoopStride = VL_FP32; | 936 | uint32_t innerLoopStride = VL_FP32; |
| 934 | __VEC_SCOPE__ | 937 | __VEC_SCOPE__ |
| 935 | { | 938 | { |
| 936 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 939 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 937 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 940 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 938 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 941 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 939 | - __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 942 | + __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 940 | uint32_t count; | 943 | uint32_t count; |
| 941 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 944 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 942 | AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg; | 945 | AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg; |
| 943 | AscendC::MicroAPI::MaskReg pMask; | 946 | AscendC::MicroAPI::MaskReg pMask; |
| 944 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 947 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 945 | count = static_cast<uint32_t>(colSize); | 948 | count = static_cast<uint32_t>(colSize); |
| 946 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 949 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 947 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i); | 950 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i); |
| 948 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 951 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 949 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 952 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 950 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride + j * innerLoopStride); | 953 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride + j * innerLoopStride); |
| 951 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); | 954 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); |
| 952 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); | 955 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); |
| 953 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); | 956 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); |
| 954 | } | 957 | } |
| 955 | } | 958 | } |
| 956 | } | 959 | } |
| @@ -30,6 +30,10 @@ | |||
| 30 | namespace LogSoftmaxV2Ops { | 30 | namespace LogSoftmaxV2Ops { |
| 31 | using namespace AscendC; | 31 | using namespace AscendC; |
| 32 | using namespace SoftmaxV2Ops; | 32 | using namespace SoftmaxV2Ops; |
| 33 | +using AscendC::Reg::LoadAlign; | ||
| 34 | +using AscendC::Reg::Move; | ||
| 35 | +using AscendC::Reg::Reduce; | ||
| 36 | +using AscendC::Reg::StoreAlign; | ||
| 33 | 37 | ||
| 34 | constexpr uint32_t DOUBLE_BUFFER = 2; | 38 | constexpr uint32_t DOUBLE_BUFFER = 2; |
| 35 | constexpr uint32_t BLOCK_SIZE = 32; // 32B | 39 | constexpr uint32_t BLOCK_SIZE = 32; // 32B |
| @@ -45,8 +49,7 @@ public: | |||
| 45 | private: | 49 | private: |
| 46 | __aicore__ inline void ProcessUB(int64_t ubA, int64_t aOffset); | 50 | __aicore__ inline void ProcessUB(int64_t ubA, int64_t aOffset); |
| 47 | 51 | ||
| 48 | - __aicore__ inline void FirstNormCompute(int64_t ubA, __local_mem__ T_in* xInAddr, | 52 | + __aicore__ inline void FirstNormCompute(int64_t ubA, __ubuf__ T_in* xInAddr, __ubuf__ float* xTmpLocalAddr); |
| 49 | - __local_mem__ float* xTmpLocalAddr); | ||
| 50 | __aicore__ inline void SecondNormCompute(const LocalTensor<T_out>& dstTensor, const LocalTensor<float>& srcTensor, | 53 | __aicore__ inline void SecondNormCompute(const LocalTensor<T_out>& dstTensor, const LocalTensor<float>& srcTensor, |
| 51 | const LocalTensor<float>& reduceSumTempTensor, const int64_t aSize, | 54 | const LocalTensor<float>& reduceSumTempTensor, const int64_t aSize, |
| 52 | const int64_t rSize, const int64_t stride); | 55 | const int64_t rSize, const int64_t stride); |
| @@ -60,9 +63,9 @@ private: | |||
| 60 | const int64_t rSize, const int64_t stride); | 63 | const int64_t rSize, const int64_t stride); |
| 61 | __aicore__ inline void CopyInX(const LocalTensor<T_in>& xInUb, int64_t ubA, int64_t offset); | 64 | __aicore__ inline void CopyInX(const LocalTensor<T_in>& xInUb, int64_t ubA, int64_t offset); |
| 62 | __aicore__ inline void CopyOutY(const LocalTensor<T_out>& yOutUb, int64_t ubA, int64_t offset); | 65 | __aicore__ inline void CopyOutY(const LocalTensor<T_out>& yOutUb, int64_t ubA, int64_t offset); |
| 63 | - __aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T_in* src, AscendC::MicroAPI::RegTensor<float>& dst, | 66 | + __aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T_in* src, AscendC::MicroAPI::RegTensor<float>& dst, |
| 64 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); | 67 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); |
| 65 | - __aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T_out* dst, AscendC::MicroAPI::RegTensor<float>& src, | 68 | + __aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T_out* dst, AscendC::MicroAPI::RegTensor<float>& src, |
| 66 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); | 69 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); |
| 67 | 70 | ||
| 68 | private: | 71 | private: |
| @@ -130,9 +133,9 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::ProcessUB(int64_t ubA, int64 | |||
| 130 | LocalTensor<float> binaryTmpLocalTensor = binaryTmpLocalBuffer_.AllocTensor<float>(); | 133 | LocalTensor<float> binaryTmpLocalTensor = binaryTmpLocalBuffer_.AllocTensor<float>(); |
| 131 | LocalTensor<float> xTmpLocalTensor = xTmpLocalBuffer_.AllocTensor<float>(); | 134 | LocalTensor<float> xTmpLocalTensor = xTmpLocalBuffer_.AllocTensor<float>(); |
| 132 | 135 | ||
| 133 | - __local_mem__ T_in* xInUbAddr = (__local_mem__ T_in*)xInUb.GetPhyAddr(); | 136 | + __ubuf__ T_in* xInUbAddr = (__ubuf__ T_in*)xInUb.GetPhyAddr(); |
| 134 | - __local_mem__ float* xTmpLocalAddr = (__local_mem__ float*)xTmpLocalTensor.GetPhyAddr(); | 137 | + __ubuf__ float* xTmpLocalAddr = (__ubuf__ float*)xTmpLocalTensor.GetPhyAddr(); |
| 135 | - __local_mem__ float* binaryTmpLocalAddr = (__local_mem__ float*)binaryTmpLocalTensor.GetPhyAddr(); | 138 | + __ubuf__ float* binaryTmpLocalAddr = (__ubuf__ float*)binaryTmpLocalTensor.GetPhyAddr(); |
| 136 | 139 | ||
| 137 | FirstNormCompute(ubA, xInUbAddr, xTmpLocalAddr); | 140 | FirstNormCompute(ubA, xInUbAddr, xTmpLocalAddr); |
| 138 | xQueue_.FreeTensor<T_in>(xInUb); | 141 | xQueue_.FreeTensor<T_in>(xInUb); |
| @@ -146,8 +149,8 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::ProcessUB(int64_t ubA, int64 | |||
| 146 | } | 149 | } |
| 147 | 150 | ||
| 148 | template <typename T_in, typename T_out> | 151 | template <typename T_in, typename T_out> |
| 149 | -__aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA, __local_mem__ T_in* xInAddr, | 152 | +__aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA, __ubuf__ T_in* xInAddr, |
| 150 | - __local_mem__ float* xTmpLocalAddr) | 153 | + __ubuf__ float* xTmpLocalAddr) |
| 151 | { | 154 | { |
| 152 | int64_t rAligned = tl_->rAligned; | 155 | int64_t rAligned = tl_->rAligned; |
| 153 | int64_t ubActualR = tl_->r; | 156 | int64_t ubActualR = tl_->r; |
| @@ -180,7 +183,7 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA | |||
| 180 | AscendC::MicroAPI::Duplicate(vreg0, static_cast<float>(-INFINITY), maskAll); | 183 | AscendC::MicroAPI::Duplicate(vreg0, static_cast<float>(-INFINITY), maskAll); |
| 181 | LoadTensorForDtypeTIn(xInAddr, vreg1, maskTail, tailAddrPtr); | 184 | LoadTensorForDtypeTIn(xInAddr, vreg1, maskTail, tailAddrPtr); |
| 182 | AscendC::MicroAPI::Max(vreg1, vreg0, vreg1, maskTail); | 185 | AscendC::MicroAPI::Max(vreg1, vreg0, vreg1, maskTail); |
| 183 | - AscendC::MicroAPI::Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vreg0, vreg1, maskTail); | 186 | + AscendC::MicroAPI::Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vreg0, vreg1, maskTail); |
| 184 | 187 | ||
| 185 | uint32_t sreg2 = static_cast<uint32_t>(ubActualR - tailUbBlockSize); | 188 | uint32_t sreg2 = static_cast<uint32_t>(ubActualR - tailUbBlockSize); |
| 186 | for (uint16_t i = 0; i < rLoopCountTmp; i++) { | 189 | for (uint16_t i = 0; i < rLoopCountTmp; i++) { |
| @@ -189,7 +192,7 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA | |||
| 189 | LoadTensorForDtypeTIn(xInAddr, vreg1, maskMax, addrPtr); | 192 | LoadTensorForDtypeTIn(xInAddr, vreg1, maskMax, addrPtr); |
| 190 | AscendC::MicroAPI::Max(vreg0, vreg0, vreg1, maskMax); | 193 | AscendC::MicroAPI::Max(vreg0, vreg0, vreg1, maskMax); |
| 191 | } | 194 | } |
| 192 | - AscendC::MicroAPI::ReduceMax(vreg2, vreg0, maskAll); | 195 | + AscendC::MicroAPI::Reduce<ReduceType::MAX>(vreg2, vreg0, maskAll); |
| 193 | AscendC::MicroAPI::Duplicate(vreg3, vreg2, maskAll); | 196 | AscendC::MicroAPI::Duplicate(vreg3, vreg2, maskAll); |
| 194 | 197 | ||
| 195 | // 求sub | 198 | // 求sub |
| @@ -199,7 +202,7 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::FirstNormCompute(int64_t ubA | |||
| 199 | uint32_t addrPtr = i * VL_FP32 + k * rAligned; | 202 | uint32_t addrPtr = i * VL_FP32 + k * rAligned; |
| 200 | LoadTensorForDtypeTIn(xInAddr, vreg1, maskSub, addrPtr); | 203 | LoadTensorForDtypeTIn(xInAddr, vreg1, maskSub, addrPtr); |
| 201 | AscendC::MicroAPI::Sub(vreg4, vreg1, vreg3, maskSub); | 204 | AscendC::MicroAPI::Sub(vreg4, vreg1, vreg3, maskSub); |
| 202 | - AscendC::MicroAPI::DataCopy(((__local_mem__ float*)xTmpLocalAddr + addrPtr), vreg4, maskSub); | 205 | + AscendC::MicroAPI::StoreAlign(((__ubuf__ float*)xTmpLocalAddr + addrPtr), vreg4, maskSub); |
| 203 | } | 206 | } |
| 204 | } | 207 | } |
| 205 | } | 208 | } |
| @@ -227,9 +230,9 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::SecondNormComputePost(const | |||
| 227 | uint16_t oriR = tl_->r; | 230 | uint16_t oriR = tl_->r; |
| 228 | uint16_t oriRAligned = tl_->rAligned; | 231 | uint16_t oriRAligned = tl_->rAligned; |
| 229 | if (rSize <= VL_FP32) { | 232 | if (rSize <= VL_FP32) { |
| 230 | - __local_mem__ T_out* dst = (__local_mem__ T_out*)dstTensor.GetPhyAddr(); | 233 | + __ubuf__ T_out* dst = (__ubuf__ T_out*)dstTensor.GetPhyAddr(); |
| 231 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 234 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 232 | - __local_mem__ float* oriSrc = (__local_mem__ float*)oriSrcTensor.GetPhyAddr(); | 235 | + __ubuf__ float* oriSrc = (__ubuf__ float*)oriSrcTensor.GetPhyAddr(); |
| 233 | 236 | ||
| 234 | __VEC_SCOPE__ | 237 | __VEC_SCOPE__ |
| 235 | { | 238 | { |
| @@ -240,25 +243,25 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::SecondNormComputePost(const | |||
| 240 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 243 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 241 | AscendC::MicroAPI::MaskReg maskOri; | 244 | AscendC::MicroAPI::MaskReg maskOri; |
| 242 | for (uint16_t i = 0; i < loopTimes; ++i) { | 245 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 243 | - DataCopy(aReg, (__local_mem__ float*)src + i * static_cast<uint32_t>(stride)); | 246 | + LoadAlign(aReg, (__ubuf__ float*)src + i * static_cast<uint32_t>(stride)); |
| 244 | - ReduceSum(bReg, aReg, pMask); | 247 | + Reduce<ReduceType::SUM>(bReg, aReg, pMask); |
| 245 | Duplicate(cReg, bReg, pFull); | 248 | Duplicate(cReg, bReg, pFull); |
| 246 | Log(cReg, cReg, pFull); | 249 | Log(cReg, cReg, pFull); |
| 247 | uint32_t sreg0 = static_cast<uint32_t>(oriR); | 250 | uint32_t sreg0 = static_cast<uint32_t>(oriR); |
| 248 | for (uint16_t j = 0; j < rLoopCount; ++j) { | 251 | for (uint16_t j = 0; j < rLoopCount; ++j) { |
| 249 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); | 252 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); |
| 250 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; | 253 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; |
| 251 | - DataCopy(dReg, ((__local_mem__ float*)oriSrc + addrPtr)); | 254 | + LoadAlign(dReg, ((__ubuf__ float*)oriSrc + addrPtr)); |
| 252 | Sub(dReg, dReg, cReg, maskOri); | 255 | Sub(dReg, dReg, cReg, maskOri); |
| 253 | StoreTensorForDtypeTOut(dst, dReg, maskOri, addrPtr); | 256 | StoreTensorForDtypeTOut(dst, dReg, maskOri, addrPtr); |
| 254 | } | 257 | } |
| 255 | } | 258 | } |
| 256 | } | 259 | } |
| 257 | } else { | 260 | } else { |
| 258 | - __local_mem__ T_out* dst = (__local_mem__ T_out*)dstTensor.GetPhyAddr(); | 261 | + __ubuf__ T_out* dst = (__ubuf__ T_out*)dstTensor.GetPhyAddr(); |
| 259 | - __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 262 | + __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 260 | - __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32; | 263 | + __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32; |
| 261 | - __local_mem__ float* oriSrc = (__local_mem__ float*)oriSrcTensor.GetPhyAddr(); | 264 | + __ubuf__ float* oriSrc = (__ubuf__ float*)oriSrcTensor.GetPhyAddr(); |
| 262 | 265 | ||
| 263 | __VEC_SCOPE__ | 266 | __VEC_SCOPE__ |
| 264 | { | 267 | { |
| @@ -269,18 +272,18 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::SecondNormComputePost(const | |||
| 269 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 272 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 270 | AscendC::MicroAPI::MaskReg maskOri; | 273 | AscendC::MicroAPI::MaskReg maskOri; |
| 271 | for (uint16_t i = 0; i < loopTimes; ++i) { | 274 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 272 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * static_cast<uint32_t>(stride)); | 275 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * static_cast<uint32_t>(stride)); |
| 273 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * static_cast<uint32_t>(stride)); | 276 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * static_cast<uint32_t>(stride)); |
| 274 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 277 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 275 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 278 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 276 | - ReduceSum(bReg, aReg, pFull); | 279 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 277 | Duplicate(dReg, bReg, pFull); | 280 | Duplicate(dReg, bReg, pFull); |
| 278 | Log(dReg, dReg, pFull); | 281 | Log(dReg, dReg, pFull); |
| 279 | uint32_t sreg0 = static_cast<uint32_t>(oriR); | 282 | uint32_t sreg0 = static_cast<uint32_t>(oriR); |
| 280 | for (uint16_t j = 0; j < rLoopCount; ++j) { | 283 | for (uint16_t j = 0; j < rLoopCount; ++j) { |
| 281 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); | 284 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); |
| 282 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; | 285 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; |
| 283 | - DataCopy(eReg, ((__local_mem__ float*)oriSrc + addrPtr)); | 286 | + LoadAlign(eReg, ((__ubuf__ float*)oriSrc + addrPtr)); |
| 284 | Sub(eReg, eReg, dReg, maskOri); | 287 | Sub(eReg, eReg, dReg, maskOri); |
| 285 | StoreTensorForDtypeTOut(dst, eReg, maskOri, addrPtr); | 288 | StoreTensorForDtypeTOut(dst, eReg, maskOri, addrPtr); |
| 286 | } | 289 | } |
| @@ -309,9 +312,9 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::SecondNormComputeWithExpPost | |||
| 309 | uint16_t oriR = tl_->r; | 312 | uint16_t oriR = tl_->r; |
| 310 | uint16_t oriRAligned = tl_->rAligned; | 313 | uint16_t oriRAligned = tl_->rAligned; |
| 311 | if (rSize <= VL_FP32) { | 314 | if (rSize <= VL_FP32) { |
| 312 | - __local_mem__ T_out* dst = (__local_mem__ T_out*)dstTensor.GetPhyAddr(); | 315 | + __ubuf__ T_out* dst = (__ubuf__ T_out*)dstTensor.GetPhyAddr(); |
| 313 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 316 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 314 | - __local_mem__ float* oriSrc = (__local_mem__ float*)oriSrcTensor.GetPhyAddr(); | 317 | + __ubuf__ float* oriSrc = (__ubuf__ float*)oriSrcTensor.GetPhyAddr(); |
| 315 | 318 | ||
| 316 | __VEC_SCOPE__ | 319 | __VEC_SCOPE__ |
| 317 | { | 320 | { |
| @@ -322,26 +325,26 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::SecondNormComputeWithExpPost | |||
| 322 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 325 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 323 | AscendC::MicroAPI::MaskReg maskOri; | 326 | AscendC::MicroAPI::MaskReg maskOri; |
| 324 | for (uint16_t i = 0; i < loopTimes; ++i) { | 327 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 325 | - DataCopy(aReg, (__local_mem__ float*)src + i * stride); | 328 | + LoadAlign(aReg, (__ubuf__ float*)src + i * stride); |
| 326 | Exp(aReg, aReg, pMask); | 329 | Exp(aReg, aReg, pMask); |
| 327 | - ReduceSum(bReg, aReg, pMask); | 330 | + Reduce<ReduceType::SUM>(bReg, aReg, pMask); |
| 328 | Duplicate(cReg, bReg, pFull); | 331 | Duplicate(cReg, bReg, pFull); |
| 329 | Log(cReg, cReg, pFull); | 332 | Log(cReg, cReg, pFull); |
| 330 | uint32_t sreg0 = static_cast<uint32_t>(oriR); | 333 | uint32_t sreg0 = static_cast<uint32_t>(oriR); |
| 331 | for (uint16_t j = 0; j < rLoopCount; ++j) { | 334 | for (uint16_t j = 0; j < rLoopCount; ++j) { |
| 332 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); | 335 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); |
| 333 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; | 336 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; |
| 334 | - DataCopy(dReg, ((__local_mem__ float*)oriSrc + addrPtr)); | 337 | + LoadAlign(dReg, ((__ubuf__ float*)oriSrc + addrPtr)); |
| 335 | Sub(dReg, dReg, cReg, maskOri); | 338 | Sub(dReg, dReg, cReg, maskOri); |
| 336 | StoreTensorForDtypeTOut(dst, dReg, maskOri, addrPtr); | 339 | StoreTensorForDtypeTOut(dst, dReg, maskOri, addrPtr); |
| 337 | } | 340 | } |
| 338 | } | 341 | } |
| 339 | } | 342 | } |
| 340 | } else { | 343 | } else { |
| 341 | - __local_mem__ T_out* dst = (__local_mem__ T_out*)dstTensor.GetPhyAddr(); | 344 | + __ubuf__ T_out* dst = (__ubuf__ T_out*)dstTensor.GetPhyAddr(); |
| 342 | - __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 345 | + __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 343 | - __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32; | 346 | + __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32; |
| 344 | - __local_mem__ float* oriSrc = (__local_mem__ float*)oriSrcTensor.GetPhyAddr(); | 347 | + __ubuf__ float* oriSrc = (__ubuf__ float*)oriSrcTensor.GetPhyAddr(); |
| 345 | __VEC_SCOPE__ | 348 | __VEC_SCOPE__ |
| 346 | { | 349 | { |
| 347 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); | 350 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); |
| @@ -351,20 +354,20 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::SecondNormComputeWithExpPost | |||
| 351 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 354 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 352 | AscendC::MicroAPI::MaskReg maskOri; | 355 | AscendC::MicroAPI::MaskReg maskOri; |
| 353 | for (uint16_t i = 0; i < loopTimes; ++i) { | 356 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 354 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * stride); | 357 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * stride); |
| 355 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * stride); | 358 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * stride); |
| 356 | Exp(aReg, aReg, pFull); | 359 | Exp(aReg, aReg, pFull); |
| 357 | Exp(bReg, bReg, pMask); | 360 | Exp(bReg, bReg, pMask); |
| 358 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 361 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 359 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 362 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 360 | - ReduceSum(bReg, aReg, pFull); | 363 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 361 | Duplicate(dReg, bReg, pFull); | 364 | Duplicate(dReg, bReg, pFull); |
| 362 | Log(dReg, dReg, pFull); | 365 | Log(dReg, dReg, pFull); |
| 363 | uint32_t sreg0 = static_cast<uint32_t>(oriR); | 366 | uint32_t sreg0 = static_cast<uint32_t>(oriR); |
| 364 | for (uint16_t j = 0; j < rLoopCount; ++j) { | 367 | for (uint16_t j = 0; j < rLoopCount; ++j) { |
| 365 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); | 368 | maskOri = AscendC::MicroAPI::UpdateMask<float>(sreg0); |
| 366 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; | 369 | uint32_t addrPtr = j * VL_FP32 + i * oriRAligned; |
| 367 | - DataCopy(eReg, ((__local_mem__ float*)oriSrc + addrPtr)); | 370 | + LoadAlign(eReg, ((__ubuf__ float*)oriSrc + addrPtr)); |
| 368 | Sub(eReg, eReg, dReg, maskOri); | 371 | Sub(eReg, eReg, dReg, maskOri); |
| 369 | StoreTensorForDtypeTOut(dst, eReg, maskOri, addrPtr); | 372 | StoreTensorForDtypeTOut(dst, eReg, maskOri, addrPtr); |
| 370 | } | 373 | } |
| @@ -412,83 +415,83 @@ __aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::SecondNormCompute(const Loca | |||
| 412 | int64_t tailSrcBOffset = floorVLCount * VL_FP32; | 415 | int64_t tailSrcBOffset = floorVLCount * VL_FP32; |
| 413 | int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32; | 416 | int64_t unFoldSrcOffset = (mainFoldLoopTimes + tailFoldLoopTimes) * VL_FP32; |
| 414 | 417 | ||
| 415 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr(); | 418 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr(); |
| 416 | - __local_mem__ float* foldSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 419 | + __ubuf__ float* foldSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 417 | - __local_mem__ float* foldSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; | 420 | + __ubuf__ float* foldSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; |
| 418 | - __local_mem__ float* tailSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; | 421 | + __ubuf__ float* tailSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; |
| 419 | - __local_mem__ float* tailSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; | 422 | + __ubuf__ float* tailSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; |
| 420 | - __local_mem__ float* unFoldSrc = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; | 423 | + __ubuf__ float* unFoldSrc = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; |
| 421 | 424 | ||
| 422 | __VEC_SCOPE__ | 425 | __VEC_SCOPE__ |
| 423 | { | 426 | { |
| 424 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 427 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 425 | - AscendC::MicroAPI::UnalignReg UReg; | 428 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 426 | 429 | ||
| 427 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 430 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 428 | - dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; | 431 | + dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; |
| 429 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { | 432 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { |
| 430 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; | 433 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; |
| 431 | - DataCopy(aReg, (__local_mem__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); | 434 | + LoadAlign(aReg, (__ubuf__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); |
| 432 | - DataCopy(bReg, (__local_mem__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); | 435 | + LoadAlign(bReg, (__ubuf__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); |
| 433 | Exp(aReg, aReg, pFull); | 436 | Exp(aReg, aReg, pFull); |
| 434 | Exp(bReg, bReg, pFull); | 437 | Exp(bReg, bReg, pFull); |
| 435 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); | 438 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); |
| 436 | - ReduceSum(dReg, cReg, pFull); | 439 | + Reduce<ReduceType::SUM>(dReg, cReg, pFull); |
| 437 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, dReg, UReg, 1); | 440 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, dReg, UReg, 1); |
| 438 | } | 441 | } |
| 439 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { | 442 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { |
| 440 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); | 443 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); |
| 441 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 444 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 442 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 445 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 443 | - DataCopy(aReg, (__local_mem__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); | 446 | + LoadAlign(aReg, (__ubuf__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); |
| 444 | - DataCopy(bReg, (__local_mem__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); | 447 | + LoadAlign(bReg, (__ubuf__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); |
| 445 | Exp(aReg, aReg, pFull); | 448 | Exp(aReg, aReg, pFull); |
| 446 | Exp(bReg, bReg, pFull); | 449 | Exp(bReg, bReg, pFull); |
| 447 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 450 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 448 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 451 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 449 | - ReduceSum(bReg, aReg, pFull); | 452 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 450 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 453 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 451 | } | 454 | } |
| 452 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { | 455 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { |
| 453 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 456 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 454 | - DataCopy(aReg, (__local_mem__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); | 457 | + LoadAlign(aReg, (__ubuf__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); |
| 455 | Exp(aReg, aReg, pFull); | 458 | Exp(aReg, aReg, pFull); |
| 456 | - ReduceSum(bReg, aReg, pFull); | 459 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 457 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 460 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 458 | } | 461 | } |
| 459 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 462 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 460 | } | 463 | } |
| 461 | } | 464 | } |
| 462 | SecondNormComputePost(dstTensor, reduceSumTempTensor, srcTensor, aSize, foldPoint, outerLoopDstStride); | 465 | SecondNormComputePost(dstTensor, reduceSumTempTensor, srcTensor, aSize, foldPoint, outerLoopDstStride); |
| 463 | } | 466 | } |
| 464 | 467 | ||
| 465 | template <typename T_in, typename T_out> | 468 | template <typename T_in, typename T_out> |
| 466 | -__aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::LoadTensorForDtypeTIn(__local_mem__ T_in* src, | 469 | +__aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::LoadTensorForDtypeTIn(__ubuf__ T_in* src, |
| 467 | AscendC::MicroAPI::RegTensor<float>& dst, | 470 | AscendC::MicroAPI::RegTensor<float>& dst, |
| 468 | AscendC::MicroAPI::MaskReg& preg, | 471 | AscendC::MicroAPI::MaskReg& preg, |
| 469 | uint32_t offset) | 472 | uint32_t offset) |
| 470 | { | 473 | { |
| 471 | if constexpr (IsSameType<T_in, float>::value) { | 474 | if constexpr (IsSameType<T_in, float>::value) { |
| 472 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(dst, src + offset); | 475 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(dst, src + offset); |
| 473 | } else { | 476 | } else { |
| 474 | AscendC::MicroAPI::RegTensor<T_in> xFp16; | 477 | AscendC::MicroAPI::RegTensor<T_in> xFp16; |
| 475 | - DataCopy<T_in, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, src + offset); | 478 | + LoadAlign<T_in, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, src + offset); |
| 476 | Cast<float, T_in, castTraitFp16ToFp32>(dst, xFp16, preg); | 479 | Cast<float, T_in, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 477 | } | 480 | } |
| 478 | } | 481 | } |
| 479 | 482 | ||
| 480 | template <typename T_in, typename T_out> | 483 | template <typename T_in, typename T_out> |
| 481 | -__aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::StoreTensorForDtypeTOut(__local_mem__ T_out* dst, | 484 | +__aicore__ inline void LogSoftmaxV2AR<T_in, T_out>::StoreTensorForDtypeTOut(__ubuf__ T_out* dst, |
| 482 | AscendC::MicroAPI::RegTensor<float>& src, | 485 | AscendC::MicroAPI::RegTensor<float>& src, |
| 483 | AscendC::MicroAPI::MaskReg& preg, | 486 | AscendC::MicroAPI::MaskReg& preg, |
| 484 | uint32_t offset) | 487 | uint32_t offset) |
| 485 | { | 488 | { |
| 486 | if constexpr (IsSameType<T_out, float>::value) { | 489 | if constexpr (IsSameType<T_out, float>::value) { |
| 487 | - DataCopy<T_out, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); | 490 | + StoreAlign<T_out, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); |
| 488 | } else { | 491 | } else { |
| 489 | AscendC::MicroAPI::RegTensor<T_out> xFp16; | 492 | AscendC::MicroAPI::RegTensor<T_out> xFp16; |
| 490 | Cast<T_out, float, castTraitFp32ToFp16>(xFp16, src, preg); | 493 | Cast<T_out, float, castTraitFp32ToFp16>(xFp16, src, preg); |
| 491 | - DataCopy<T_out, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); | 494 | + StoreAlign<T_out, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); |
| 492 | } | 495 | } |
| 493 | } | 496 | } |
| 494 | 497 | ||
| @@ -44,15 +44,14 @@ public: | |||
| 44 | __aicore__ inline void Process(); | 44 | __aicore__ inline void Process(); |
| 45 | 45 | ||
| 46 | private: | 46 | private: |
| 47 | - __aicore__ inline void CalculateMaxVF(__local_mem__ float*& xMaxPtr, __local_mem__ Tx*& xPtr, uint32_t aSize, | 47 | + __aicore__ inline void CalculateMaxVF(__ubuf__ float*& xMaxPtr, __ubuf__ Tx*& xPtr, uint32_t aSize, |
| 48 | uint32_t ubFactor); | 48 | uint32_t ubFactor); |
| 49 | - __aicore__ inline void CalculateOutVF(__local_mem__ Ty*& yPtr, __local_mem__ Tx*& xPtr, | 49 | + __aicore__ inline void CalculateOutVF(__ubuf__ Ty*& yPtr, __ubuf__ Tx*& xPtr, __ubuf__ float*& xMaxPtr, |
| 50 | - __local_mem__ float*& xMaxPtr, __local_mem__ float*& xSumPtr, uint32_t a, | 50 | + __ubuf__ float*& xSumPtr, uint32_t a, uint32_t ubFactor); |
| 51 | - uint32_t ubFactor); | 51 | + __aicore__ inline void MainBlockCastSubExpVF(__ubuf__ float*& xFp32Ptr, __ubuf__ Tx*& xPtr, |
| 52 | - __aicore__ inline void MainBlockCastSubExpVF(__local_mem__ float*& xFp32Ptr, __local_mem__ Tx*& xPtr, | 52 | + __ubuf__ float*& xMaxPtr, uint32_t a, uint32_t ubFactor); |
| 53 | - __local_mem__ float*& xMaxPtr, uint32_t a, uint32_t ubFactor); | 53 | + __aicore__ inline void FoldBlockCastSubExpVF(__ubuf__ float*& dstPtr, __ubuf__ Tx*& xPtr, __ubuf__ float*& xMaxPtr, |
| 54 | - __aicore__ inline void FoldBlockCastSubExpVF(__local_mem__ float*& dstPtr, __local_mem__ Tx*& xPtr, | 54 | + uint32_t a, uint32_t ubFactor); |
| 55 | - __local_mem__ float*& xMaxPtr, uint32_t a, uint32_t ubFactor); | ||
| 56 | __aicore__ inline int64_t GetCacheId(const int64_t idx); | 55 | __aicore__ inline int64_t GetCacheId(const int64_t idx); |
| 57 | __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor, | 56 | __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const LocalTensor<float>& srcTensor, |
| 58 | const int64_t cacheId, const int64_t stride, const int64_t count); | 57 | const int64_t cacheId, const int64_t stride, const int64_t count); |
| @@ -136,7 +135,7 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 136 | x1DataCopyExtParams.srcStride = 0; | 135 | x1DataCopyExtParams.srcStride = 0; |
| 137 | x1DataCopyExtParams.dstStride = 0; | 136 | x1DataCopyExtParams.dstStride = 0; |
| 138 | 137 | ||
| 139 | - __local_mem__ float* xMaxPtr = (__local_mem__ float*)xMaxLocal.GetPhyAddr(); | 138 | + __ubuf__ float* xMaxPtr = (__ubuf__ float*)xMaxLocal.GetPhyAddr(); |
| 140 | // 每行：对R循环，求整行R的最大值 | 139 | // 每行：对R循环，求整行R的最大值 |
| 141 | for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) { | 140 | for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) { |
| 142 | int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; // 每个UB循环的偏移量 | 141 | int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; // 每个UB循环的偏移量 |
| @@ -151,14 +150,14 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 151 | xQueue_.EnQue<Tx>(xLocal); | 150 | xQueue_.EnQue<Tx>(xLocal); |
| 152 | xLocal = xQueue_.DeQue<Tx>(); | 151 | xLocal = xQueue_.DeQue<Tx>(); |
| 153 | 152 | ||
| 154 | - __local_mem__ Tx* xPtr = (__local_mem__ Tx*)xLocal.GetPhyAddr(); | 153 | + __ubuf__ Tx* xPtr = (__ubuf__ Tx*)xLocal.GetPhyAddr(); |
| 155 | CalculateMaxVF(xMaxPtr, xPtr, A_IN_IN, ubFactor); | 154 | CalculateMaxVF(xMaxPtr, xPtr, A_IN_IN, ubFactor); |
| 156 | xQueue_.FreeTensor(xLocal); | 155 | xQueue_.FreeTensor(xLocal); |
| 157 | } | 156 | } |
| 158 | 157 | ||
| 159 | // UB间二分累加：计算每行的Σe^(x - max) | 158 | // UB间二分累加：计算每行的Σe^(x - max) |
| 160 | LocalTensor<float> xTmpLocal = xTmpBuffer.Get<float>(); | 159 | LocalTensor<float> xTmpLocal = xTmpBuffer.Get<float>(); |
| 161 | - __local_mem__ float* xTmpFp32Ptr = (__local_mem__ float*)xTmpLocal.GetPhyAddr(); | 160 | + __ubuf__ float* xTmpFp32Ptr = (__ubuf__ float*)xTmpLocal.GetPhyAddr(); |
| 162 | 161 | ||
| 163 | x1DataCopyExtParams.blockLen = tl_->ubFactor * sizeof(Tx); | 162 | x1DataCopyExtParams.blockLen = tl_->ubFactor * sizeof(Tx); |
| 164 | 163 | ||
| @@ -179,14 +178,14 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 179 | xQueue_.EnQue<Tx>(x1Local); | 178 | xQueue_.EnQue<Tx>(x1Local); |
| 180 | x1Local = xQueue_.DeQue<Tx>(); | 179 | x1Local = xQueue_.DeQue<Tx>(); |
| 181 | 180 | ||
| 182 | - __local_mem__ Tx* x1Ptr = (__local_mem__ Tx*)x1Local.GetPhyAddr(); | 181 | + __ubuf__ Tx* x1Ptr = (__ubuf__ Tx*)x1Local.GetPhyAddr(); |
| 183 | MainBlockCastSubExpVF(xTmpFp32Ptr, x1Ptr, xMaxPtr, A_IN_IN, tl_->ubFactor); | 182 | MainBlockCastSubExpVF(xTmpFp32Ptr, x1Ptr, xMaxPtr, A_IN_IN, tl_->ubFactor); |
| 184 | xQueue_.FreeTensor(x1Local); | 183 | xQueue_.FreeTensor(x1Local); |
| 185 | 184 | ||
| 186 | // 折叠部分：X2折叠到X1上 | 185 | // 折叠部分：X2折叠到X1上 |
| 187 | if (basicBlockIdx < tl_->mainFoldCount) { | 186 | if (basicBlockIdx < tl_->mainFoldCount) { |
| 188 | LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>(); | 187 | LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>(); |
| 189 | - __local_mem__ Tx* x2Ptr = (__local_mem__ Tx*)x2Local.GetPhyAddr(); | 188 | + __ubuf__ Tx* x2Ptr = (__ubuf__ Tx*)x2Local.GetPhyAddr(); |
| 190 | DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams); | 189 | DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams); |
| 191 | xQueue_.EnQue<Tx>(x2Local); | 190 | xQueue_.EnQue<Tx>(x2Local); |
| 192 | x2Local = xQueue_.DeQue<Tx>(); | 191 | x2Local = xQueue_.DeQue<Tx>(); |
| @@ -195,7 +194,7 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 195 | xQueue_.FreeTensor(x2Local); | 194 | xQueue_.FreeTensor(x2Local); |
| 196 | } else if ((basicBlockIdx == tl_->mainFoldCount) && (tl_->ubFactorTail > 0)) { | 195 | } else if ((basicBlockIdx == tl_->mainFoldCount) && (tl_->ubFactorTail > 0)) { |
| 197 | LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>(); | 196 | LocalTensor<Tx> x2Local = xQueue_.AllocTensor<Tx>(); |
| 198 | - __local_mem__ Tx* x2Ptr = (__local_mem__ Tx*)x2Local.GetPhyAddr(); | 197 | + __ubuf__ Tx* x2Ptr = (__ubuf__ Tx*)x2Local.GetPhyAddr(); |
| 199 | x2DataCopyExtParams.blockLen = tl_->ubFactorTail * sizeof(Tx); // 这里的x2为尾块 | 198 | x2DataCopyExtParams.blockLen = tl_->ubFactorTail * sizeof(Tx); // 这里的x2为尾块 |
| 200 | DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams); | 199 | DataCopyPad(x2Local[0], xGm_[xUbOffset2], x2DataCopyExtParams, padExtParams); |
| 201 | xQueue_.EnQue<Tx>(x2Local); | 200 | xQueue_.EnQue<Tx>(x2Local); |
| @@ -216,7 +215,7 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 216 | // R很小，不需要做UB间二分累加 | 215 | // R很小，不需要做UB间二分累加 |
| 217 | if (tl_->basicBlockLoop == 0) { | 216 | if (tl_->basicBlockLoop == 0) { |
| 218 | LocalTensor<Tx> x1Local = xQueue_.AllocTensor<Tx>(); | 217 | LocalTensor<Tx> x1Local = xQueue_.AllocTensor<Tx>(); |
| 219 | - __local_mem__ Tx* x1Ptr = (__local_mem__ Tx*)x1Local.GetPhyAddr(); | 218 | + __ubuf__ Tx* x1Ptr = (__ubuf__ Tx*)x1Local.GetPhyAddr(); |
| 220 | DataCopyPad(x1Local[0], xGm_[xDimOffset], x1DataCopyExtParams, padExtParams); | 219 | DataCopyPad(x1Local[0], xGm_[xDimOffset], x1DataCopyExtParams, padExtParams); |
| 221 | xQueue_.EnQue<Tx>(x1Local); | 220 | xQueue_.EnQue<Tx>(x1Local); |
| 222 | x1Local = xQueue_.DeQue<Tx>(); | 221 | x1Local = xQueue_.DeQue<Tx>(); |
| @@ -234,7 +233,7 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 234 | yDataCopyExtParams.srcStride = 0; | 233 | yDataCopyExtParams.srcStride = 0; |
| 235 | yDataCopyExtParams.dstStride = 0; | 234 | yDataCopyExtParams.dstStride = 0; |
| 236 | 235 | ||
| 237 | - __local_mem__ float* xSumPtr = (__local_mem__ float*)totalSumLocal_.GetPhyAddr(); | 236 | + __ubuf__ float* xSumPtr = (__ubuf__ float*)totalSumLocal_.GetPhyAddr(); |
| 238 | // 每行：遍历UB块，计算除法 | 237 | // 每行：遍历UB块，计算除法 |
| 239 | for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) { | 238 | for (uint64_t ubIdx = 0; ubIdx < tl_->aLoopCountCeil; ubIdx++) { |
| 240 | int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; | 239 | int64_t xUbOffset = xDimOffset + tl_->ubFactor * ubIdx; |
| @@ -245,8 +244,8 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 245 | 244 | ||
| 246 | LocalTensor<Tx> xLocal = xQueue_.AllocTensor<Tx>(); | 245 | LocalTensor<Tx> xLocal = xQueue_.AllocTensor<Tx>(); |
| 247 | LocalTensor<Ty> yLocal = yQueue_.AllocTensor<Ty>(); | 246 | LocalTensor<Ty> yLocal = yQueue_.AllocTensor<Ty>(); |
| 248 | - __local_mem__ Tx* xPtr = (__local_mem__ Tx*)xLocal.GetPhyAddr(); | 247 | + __ubuf__ Tx* xPtr = (__ubuf__ Tx*)xLocal.GetPhyAddr(); |
| 249 | - __local_mem__ Ty* yPtr = (__local_mem__ Ty*)yLocal.GetPhyAddr(); | 248 | + __ubuf__ Ty* yPtr = (__ubuf__ Ty*)yLocal.GetPhyAddr(); |
| 250 | 249 | ||
| 251 | x1DataCopyExtParams.blockLen = ubFactor * sizeof(Tx); | 250 | x1DataCopyExtParams.blockLen = ubFactor * sizeof(Tx); |
| 252 | DataCopyPad(xLocal[0], xGm_[xUbOffset], x1DataCopyExtParams, padExtParams); | 251 | DataCopyPad(xLocal[0], xGm_[xUbOffset], x1DataCopyExtParams, padExtParams); |
| @@ -266,9 +265,8 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::Process() | |||
| 266 | } | 265 | } |
| 267 | 266 | ||
| 268 | template <typename Tx, typename Ty> | 267 | template <typename Tx, typename Ty> |
| 269 | -__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::CalculateMaxVF(__local_mem__ float*& xMaxPtr, | 268 | +__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::CalculateMaxVF(__ubuf__ float*& xMaxPtr, __ubuf__ Tx*& xPtr, |
| 270 | - __local_mem__ Tx*& xPtr, uint32_t aSize, | 269 | + uint32_t aSize, uint32_t ubFactor) |
| 271 | - uint32_t ubFactor) | ||
| 272 | { | 270 | { |
| 273 | __VEC_SCOPE__ | 271 | __VEC_SCOPE__ |
| 274 | { | 272 | { |
| @@ -292,37 +290,37 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::CalculateMaxVF(__local_m | |||
| 292 | AscendC::MicroAPI::Duplicate(maxReg, CONST_FP32_MIN); | 290 | AscendC::MicroAPI::Duplicate(maxReg, CONST_FP32_MIN); |
| 293 | 291 | ||
| 294 | if constexpr (xToFp32_) { | 292 | if constexpr (xToFp32_) { |
| 295 | - AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr); | 293 | + AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr); |
| 296 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskTail); | 294 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskTail); |
| 297 | } else { | 295 | } else { |
| 298 | - AscendC::MicroAPI::DataCopy(vreg1, xAddr); | 296 | + AscendC::MicroAPI::LoadAlign(vreg1, xAddr); |
| 299 | } | 297 | } |
| 300 | AscendC::MicroAPI::Max(vreg1, maxReg, vreg1, maskTail); | 298 | AscendC::MicroAPI::Max(vreg1, maxReg, vreg1, maskTail); |
| 301 | - AscendC::MicroAPI::Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(maxReg, vreg1, maskTail); | 299 | + AscendC::MicroAPI::Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(maxReg, vreg1, maskTail); |
| 302 | 300 | ||
| 303 | // 整块处理 | 301 | // 整块处理 |
| 304 | for (uint16_t j = 0; j < repeatTimesTmp; j++) { | 302 | for (uint16_t j = 0; j < repeatTimesTmp; j++) { |
| 305 | auto xAddr = xPtr + j * VL_FP32; | 303 | auto xAddr = xPtr + j * VL_FP32; |
| 306 | if constexpr (xToFp32_) { | 304 | if constexpr (xToFp32_) { |
| 307 | - AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr); | 305 | + AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg3, xAddr); |
| 308 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskFull); | 306 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg3, maskFull); |
| 309 | } else { | 307 | } else { |
| 310 | - AscendC::MicroAPI::DataCopy(vreg1, xAddr); | 308 | + AscendC::MicroAPI::LoadAlign(vreg1, xAddr); |
| 311 | } | 309 | } |
| 312 | AscendC::MicroAPI::Max(maxReg, maxReg, vreg1, maskFull); | 310 | AscendC::MicroAPI::Max(maxReg, maxReg, vreg1, maskFull); |
| 313 | } | 311 | } |
| 314 | - AscendC::MicroAPI::DataCopy(vreg2, xMaxPtr); | 312 | + AscendC::MicroAPI::LoadAlign(vreg2, xMaxPtr); |
| 315 | 313 | ||
| 316 | - AscendC::MicroAPI::ReduceMax(maxReg, maxReg, maskFull); | 314 | + AscendC::MicroAPI::Reduce<ReduceType::MAX>(maxReg, maxReg, maskFull); |
| 317 | AscendC::MicroAPI::Max(maxReg, maxReg, vreg2, maskOne); | 315 | AscendC::MicroAPI::Max(maxReg, maxReg, vreg2, maskOne); |
| 318 | - AscendC::MicroAPI::DataCopy(xMaxPtr, maxReg, maskOne); | 316 | + AscendC::MicroAPI::StoreAlign(xMaxPtr, maxReg, maskOne); |
| 319 | } | 317 | } |
| 320 | } | 318 | } |
| 321 | 319 | ||
| 322 | template <typename Tx, typename Ty> | 320 | template <typename Tx, typename Ty> |
| 323 | -__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__local_mem__ Ty*& yPtr, __local_mem__ Tx*& xPtr, | 321 | +__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__ubuf__ Ty*& yPtr, __ubuf__ Tx*& xPtr, |
| 324 | - __local_mem__ float*& xMaxPtr, | 322 | + __ubuf__ float*& xMaxPtr, |
| 325 | - __local_mem__ float*& xSumPtr, uint32_t a, | 323 | + __ubuf__ float*& xSumPtr, uint32_t a, |
| 326 | uint32_t ubFactor) | 324 | uint32_t ubFactor) |
| 327 | { | 325 | { |
| 328 | __VEC_SCOPE__ | 326 | __VEC_SCOPE__ |
| @@ -336,8 +334,8 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__local_m | |||
| 336 | uint32_t width = ubFactor; | 334 | uint32_t width = ubFactor; |
| 337 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); | 335 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); |
| 338 | 336 | ||
| 339 | - AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr); | 337 | + AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr); |
| 340 | - AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, xSumPtr); | 338 | + AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sumReg, xSumPtr); |
| 341 | AscendC::MicroAPI::Ln(lnReg, sumReg, maskAll); | 339 | AscendC::MicroAPI::Ln(lnReg, sumReg, maskAll); |
| 342 | 340 | ||
| 343 | for (uint16_t j = 0; j < repeatTimes; j++) { | 341 | for (uint16_t j = 0; j < repeatTimes; j++) { |
| @@ -346,29 +344,29 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::CalculateOutVF(__local_m | |||
| 346 | auto yAddr = yPtr + j * VL_FP32; | 344 | auto yAddr = yPtr + j * VL_FP32; |
| 347 | 345 | ||
| 348 | if constexpr (xToFp32_) { | 346 | if constexpr (xToFp32_) { |
| 349 | - AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr); | 347 | + AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr); |
| 350 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask); | 348 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask); |
| 351 | } else { | 349 | } else { |
| 352 | - AscendC::MicroAPI::DataCopy(vreg1, xAddr); | 350 | + AscendC::MicroAPI::LoadAlign(vreg1, xAddr); |
| 353 | } | 351 | } |
| 354 | 352 | ||
| 355 | AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask); | 353 | AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask); |
| 356 | AscendC::MicroAPI::Sub(vreg3, vreg2, lnReg, mask); | 354 | AscendC::MicroAPI::Sub(vreg3, vreg2, lnReg, mask); |
| 357 | 355 | ||
| 358 | if constexpr (yToFp32_) { | 356 | if constexpr (yToFp32_) { |
| 359 | - AscendC::MicroAPI::DataCopy(yAddr, vreg3, mask); | 357 | + AscendC::MicroAPI::StoreAlign(yAddr, vreg3, mask); |
| 360 | } else { | 358 | } else { |
| 361 | AscendC::MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(vreg4, vreg3, mask); | 359 | AscendC::MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(vreg4, vreg3, mask); |
| 362 | - AscendC::MicroAPI::DataCopy<Ty, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(yAddr, vreg4, mask); | 360 | + AscendC::MicroAPI::StoreAlign<Ty, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(yAddr, vreg4, mask); |
| 363 | } | 361 | } |
| 364 | } | 362 | } |
| 365 | } | 363 | } |
| 366 | } | 364 | } |
| 367 | 365 | ||
| 368 | template <typename Tx, typename Ty> | 366 | template <typename Tx, typename Ty> |
| 369 | -__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::MainBlockCastSubExpVF(__local_mem__ float*& xFp32Ptr, | 367 | +__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::MainBlockCastSubExpVF(__ubuf__ float*& xFp32Ptr, |
| 370 | - __local_mem__ Tx*& xPtr, | 368 | + __ubuf__ Tx*& xPtr, |
| 371 | - __local_mem__ float*& xMaxPtr, uint32_t a, | 369 | + __ubuf__ float*& xMaxPtr, uint32_t a, |
| 372 | uint32_t ubFactor) | 370 | uint32_t ubFactor) |
| 373 | { | 371 | { |
| 374 | __VEC_SCOPE__ | 372 | __VEC_SCOPE__ |
| @@ -380,31 +378,31 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::MainBlockCastSubExpVF(__ | |||
| 380 | uint32_t width = ubFactor; | 378 | uint32_t width = ubFactor; |
| 381 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); | 379 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); |
| 382 | 380 | ||
| 383 | - AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr); | 381 | + AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr); |
| 384 | for (uint16_t j = 0; j < repeatTimes; j++) { | 382 | for (uint16_t j = 0; j < repeatTimes; j++) { |
| 385 | mask = AscendC::MicroAPI::UpdateMask<float>(width); | 383 | mask = AscendC::MicroAPI::UpdateMask<float>(width); |
| 386 | auto xAddr = xPtr + j * VL_FP32; | 384 | auto xAddr = xPtr + j * VL_FP32; |
| 387 | auto xFp32Addr = xFp32Ptr + j * VL_FP32; | 385 | auto xFp32Addr = xFp32Ptr + j * VL_FP32; |
| 388 | 386 | ||
| 389 | if constexpr (xToFp32_) { | 387 | if constexpr (xToFp32_) { |
| 390 | - AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr); | 388 | + AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr); |
| 391 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask); | 389 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask); |
| 392 | } else { | 390 | } else { |
| 393 | - AscendC::MicroAPI::DataCopy(vreg1, xAddr); | 391 | + AscendC::MicroAPI::LoadAlign(vreg1, xAddr); |
| 394 | } | 392 | } |
| 395 | 393 | ||
| 396 | AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask); | 394 | AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask); |
| 397 | AscendC::MicroAPI::Exp(vreg3, vreg2, mask); | 395 | AscendC::MicroAPI::Exp(vreg3, vreg2, mask); |
| 398 | 396 | ||
| 399 | - AscendC::MicroAPI::DataCopy(xFp32Addr, vreg3, mask); | 397 | + AscendC::MicroAPI::StoreAlign(xFp32Addr, vreg3, mask); |
| 400 | } | 398 | } |
| 401 | } | 399 | } |
| 402 | } | 400 | } |
| 403 | 401 | ||
| 404 | template <typename Tx, typename Ty> | 402 | template <typename Tx, typename Ty> |
| 405 | -__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::FoldBlockCastSubExpVF(__local_mem__ float*& dstPtr, | 403 | +__aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::FoldBlockCastSubExpVF(__ubuf__ float*& dstPtr, |
| 406 | - __local_mem__ Tx*& xPtr, | 404 | + __ubuf__ Tx*& xPtr, |
| 407 | - __local_mem__ float*& xMaxPtr, uint32_t a, | 405 | + __ubuf__ float*& xMaxPtr, uint32_t a, |
| 408 | uint32_t ubFactor) | 406 | uint32_t ubFactor) |
| 409 | { | 407 | { |
| 410 | __VEC_SCOPE__ | 408 | __VEC_SCOPE__ |
| @@ -416,26 +414,26 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::FoldBlockCastSubExpVF(__ | |||
| 416 | uint32_t width = ubFactor; | 414 | uint32_t width = ubFactor; |
| 417 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); | 415 | uint16_t repeatTimes = CeilDivision(ubFactor, VL_FP32); |
| 418 | 416 | ||
| 419 | - AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr); | 417 | + AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(maxReg, xMaxPtr); |
| 420 | for (uint16_t j = 0; j < repeatTimes; j++) { | 418 | for (uint16_t j = 0; j < repeatTimes; j++) { |
| 421 | mask = AscendC::MicroAPI::UpdateMask<float>(width); | 419 | mask = AscendC::MicroAPI::UpdateMask<float>(width); |
| 422 | auto xAddr = xPtr + j * VL_FP32; | 420 | auto xAddr = xPtr + j * VL_FP32; |
| 423 | auto dstAddr = dstPtr + j * VL_FP32; | 421 | auto dstAddr = dstPtr + j * VL_FP32; |
| 424 | 422 | ||
| 425 | if constexpr (xToFp32_) { | 423 | if constexpr (xToFp32_) { |
| 426 | - AscendC::MicroAPI::DataCopy<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr); | 424 | + AscendC::MicroAPI::LoadAlign<Tx, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(vreg0, xAddr); |
| 427 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask); | 425 | AscendC::MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(vreg1, vreg0, mask); |
| 428 | } else { | 426 | } else { |
| 429 | - AscendC::MicroAPI::DataCopy(vreg1, xAddr); | 427 | + AscendC::MicroAPI::LoadAlign(vreg1, xAddr); |
| 430 | } | 428 | } |
| 431 | 429 | ||
| 432 | AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask); | 430 | AscendC::MicroAPI::Sub(vreg2, vreg1, maxReg, mask); |
| 433 | AscendC::MicroAPI::Exp(vreg3, vreg2, mask); | 431 | AscendC::MicroAPI::Exp(vreg3, vreg2, mask); |
| 434 | 432 | ||
| 435 | - AscendC::MicroAPI::DataCopy(dstReg, dstAddr); | 433 | + AscendC::MicroAPI::LoadAlign(dstReg, dstAddr); |
| 436 | AscendC::MicroAPI::Add(dstReg, dstReg, vreg3, mask); | 434 | AscendC::MicroAPI::Add(dstReg, dstReg, vreg3, mask); |
| 437 | 435 | ||
| 438 | - AscendC::MicroAPI::DataCopy(dstAddr, dstReg, mask); | 436 | + AscendC::MicroAPI::StoreAlign(dstAddr, dstReg, mask); |
| 439 | } | 437 | } |
| 440 | } | 438 | } |
| 441 | } | 439 | } |
| @@ -457,9 +455,9 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::UpdateCache(const LocalT | |||
| 457 | uint16_t innerLoopTimes = cacheId; | 455 | uint16_t innerLoopTimes = cacheId; |
| 458 | uint32_t outerLoopStride = VL_FP32; | 456 | uint32_t outerLoopStride = VL_FP32; |
| 459 | uint32_t innerLoopStride = stride; | 457 | uint32_t innerLoopStride = stride; |
| 460 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 458 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 461 | - __local_mem__ float* cache = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheId * stride; | 459 | + __ubuf__ float* cache = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheId * stride; |
| 462 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 460 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 463 | 461 | ||
| 464 | __VEC_SCOPE__ | 462 | __VEC_SCOPE__ |
| 465 | { | 463 | { |
| @@ -468,15 +466,14 @@ __aicore__ inline void LogSoftmaxV2ArRecompute<Tx, Ty>::UpdateCache(const LocalT | |||
| 468 | AscendC::MicroAPI::MaskReg pMask; | 466 | AscendC::MicroAPI::MaskReg pMask; |
| 469 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 467 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 470 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 468 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 471 | - AscendC::MicroAPI::DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 469 | + AscendC::MicroAPI::LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 472 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 470 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 473 | - AscendC::MicroAPI::DataCopy(bReg, | 471 | + AscendC::MicroAPI::LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride); |
| 474 | - (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride); | ||
| 475 | AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 472 | AscendC::MicroAPI::Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 476 | } | 473 | } |
| 477 | - AscendC::MicroAPI::DataCopy((__local_mem__ float*)cache + i * outerLoopStride, aReg, pMask); | 474 | + AscendC::MicroAPI::StoreAlign((__ubuf__ float*)cache + i * outerLoopStride, aReg, pMask); |
| 478 | } | 475 | } |
| 479 | } | 476 | } |
| 480 | } | 477 | } |
| 481 | } // namespace LogSoftmaxV2Ops | 478 | } // namespace LogSoftmaxV2Ops |
| 482 | -#endif // SOFTMAX_V2_AR_RECOMPUTE_H | 479 | +#endif // SOFTMAX_V2_AR_RECOMPUTE_H |
| @@ -110,9 +110,9 @@ private: | |||
| 110 | __aicore__ inline void CalcMaxSubExp(uint32_t curTileA0Len, uint32_t totalRLen) | 110 | __aicore__ inline void CalcMaxSubExp(uint32_t curTileA0Len, uint32_t totalRLen) |
| 111 | { | 111 | { |
| 112 | LocalTensor<Tx> xLocal_ = xQueue_.DeQue<Tx>(); | 112 | LocalTensor<Tx> xLocal_ = xQueue_.DeQue<Tx>(); |
| 113 | - __local_mem__ Tx* xAddr = (__local_mem__ Tx*)xLocal_.GetPhyAddr(); | 113 | + __ubuf__ Tx* xAddr = (__ubuf__ Tx*)xLocal_.GetPhyAddr(); |
| 114 | - __local_mem__ float* tmpAddr = (__local_mem__ float*)tmpLocal_.GetPhyAddr(); | 114 | + __ubuf__ float* tmpAddr = (__ubuf__ float*)tmpLocal_.GetPhyAddr(); |
| 115 | - __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr(); | 115 | + __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr(); |
| 116 | 116 | ||
| 117 | uint16_t aLoopTimes = Ops::Base::CeilDiv(curTileA0Len, VL_FP32); | 117 | uint16_t aLoopTimes = Ops::Base::CeilDiv(curTileA0Len, VL_FP32); |
| 118 | uint16_t rLoopTimes = static_cast<uint16_t>(totalRLen); | 118 | uint16_t rLoopTimes = static_cast<uint16_t>(totalRLen); |
| @@ -139,9 +139,9 @@ private: | |||
| 139 | uint32_t offset = j * VL_FP32 + i * tileA0LenLocal; | 139 | uint32_t offset = j * VL_FP32 + i * tileA0LenLocal; |
| 140 | LoadTensorForDtypeT(xAddr, reg2, mask, offset); | 140 | LoadTensorForDtypeT(xAddr, reg2, mask, offset); |
| 141 | MicroAPI::Sub(reg2, reg2, maxReg, mask); | 141 | MicroAPI::Sub(reg2, reg2, maxReg, mask); |
| 142 | - MicroAPI::DataCopy(tmpAddr2 + offset, reg2, mask); | 142 | + MicroAPI::StoreAlign(tmpAddr2 + offset, reg2, mask); |
| 143 | MicroAPI::Exp(reg2, reg2, mask); | 143 | MicroAPI::Exp(reg2, reg2, mask); |
| 144 | - MicroAPI::DataCopy(tmpAddr + offset, reg2, mask); | 144 | + MicroAPI::StoreAlign(tmpAddr + offset, reg2, mask); |
| 145 | } | 145 | } |
| 146 | } | 146 | } |
| 147 | } | 147 | } |
| @@ -157,10 +157,10 @@ private: | |||
| 157 | 157 | ||
| 158 | __aicore__ inline void CalcOutput(uint32_t curTileA0Len, uint32_t totalRLen) | 158 | __aicore__ inline void CalcOutput(uint32_t curTileA0Len, uint32_t totalRLen) |
| 159 | { | 159 | { |
| 160 | - __local_mem__ float* sumAddr = (__local_mem__ float*)sumLocal_.GetPhyAddr(); | 160 | + __ubuf__ float* sumAddr = (__ubuf__ float*)sumLocal_.GetPhyAddr(); |
| 161 | - __local_mem__ float* tmpAddr2 = (__local_mem__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr(); | 161 | + __ubuf__ float* tmpAddr2 = (__ubuf__ float*)tmpLocal_[tl_->tileA0Len * tl_->rAligned].GetPhyAddr(); |
| 162 | tmpLocalTy_ = tmpLocal_.template ReinterpretCast<Ty>(); | 162 | tmpLocalTy_ = tmpLocal_.template ReinterpretCast<Ty>(); |
| 163 | - __local_mem__ Ty* tmpAddrTy = (__local_mem__ Ty*)tmpLocalTy_.GetPhyAddr(); | 163 | + __ubuf__ Ty* tmpAddrTy = (__ubuf__ Ty*)tmpLocalTy_.GetPhyAddr(); |
| 164 | 164 | ||
| 165 | uint16_t aLoopTimes = static_cast<uint16_t>(Ops::Base::CeilDiv(curTileA0Len, VL_FP32)); | 165 | uint16_t aLoopTimes = static_cast<uint16_t>(Ops::Base::CeilDiv(curTileA0Len, VL_FP32)); |
| 166 | uint16_t rLoopTimes = static_cast<uint16_t>(tl_->totalRLen); | 166 | uint16_t rLoopTimes = static_cast<uint16_t>(tl_->totalRLen); |
| @@ -176,22 +176,22 @@ private: | |||
| 176 | 176 | ||
| 177 | for (uint16_t j = 0; j < aLoopTimes; j++) { // 列 | 177 | for (uint16_t j = 0; j < aLoopTimes; j++) { // 列 |
| 178 | mask = MicroAPI::UpdateMask<float>(sreg); | 178 | mask = MicroAPI::UpdateMask<float>(sreg); |
| 179 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(sumReg, | 179 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(sumReg, |
| 180 | - (__local_mem__ float*)sumAddr + j * VL_FP32); | 180 | + (__ubuf__ float*)sumAddr + j * VL_FP32); |
| 181 | MicroAPI::Ln(lnReg, sumReg, mask); | 181 | MicroAPI::Ln(lnReg, sumReg, mask); |
| 182 | 182 | ||
| 183 | for (uint16_t i = 0; i < rLoopTimes; i++) { // 行 | 183 | for (uint16_t i = 0; i < rLoopTimes; i++) { // 行 |
| 184 | uint32_t offset = j * VL_FP32 + i * tileA0LenLocal; | 184 | uint32_t offset = j * VL_FP32 + i * tileA0LenLocal; |
| 185 | 185 | ||
| 186 | - MicroAPI::DataCopy(reg1, tmpAddr2 + offset); | 186 | + MicroAPI::LoadAlign(reg1, tmpAddr2 + offset); |
| 187 | MicroAPI::Sub(reg1, reg1, lnReg, mask); | 187 | MicroAPI::Sub(reg1, reg1, lnReg, mask); |
| 188 | 188 | ||
| 189 | if constexpr (yToFp32_) { | 189 | if constexpr (yToFp32_) { |
| 190 | - MicroAPI::DataCopy(tmpAddrTy + offset, reg1, mask); | 190 | + MicroAPI::StoreAlign(tmpAddrTy + offset, reg1, mask); |
| 191 | } else { // fp16、bf16 | 191 | } else { // fp16、bf16 |
| 192 | MicroAPI::RegTensor<Ty> xFp16; | 192 | MicroAPI::RegTensor<Ty> xFp16; |
| 193 | MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(xFp16, reg1, mask); | 193 | MicroAPI::Cast<Ty, float, castTraitFp32ToFp16>(xFp16, reg1, mask); |
| 194 | - MicroAPI::DataCopy<Ty, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + offset, xFp16, mask); | 194 | + MicroAPI::StoreAlign<Ty, MicroAPI::StoreDist::DIST_PACK_B32>(tmpAddrTy + offset, xFp16, mask); |
| 195 | } | 195 | } |
| 196 | } | 196 | } |
| 197 | } | 197 | } |
| @@ -263,29 +263,29 @@ private: | |||
| 263 | yQueue_.EnQue(yLocal); | 263 | yQueue_.EnQue(yLocal); |
| 264 | } | 264 | } |
| 265 | 265 | ||
| 266 | - __aicore__ inline void LoadTensorForDtypeT(const __local_mem__ Tx* src, RegTensor<float>& dst, MaskReg& preg, | 266 | + __aicore__ inline void LoadTensorForDtypeT(const __ubuf__ Tx* src, RegTensor<float>& dst, MaskReg& preg, |
| 267 | uint32_t offset) | 267 | uint32_t offset) |
| 268 | { | 268 | { |
| 269 | if constexpr (xToFp32_) { | 269 | if constexpr (xToFp32_) { |
| 270 | MicroAPI::RegTensor<Tx> xFp16; | 270 | MicroAPI::RegTensor<Tx> xFp16; |
| 271 | - MicroAPI::DataCopy<Tx, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ Tx*)src + offset)); | 271 | + MicroAPI::LoadAlign<Tx, MicroAPI::LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ Tx*)src + offset)); |
| 272 | MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(dst, xFp16, preg); | 272 | MicroAPI::Cast<float, Tx, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 273 | } else { | 273 | } else { |
| 274 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset); | 274 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset); |
| 275 | } | 275 | } |
| 276 | } | 276 | } |
| 277 | 277 | ||
| 278 | __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen) | 278 | __aicore__ inline void CopyInAndTransPose(int64_t xGmOffset, uint32_t curTileA0Len, uint32_t totalRLen) |
| 279 | { | 279 | { |
| 280 | - static constexpr MultiCopyConfig config = {false}; | 280 | + static constexpr NdDmaConfig config = {false}; |
| 281 | - MultiCopyLoopInfo<CONST_TWO> copyLoopInfo; | 281 | + NdDmaLoopInfo<CONST_TWO> copyLoopInfo; |
| 282 | copyLoopInfo.loopSrcStride[0] = 1; | 282 | copyLoopInfo.loopSrcStride[0] = 1; |
| 283 | copyLoopInfo.loopSrcStride[1] = totalRLen; | 283 | copyLoopInfo.loopSrcStride[1] = totalRLen; |
| 284 | copyLoopInfo.loopDstStride[0] = tl_->tileA0Len; | 284 | copyLoopInfo.loopDstStride[0] = tl_->tileA0Len; |
| 285 | copyLoopInfo.loopDstStride[1] = 1; | 285 | copyLoopInfo.loopDstStride[1] = 1; |
| 286 | copyLoopInfo.loopSize[0] = totalRLen; | 286 | copyLoopInfo.loopSize[0] = totalRLen; |
| 287 | copyLoopInfo.loopSize[1] = curTileA0Len; | 287 | copyLoopInfo.loopSize[1] = curTileA0Len; |
| 288 | - MultiCopyParams<Tx, CONST_TWO> params = {copyLoopInfo, 0}; | 288 | + NdDmaParams<Tx, CONST_TWO> params = {copyLoopInfo, 0}; |
| 289 | 289 | ||
| 290 | LocalTensor<Tx> xLocal_ = xQueue_.AllocTensor<Tx>(); | 290 | LocalTensor<Tx> xLocal_ = xQueue_.AllocTensor<Tx>(); |
| 291 | DataCopy<Tx, CONST_TWO, config>(xLocal_, xGm_[xGmOffset], params); | 291 | DataCopy<Tx, CONST_TWO, config>(xLocal_, xGm_[xGmOffset], params); |
| @@ -349,4 +349,4 @@ private: | |||
| 349 | 349 | ||
| 350 | } // namespace LogSoftmaxV2Ops | 350 | } // namespace LogSoftmaxV2Ops |
| 351 | 351 | ||
| 352 | -#endif | 352 | +#endif |
| @@ -137,10 +137,10 @@ private: | |||
| 137 | __aicore__ inline void Compute(int64_t curTileRLen, uint32_t curTileA0Len) | 137 | __aicore__ inline void Compute(int64_t curTileRLen, uint32_t curTileA0Len) |
| 138 | { | 138 | { |
| 139 | LocalTensor<T1> x = xQueue_.DeQue<T1>(); | 139 | LocalTensor<T1> x = xQueue_.DeQue<T1>(); |
| 140 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)x.GetPhyAddr(); | 140 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)x.GetPhyAddr(); |
| 141 | 141 | ||
| 142 | LocalTensor<float> xTmpTensor = xTmpBuf_.Get<float>(); | 142 | LocalTensor<float> xTmpTensor = xTmpBuf_.Get<float>(); |
| 143 | - __local_mem__ float* xTmpLocal = (__local_mem__ float*)xTmpTensor.GetPhyAddr(); | 143 | + __ubuf__ float* xTmpLocal = (__ubuf__ float*)xTmpTensor.GetPhyAddr(); |
| 144 | 144 | ||
| 145 | uint16_t loopA0Num = Ops::Base::CeilDiv(curTileA0Len, VL_FP32); | 145 | uint16_t loopA0Num = Ops::Base::CeilDiv(curTileA0Len, VL_FP32); |
| 146 | 146 | ||
| @@ -148,10 +148,10 @@ private: | |||
| 148 | xQueue_.FreeTensor<T1>(x); | 148 | xQueue_.FreeTensor<T1>(x); |
| 149 | 149 | ||
| 150 | LocalTensor<float> y = yQueue_.AllocTensor<float>(); | 150 | LocalTensor<float> y = yQueue_.AllocTensor<float>(); |
| 151 | - __local_mem__ float* yLocal = (__local_mem__ float*)y.GetPhyAddr(); | 151 | + __ubuf__ float* yLocal = (__ubuf__ float*)y.GetPhyAddr(); |
| 152 | 152 | ||
| 153 | LocalTensor<float> xReduceTensor = xReduceBuf_.Get<float>(); | 153 | LocalTensor<float> xReduceTensor = xReduceBuf_.Get<float>(); |
| 154 | - __local_mem__ float* xReduceLocal = (__local_mem__ float*)xReduceTensor.GetPhyAddr(); | 154 | + __ubuf__ float* xReduceLocal = (__ubuf__ float*)xReduceTensor.GetPhyAddr(); |
| 155 | 155 | ||
| 156 | VFReduceSum(xReduceLocal, xTmpLocal, yLocal, curTileRLen, curTileA0Len); | 156 | VFReduceSum(xReduceLocal, xTmpLocal, yLocal, curTileRLen, curTileA0Len); |
| 157 | 157 | ||
| @@ -160,7 +160,7 @@ private: | |||
| 160 | yQueue_.EnQue(y); | 160 | yQueue_.EnQue(y); |
| 161 | } | 161 | } |
| 162 | 162 | ||
| 163 | - __aicore__ inline void VFShiftVector(__local_mem__ float* xTmpLocal, __local_mem__ T1* xLocal, uint16_t curTileRLen, | 163 | + __aicore__ inline void VFShiftVector(__ubuf__ float* xTmpLocal, __ubuf__ T1* xLocal, uint16_t curTileRLen, |
| 164 | uint16_t curTileA0Len, uint16_t loopA0Num) | 164 | uint16_t curTileA0Len, uint16_t loopA0Num) |
| 165 | { | 165 | { |
| 166 | uint32_t tileA0Len = tilingData_->tileA0Len; | 166 | uint32_t tileA0Len = tilingData_->tileA0Len; |
| @@ -187,14 +187,14 @@ private: | |||
| 187 | uint32_t xOffset1 = i * tileA0Len + k * VL_FP32; | 187 | uint32_t xOffset1 = i * tileA0Len + k * VL_FP32; |
| 188 | LoadTensorForDtypeT1(xLocal, x, pregMask, xOffset1); | 188 | LoadTensorForDtypeT1(xLocal, x, pregMask, xOffset1); |
| 189 | Sub(x, x, maxReg, pregMask); | 189 | Sub(x, x, maxReg, pregMask); |
| 190 | - DataCopy(((__local_mem__ float*)xTmpLocal) + xOffset1, x, pregMask); | 190 | + StoreAlign(((__ubuf__ float*)xTmpLocal) + xOffset1, x, pregMask); |
| 191 | } | 191 | } |
| 192 | } | 192 | } |
| 193 | } | 193 | } |
| 194 | } | 194 | } |
| 195 | 195 | ||
| 196 | - __aicore__ inline void VFReduceSum(__local_mem__ float* xReduceLocal, __local_mem__ float* xTmpLocal, | 196 | + __aicore__ inline void VFReduceSum(__ubuf__ float* xReduceLocal, __ubuf__ float* xTmpLocal, __ubuf__ float* yInUb, |
| 197 | - __local_mem__ float* yInUb, uint16_t curTileRLen, uint16_t curTileA0Len) | 197 | + uint16_t curTileRLen, uint16_t curTileA0Len) |
| 198 | { | 198 | { |
| 199 | if (tilingData_->totalRLen <= SCALE_COEF_TWO) { | 199 | if (tilingData_->totalRLen <= SCALE_COEF_TWO) { |
| 200 | SumRLessThan2(xTmpLocal, xReduceLocal, curTileA0Len); | 200 | SumRLessThan2(xTmpLocal, xReduceLocal, curTileA0Len); |
| @@ -207,8 +207,7 @@ private: | |||
| 207 | } | 207 | } |
| 208 | } | 208 | } |
| 209 | 209 | ||
| 210 | - __aicore__ inline void SumRLessThan2(__local_mem__ float* xTmpLocal, __local_mem__ float* xReduceLocal, | 210 | + __aicore__ inline void SumRLessThan2(__ubuf__ float* xTmpLocal, __ubuf__ float* xReduceLocal, uint32_t curTileA0Len) |
| 211 | - uint32_t curTileA0Len) | ||
| 212 | { | 211 | { |
| 213 | uint32_t rStride = tilingData_->tileA0Len; | 212 | uint32_t rStride = tilingData_->tileA0Len; |
| 214 | uint16_t rLoopCount = tilingData_->totalRLen; | 213 | uint16_t rLoopCount = tilingData_->totalRLen; |
| @@ -225,17 +224,16 @@ private: | |||
| 225 | pregLoop = UpdateMask<float>(sreg0); | 224 | pregLoop = UpdateMask<float>(sreg0); |
| 226 | Duplicate(sum, 0.0, pregLoop); | 225 | Duplicate(sum, 0.0, pregLoop); |
| 227 | for (uint16_t i = 0; i < rLoopCount; i++) { | 226 | for (uint16_t i = 0; i < rLoopCount; i++) { |
| 228 | - DataCopy(xld, ((__local_mem__ float*)xTmpLocal + i * rStride + k * VL_FP32)); | 227 | + LoadAlign(xld, ((__ubuf__ float*)xTmpLocal + i * rStride + k * VL_FP32)); |
| 229 | Exp(xld, xld, pregLoop); | 228 | Exp(xld, xld, pregLoop); |
| 230 | Add(sum, sum, xld, pregLoop); | 229 | Add(sum, sum, xld, pregLoop); |
| 231 | } | 230 | } |
| 232 | - DataCopy(((__local_mem__ float*)xReduceLocal + k * VL_FP32), sum, pregLoop); | 231 | + StoreAlign(((__ubuf__ float*)xReduceLocal + k * VL_FP32), sum, pregLoop); |
| 233 | } | 232 | } |
| 234 | } | 233 | } |
| 235 | } | 234 | } |
| 236 | 235 | ||
| 237 | - __aicore__ inline void SumRLessThan4(__local_mem__ float* xTmpLocal, __local_mem__ float* xReduceLocal, | 236 | + __aicore__ inline void SumRLessThan4(__ubuf__ float* xTmpLocal, __ubuf__ float* xReduceLocal, uint32_t curTileA0Len) |
| 238 | - uint32_t curTileA0Len) | ||
| 239 | { | 237 | { |
| 240 | uint32_t remainderOffset = SCALE_COEF_TWO * tilingData_->tileA0Len; | 238 | uint32_t remainderOffset = SCALE_COEF_TWO * tilingData_->tileA0Len; |
| 241 | uint32_t aLength = tilingData_->tileA0Len; | 239 | uint32_t aLength = tilingData_->tileA0Len; |
| @@ -262,50 +260,49 @@ private: | |||
| 262 | for (uint16_t k = 0; k < aLoopCount; k++) { | 260 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 263 | pregLoop = UpdateMask<float>(sreg0); | 261 | pregLoop = UpdateMask<float>(sreg0); |
| 264 | uint32_t aLoopOffset = k * VL_FP32; | 262 | uint32_t aLoopOffset = k * VL_FP32; |
| 265 | - DataCopy(((__local_mem__ float*)xTmpLocal + validNumInXUb + aLoopOffset), nInf, pregLoop); | 263 | + StoreAlign(((__ubuf__ float*)xTmpLocal + validNumInXUb + aLoopOffset), nInf, pregLoop); |
| 266 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 264 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 267 | TwoRowAddExpWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 265 | TwoRowAddExpWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 268 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, | 266 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, |
| 269 | remNextRow); | 267 | remNextRow); |
| 270 | - DataCopy(((__local_mem__ float*)xReduceLocal + aLoopOffset), x1, pregLoop); | 268 | + StoreAlign(((__ubuf__ float*)xReduceLocal + aLoopOffset), x1, pregLoop); |
| 271 | } | 269 | } |
| 272 | } | 270 | } |
| 273 | } | 271 | } |
| 274 | 272 | ||
| 275 | - __aicore__ inline void TwoRowAddExpWithTail(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, | 273 | + __aicore__ inline void TwoRowAddExpWithTail(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 276 | uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, | 274 | uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, |
| 277 | RegTensor<float>& rem, RegTensor<float>& nextRow, | 275 | RegTensor<float>& rem, RegTensor<float>& nextRow, |
| 278 | RegTensor<float>& remNextRow) | 276 | RegTensor<float>& remNextRow) |
| 279 | { | 277 | { |
| 280 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 278 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 281 | - DataCopy(rem, ((__local_mem__ float*)(input) + (offset2))); | 279 | + LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2))); |
| 282 | Exp(dst, dst, preg); | 280 | Exp(dst, dst, preg); |
| 283 | Exp(rem, rem, preg); | 281 | Exp(rem, rem, preg); |
| 284 | Add(dst, dst, rem, preg); | 282 | Add(dst, dst, rem, preg); |
| 285 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset3))); | 283 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset3))); |
| 286 | - DataCopy(remNextRow, ((__local_mem__ float*)(input) + (offset4))); | 284 | + LoadAlign(remNextRow, ((__ubuf__ float*)(input) + (offset4))); |
| 287 | Exp(nextRow, nextRow, preg); | 285 | Exp(nextRow, nextRow, preg); |
| 288 | Exp(remNextRow, remNextRow, preg); | 286 | Exp(remNextRow, remNextRow, preg); |
| 289 | Add(nextRow, nextRow, remNextRow, preg); | 287 | Add(nextRow, nextRow, remNextRow, preg); |
| 290 | Add(dst, dst, nextRow, preg); | 288 | Add(dst, dst, nextRow, preg); |
| 291 | } | 289 | } |
| 292 | 290 | ||
| 293 | - __aicore__ inline void TwoRowAddWithTail(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, | 291 | + __aicore__ inline void TwoRowAddWithTail(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 294 | uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, | 292 | uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, |
| 295 | RegTensor<float>& rem, RegTensor<float>& nextRow, | 293 | RegTensor<float>& rem, RegTensor<float>& nextRow, |
| 296 | RegTensor<float>& remNextRow) | 294 | RegTensor<float>& remNextRow) |
| 297 | { | 295 | { |
| 298 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 296 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 299 | - DataCopy(rem, ((__local_mem__ float*)(input) + (offset2))); | 297 | + LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2))); |
| 300 | Add(dst, dst, rem, preg); | 298 | Add(dst, dst, rem, preg); |
| 301 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset3))); | 299 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset3))); |
| 302 | - DataCopy(remNextRow, ((__local_mem__ float*)(input) + (offset4))); | 300 | + LoadAlign(remNextRow, ((__ubuf__ float*)(input) + (offset4))); |
| 303 | Add(nextRow, nextRow, remNextRow, preg); | 301 | Add(nextRow, nextRow, remNextRow, preg); |
| 304 | Add(dst, dst, nextRow, preg); | 302 | Add(dst, dst, nextRow, preg); |
| 305 | } | 303 | } |
| 306 | 304 | ||
| 307 | - __aicore__ inline void SumRLessThan8(__local_mem__ float* xTmpLocal, __local_mem__ float* xReduceLocal, | 305 | + __aicore__ inline void SumRLessThan8(__ubuf__ float* xTmpLocal, __ubuf__ float* xReduceLocal, uint32_t curTileA0Len) |
| 308 | - uint32_t curTileA0Len) | ||
| 309 | { | 306 | { |
| 310 | uint32_t remainderOffset = SCALE_COEF_FOUR * tilingData_->tileA0Len; | 307 | uint32_t remainderOffset = SCALE_COEF_FOUR * tilingData_->tileA0Len; |
| 311 | uint32_t aLength = tilingData_->tileA0Len; | 308 | uint32_t aLength = tilingData_->tileA0Len; |
| @@ -337,7 +334,7 @@ private: | |||
| 337 | for (uint16_t k = 0; k < aLoopCount; k++) { | 334 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 338 | pregLoop = UpdateMask<float>(sreg0); | 335 | pregLoop = UpdateMask<float>(sreg0); |
| 339 | uint32_t aLoopOffset = k * VL_FP32; | 336 | uint32_t aLoopOffset = k * VL_FP32; |
| 340 | - DataCopy(((__local_mem__ float*)xTmpLocal + validNumInXUb + aLoopOffset), nInf, pregLoop); | 337 | + StoreAlign(((__ubuf__ float*)xTmpLocal + validNumInXUb + aLoopOffset), nInf, pregLoop); |
| 341 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 338 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 342 | TwoRowAddExpWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 339 | TwoRowAddExpWithTail(x1, xTmpLocal, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 343 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, | 340 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, |
| @@ -346,13 +343,13 @@ private: | |||
| 346 | remainderTailOffset2 + aLoopOffset, ROW_THREE_OFFSET * aLength + aLoopOffset, | 343 | remainderTailOffset2 + aLoopOffset, ROW_THREE_OFFSET * aLength + aLoopOffset, |
| 347 | remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow); | 344 | remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow); |
| 348 | Add(x1, x1, x2, pregLoop); | 345 | Add(x1, x1, x2, pregLoop); |
| 349 | - DataCopy(((__local_mem__ float*)xReduceLocal + aLoopOffset), x1, pregLoop); | 346 | + StoreAlign(((__ubuf__ float*)xReduceLocal + aLoopOffset), x1, pregLoop); |
| 350 | } | 347 | } |
| 351 | } | 348 | } |
| 352 | } | 349 | } |
| 353 | 350 | ||
| 354 | - __aicore__ inline void SumRMoreThan8(__local_mem__ float* xInUb, __local_mem__ float* yInUb, | 351 | + __aicore__ inline void SumRMoreThan8(__ubuf__ float* xInUb, __ubuf__ float* yInUb, __ubuf__ float* xReduceLocal, |
| 355 | - __local_mem__ float* xReduceLocal, uint32_t curTileA0Len) | 352 | + uint32_t curTileA0Len) |
| 356 | { | 353 | { |
| 357 | uint16_t remainderLoopCount = tilingData_->remainderLoopCount; | 354 | uint16_t remainderLoopCount = tilingData_->remainderLoopCount; |
| 358 | uint16_t remainderLoopCountTmp = remainderLoopCount - 1; | 355 | uint16_t remainderLoopCountTmp = remainderLoopCount - 1; |
| @@ -401,7 +398,7 @@ private: | |||
| 401 | for (uint16_t k = 0; k < aLoopCount; k++) { | 398 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 402 | pregLoop = UpdateMask<float>(sreg0); | 399 | pregLoop = UpdateMask<float>(sreg0); |
| 403 | uint32_t aLoopOffset = k * VL_FP32; | 400 | uint32_t aLoopOffset = k * VL_FP32; |
| 404 | - DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), nInf, pregLoop); | 401 | + StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), nInf, pregLoop); |
| 405 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 402 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 406 | // 前半部分与后半部分中，都为8行的部分 | 403 | // 前半部分与后半部分中，都为8行的部分 |
| 407 | for (uint16_t i = 0; i < remainderLoopCountTmp; i++) { | 404 | for (uint16_t i = 0; i < remainderLoopCountTmp; i++) { |
| @@ -421,7 +418,7 @@ private: | |||
| 421 | remOffset + ROW_SEVEN_OFFSET * aLength, rem, nextRow, remNextRow); | 418 | remOffset + ROW_SEVEN_OFFSET * aLength, rem, nextRow, remNextRow); |
| 422 | Add(x3, x3, x4, pregLoop); | 419 | Add(x3, x3, x4, pregLoop); |
| 423 | Add(x1, x1, x3, pregLoop); | 420 | Add(x1, x1, x3, pregLoop); |
| 424 | - DataCopy(((__local_mem__ float*)yInUb + i * aLength + aLoopOffset), x1, pregLoop); | 421 | + StoreAlign(((__ubuf__ float*)yInUb + i * aLength + aLoopOffset), x1, pregLoop); |
| 425 | } | 422 | } |
| 426 | // 前半部分为8行，后半部分可能不足8行 | 423 | // 前半部分为8行，后半部分可能不足8行 |
| 427 | { | 424 | { |
| @@ -446,8 +443,8 @@ private: | |||
| 446 | remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow); | 443 | remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow); |
| 447 | Add(x3, x3, x4, pregLoop); | 444 | Add(x3, x3, x4, pregLoop); |
| 448 | Add(x1, x1, x3, pregLoop); | 445 | Add(x1, x1, x3, pregLoop); |
| 449 | - DataCopy(((__local_mem__ float*)yInUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, | 446 | + StoreAlign(((__ubuf__ float*)yInUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, |
| 450 | - pregLoop); | 447 | + pregLoop); |
| 451 | } | 448 | } |
| 452 | // 剩余的前半部分，一次for循环，处理8行 | 449 | // 剩余的前半部分，一次for循环，处理8行 |
| 453 | for (uint16_t i = 0; i < quotientLoopCount; i++) { | 450 | for (uint16_t i = 0; i < quotientLoopCount; i++) { |
| @@ -462,29 +459,29 @@ private: | |||
| 462 | baseOffset + ROW_SEVEN_OFFSET * aLength, nextRow); | 459 | baseOffset + ROW_SEVEN_OFFSET * aLength, nextRow); |
| 463 | Add(x3, x3, x4, pregLoop); | 460 | Add(x3, x3, x4, pregLoop); |
| 464 | Add(x1, x1, x3, pregLoop); | 461 | Add(x1, x1, x3, pregLoop); |
| 465 | - DataCopy(((__local_mem__ float*)yInUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1, | 462 | + StoreAlign(((__ubuf__ float*)yInUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1, |
| 466 | - pregLoop); | 463 | + pregLoop); |
| 467 | } | 464 | } |
| 468 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 465 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 469 | - BinaryAddVF((__local_mem__ float*)yInUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop, | 466 | + BinaryAddVF((__ubuf__ float*)yInUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop, |
| 470 | binaryAddLastLoop, pregLoop, x1, x2, x3, x4); | 467 | binaryAddLastLoop, pregLoop, x1, x2, x3, x4); |
| 471 | - DataCopy(x1, ((__local_mem__ float*)yInUb + aLoopOffset)); | 468 | + LoadAlign(x1, ((__ubuf__ float*)yInUb + aLoopOffset)); |
| 472 | - DataCopy(((__local_mem__ float*)xReduceLocal + aLoopOffset), x1, pregLoop); | 469 | + StoreAlign(((__ubuf__ float*)xReduceLocal + aLoopOffset), x1, pregLoop); |
| 473 | } | 470 | } |
| 474 | } | 471 | } |
| 475 | } | 472 | } |
| 476 | 473 | ||
| 477 | - __aicore__ inline void TwoRowAdd(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, uint32_t offset1, | 474 | + __aicore__ inline void TwoRowAdd(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, uint32_t offset1, |
| 478 | uint32_t offset2, RegTensor<float>& nextRow) | 475 | uint32_t offset2, RegTensor<float>& nextRow) |
| 479 | { | 476 | { |
| 480 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 477 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 481 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset2))); | 478 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset2))); |
| 482 | Exp(dst, dst, preg); | 479 | Exp(dst, dst, preg); |
| 483 | Exp(nextRow, nextRow, preg); | 480 | Exp(nextRow, nextRow, preg); |
| 484 | Add(dst, dst, nextRow, preg); | 481 | Add(dst, dst, nextRow, preg); |
| 485 | } | 482 | } |
| 486 | 483 | ||
| 487 | - __aicore__ inline void BinaryAddVF(__local_mem__ float* binaryAddTmpAddr, uint32_t rLoopStride, uint32_t offset, | 484 | + __aicore__ inline void BinaryAddVF(__ubuf__ float* binaryAddTmpAddr, uint32_t rLoopStride, uint32_t offset, |
| 488 | uint16_t binaryAddKLoop, uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop, | 485 | uint16_t binaryAddKLoop, uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop, |
| 489 | MaskReg& pregLoop, RegTensor<float>& x1, RegTensor<float>& x2, | 486 | MaskReg& pregLoop, RegTensor<float>& x1, RegTensor<float>& x2, |
| 490 | RegTensor<float>& x3, RegTensor<float>& x4) | 487 | RegTensor<float>& x3, RegTensor<float>& x4) |
| @@ -493,32 +490,31 @@ private: | |||
| 493 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { | 490 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { |
| 494 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET; | 491 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET; |
| 495 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { | 492 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { |
| 496 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset)); | 493 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset)); |
| 497 | - DataCopy(x2, | 494 | + LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset)); |
| 498 | - ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset)); | ||
| 499 | Add(x1, x1, x2, pregLoop); | 495 | Add(x1, x1, x2, pregLoop); |
| 500 | - DataCopy(x3, ((__local_mem__ float*)binaryAddTmpAddr + | 496 | + LoadAlign(x3, ((__ubuf__ float*)binaryAddTmpAddr + |
| 501 | - (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset)); | 497 | + (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset)); |
| 502 | - DataCopy(x4, ((__local_mem__ float*)binaryAddTmpAddr + | 498 | + LoadAlign(x4, ((__ubuf__ float*)binaryAddTmpAddr + |
| 503 | - (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset)); | 499 | + (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset)); |
| 504 | Add(x3, x3, x4, pregLoop); | 500 | Add(x3, x3, x4, pregLoop); |
| 505 | Add(x1, x1, x3, pregLoop); | 501 | Add(x1, x1, x3, pregLoop); |
| 506 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop); | 502 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop); |
| 507 | } | 503 | } |
| 508 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 504 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 509 | } | 505 | } |
| 510 | for (uint16_t i = 0; i < binaryAddLastLoop; i++) { | 506 | for (uint16_t i = 0; i < binaryAddLastLoop; i++) { |
| 511 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + offset)); | 507 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + offset)); |
| 512 | - DataCopy(x2, ((__local_mem__ float*)binaryAddTmpAddr + rLoopStride + offset)); | 508 | + LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + rLoopStride + offset)); |
| 513 | Add(x1, x1, x2, pregLoop); | 509 | Add(x1, x1, x2, pregLoop); |
| 514 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + offset), x1, pregLoop); | 510 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + offset), x1, pregLoop); |
| 515 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 511 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 516 | } | 512 | } |
| 517 | } | 513 | } |
| 518 | 514 | ||
| 519 | - __aicore__ inline void VFCalculateOutput(__local_mem__ float* yLocal, __local_mem__ float* xTmpLocal, | 515 | + __aicore__ inline void VFCalculateOutput(__ubuf__ float* yLocal, __ubuf__ float* xTmpLocal, |
| 520 | - __local_mem__ float* xReduceLocal, uint16_t curTileRLen, | 516 | + __ubuf__ float* xReduceLocal, uint16_t curTileRLen, uint16_t curTileA0Len, |
| 521 | - uint16_t curTileA0Len, uint16_t loopA0Num) | 517 | + uint16_t loopA0Num) |
| 522 | { | 518 | { |
| 523 | uint32_t tileA0Len = tilingData_->tileA0Len; | 519 | uint32_t tileA0Len = tilingData_->tileA0Len; |
| 524 | __VEC_SCOPE__ | 520 | __VEC_SCOPE__ |
| @@ -532,35 +528,34 @@ private: | |||
| 532 | 528 | ||
| 533 | for (uint16_t k = 0; k < loopA0Num; k++) { | 529 | for (uint16_t k = 0; k < loopA0Num; k++) { |
| 534 | pregMask = UpdateMask<float>(sreg); | 530 | pregMask = UpdateMask<float>(sreg); |
| 535 | - DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xReduceLocal + k * VL_FP32); | 531 | + LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xReduceLocal + k * VL_FP32); |
| 536 | Log(sumReg, sumReg, pregMask); | 532 | Log(sumReg, sumReg, pregMask); |
| 537 | for (uint16_t i = 0; i < curTileRLen; i++) { | 533 | for (uint16_t i = 0; i < curTileRLen; i++) { |
| 538 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; | 534 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; |
| 539 | 535 | ||
| 540 | - DataCopy<float, LoadDist::DIST_NORM>(xReg, (__local_mem__ float*)xTmpLocal + xOffset); | 536 | + LoadAlign<float, LoadDist::DIST_NORM>(xReg, (__ubuf__ float*)xTmpLocal + xOffset); |
| 541 | Sub(yReg, xReg, sumReg, pregMask); | 537 | Sub(yReg, xReg, sumReg, pregMask); |
| 542 | 538 | ||
| 543 | // copy out | 539 | // copy out |
| 544 | if constexpr (IsSameType<T2, float>::value) { | 540 | if constexpr (IsSameType<T2, float>::value) { |
| 545 | - DataCopy(((__local_mem__ float*)yLocal) + xOffset, yReg, pregMask); | 541 | + StoreAlign(((__ubuf__ float*)yLocal) + xOffset, yReg, pregMask); |
| 546 | } else { // fp16、bf16 | 542 | } else { // fp16、bf16 |
| 547 | RegTensor<T2> xFp16; | 543 | RegTensor<T2> xFp16; |
| 548 | Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask); | 544 | Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask); |
| 549 | - DataCopy<T2, StoreDist::DIST_PACK_B32>(((__local_mem__ T2*)yLocal) + xOffset, xFp16, pregMask); | 545 | + StoreAlign<T2, StoreDist::DIST_PACK_B32>(((__ubuf__ T2*)yLocal) + xOffset, xFp16, pregMask); |
| 550 | } | 546 | } |
| 551 | } | 547 | } |
| 552 | } | 548 | } |
| 553 | } | 549 | } |
| 554 | } | 550 | } |
| 555 | 551 | ||
| 556 | - __aicore__ inline void LoadTensorForDtypeT1(__local_mem__ T1* src, RegTensor<float>& dst, MaskReg& preg, | 552 | + __aicore__ inline void LoadTensorForDtypeT1(__ubuf__ T1* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset) |
| 557 | - uint32_t offset) | ||
| 558 | { | 553 | { |
| 559 | if constexpr (IsSameType<T1, float>::value) { | 554 | if constexpr (IsSameType<T1, float>::value) { |
| 560 | - DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset); | 555 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset); |
| 561 | } else { // fp16、bf16 | 556 | } else { // fp16、bf16 |
| 562 | RegTensor<T1> xFp16; | 557 | RegTensor<T1> xFp16; |
| 563 | - DataCopy<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T1*)src + offset)); | 558 | + LoadAlign<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T1*)src + offset)); |
| 564 | Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg); | 559 | Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 565 | } | 560 | } |
| 566 | } | 561 | } |
| @@ -112,7 +112,7 @@ private: | |||
| 112 | { | 112 | { |
| 113 | // max 初始化 | 113 | // max 初始化 |
| 114 | LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>(); | 114 | LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>(); |
| 115 | - __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr(); | 115 | + __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr(); |
| 116 | 116 | ||
| 117 | __VEC_SCOPE__ | 117 | __VEC_SCOPE__ |
| 118 | { | 118 | { |
| @@ -122,7 +122,7 @@ private: | |||
| 122 | for (uint16_t k = 0; k < loopA0Num; k++) { | 122 | for (uint16_t k = 0; k < loopA0Num; k++) { |
| 123 | pregMask = UpdateMask<float>(sreg); | 123 | pregMask = UpdateMask<float>(sreg); |
| 124 | Duplicate(maxReg, static_cast<float>(-INFINITY), pregMask); | 124 | Duplicate(maxReg, static_cast<float>(-INFINITY), pregMask); |
| 125 | - DataCopy(((__local_mem__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask); | 125 | + StoreAlign(((__ubuf__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask); |
| 126 | } | 126 | } |
| 127 | } | 127 | } |
| 128 | 128 | ||
| @@ -136,7 +136,7 @@ private: | |||
| 136 | 136 | ||
| 137 | CopyInX(xOffset, curTileRLen, curTileA0Len); | 137 | CopyInX(xOffset, curTileRLen, curTileA0Len); |
| 138 | LocalTensor<T1> x = xQueue_.DeQue<T1>(); | 138 | LocalTensor<T1> x = xQueue_.DeQue<T1>(); |
| 139 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)x.GetPhyAddr(); | 139 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)x.GetPhyAddr(); |
| 140 | 140 | ||
| 141 | __VEC_SCOPE__ | 141 | __VEC_SCOPE__ |
| 142 | { | 142 | { |
| @@ -149,7 +149,7 @@ private: | |||
| 149 | for (uint16_t k = 0; k < loopA0Num; k++) { | 149 | for (uint16_t k = 0; k < loopA0Num; k++) { |
| 150 | pregMask = UpdateMask<float>(sreg); | 150 | pregMask = UpdateMask<float>(sreg); |
| 151 | // load max | 151 | // load max |
| 152 | - DataCopy<float, LoadDist::DIST_NORM>(maxReg, (__local_mem__ float*)xMaxLocal + k * VL_FP32); | 152 | + LoadAlign<float, LoadDist::DIST_NORM>(maxReg, (__ubuf__ float*)xMaxLocal + k * VL_FP32); |
| 153 | 153 | ||
| 154 | for (uint16_t i = 0; i < curTileRLenVl; i++) { | 154 | for (uint16_t i = 0; i < curTileRLenVl; i++) { |
| 155 | uint32_t offset = i * tileA0Len + k * VL_FP32; | 155 | uint32_t offset = i * tileA0Len + k * VL_FP32; |
| @@ -157,7 +157,7 @@ private: | |||
| 157 | LoadTensorForDtypeT1(xLocal, x, pregMask, offset); | 157 | LoadTensorForDtypeT1(xLocal, x, pregMask, offset); |
| 158 | Max(maxReg, maxReg, x, pregMask); | 158 | Max(maxReg, maxReg, x, pregMask); |
| 159 | } | 159 | } |
| 160 | - DataCopy(((__local_mem__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask); | 160 | + StoreAlign(((__ubuf__ float*)xMaxLocal) + k * VL_FP32, maxReg, pregMask); |
| 161 | } | 161 | } |
| 162 | } | 162 | } |
| 163 | 163 | ||
| @@ -214,9 +214,9 @@ private: | |||
| 214 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); | 214 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); |
| 215 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; | 215 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; |
| 216 | 216 | ||
| 217 | - __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr(); | 217 | + __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr(); |
| 218 | - __local_mem__ T1* src = (__local_mem__ T1*)xMain_.GetPhyAddr(); | 218 | + __ubuf__ T1* src = (__ubuf__ T1*)xMain_.GetPhyAddr(); |
| 219 | - __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr(); | 219 | + __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr(); |
| 220 | 220 | ||
| 221 | __VEC_SCOPE__ | 221 | __VEC_SCOPE__ |
| 222 | { | 222 | { |
| @@ -229,13 +229,13 @@ private: | |||
| 229 | 229 | ||
| 230 | for (uint16_t j = 0; j < loopA0Num; ++j) { | 230 | for (uint16_t j = 0; j < loopA0Num; ++j) { |
| 231 | pregMask = UpdateMask<float>(sreg); | 231 | pregMask = UpdateMask<float>(sreg); |
| 232 | - DataCopy<float, LoadDist::DIST_NORM>(maxReg, (__local_mem__ float*)xMaxLocal + j * VL_FP32); | 232 | + LoadAlign<float, LoadDist::DIST_NORM>(maxReg, (__ubuf__ float*)xMaxLocal + j * VL_FP32); |
| 233 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 233 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 234 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; | 234 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; |
| 235 | LoadTensorForDtypeT1(src, srcReg, pregMask, xOffset); | 235 | LoadTensorForDtypeT1(src, srcReg, pregMask, xOffset); |
| 236 | Sub(dstReg, srcReg, maxReg, pregMask); | 236 | Sub(dstReg, srcReg, maxReg, pregMask); |
| 237 | Exp(dstReg, dstReg, pregMask); | 237 | Exp(dstReg, dstReg, pregMask); |
| 238 | - DataCopy((__local_mem__ float*)dst + xOffset, dstReg, pregMask); | 238 | + StoreAlign((__ubuf__ float*)dst + xOffset, dstReg, pregMask); |
| 239 | } | 239 | } |
| 240 | } | 240 | } |
| 241 | } | 241 | } |
| @@ -256,9 +256,9 @@ private: | |||
| 256 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); | 256 | uint16_t outerLoopTimes = static_cast<uint16_t>(curTileRLen); |
| 257 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; | 257 | uint32_t outerLoopSrcStride = tilingData_->tileA0Len; |
| 258 | 258 | ||
| 259 | - __local_mem__ float* dst = (__local_mem__ float*)yMain_.GetPhyAddr(); | 259 | + __ubuf__ float* dst = (__ubuf__ float*)yMain_.GetPhyAddr(); |
| 260 | - __local_mem__ T1* src = (__local_mem__ T1*)xFold.GetPhyAddr(); | 260 | + __ubuf__ T1* src = (__ubuf__ T1*)xFold.GetPhyAddr(); |
| 261 | - __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr(); | 261 | + __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr(); |
| 262 | 262 | ||
| 263 | __VEC_SCOPE__ | 263 | __VEC_SCOPE__ |
| 264 | { | 264 | { |
| @@ -271,15 +271,15 @@ private: | |||
| 271 | 271 | ||
| 272 | for (uint16_t j = 0; j < loopA0Num; ++j) { | 272 | for (uint16_t j = 0; j < loopA0Num; ++j) { |
| 273 | pregMask = UpdateMask<float>(sreg); | 273 | pregMask = UpdateMask<float>(sreg); |
| 274 | - DataCopy<float, LoadDist::DIST_NORM>(maxReg, (__local_mem__ float*)xMaxLocal + j * VL_FP32); | 274 | + LoadAlign<float, LoadDist::DIST_NORM>(maxReg, (__ubuf__ float*)xMaxLocal + j * VL_FP32); |
| 275 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 275 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 276 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; | 276 | uint32_t xOffset = i * outerLoopSrcStride + j * VL_FP32; |
| 277 | LoadTensorForDtypeT1(src, srcReg, pregMask, i * outerLoopSrcStride + j * VL_FP32); | 277 | LoadTensorForDtypeT1(src, srcReg, pregMask, i * outerLoopSrcStride + j * VL_FP32); |
| 278 | Sub(dstReg, srcReg, maxReg, pregMask); | 278 | Sub(dstReg, srcReg, maxReg, pregMask); |
| 279 | Exp(dstReg, dstReg, pregMask); | 279 | Exp(dstReg, dstReg, pregMask); |
| 280 | - DataCopy(srcReg, (__local_mem__ float*)dst + xOffset); | 280 | + LoadAlign(srcReg, (__ubuf__ float*)dst + xOffset); |
| 281 | Add(dstReg, dstReg, srcReg, pregMask); | 281 | Add(dstReg, dstReg, srcReg, pregMask); |
| 282 | - DataCopy((__local_mem__ float*)dst + xOffset, dstReg, pregMask); | 282 | + StoreAlign((__ubuf__ float*)dst + xOffset, dstReg, pregMask); |
| 283 | } | 283 | } |
| 284 | } | 284 | } |
| 285 | } | 285 | } |
| @@ -330,16 +330,16 @@ private: | |||
| 330 | __aicore__ inline void CalcOutput(int64_t curTileRLen, uint32_t curTileA0Len, uint16_t loopA0Num) | 330 | __aicore__ inline void CalcOutput(int64_t curTileRLen, uint32_t curTileA0Len, uint16_t loopA0Num) |
| 331 | { | 331 | { |
| 332 | LocalTensor<T1> x = xQueue_.DeQue<T1>(); | 332 | LocalTensor<T1> x = xQueue_.DeQue<T1>(); |
| 333 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)x.GetPhyAddr(); | 333 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)x.GetPhyAddr(); |
| 334 | 334 | ||
| 335 | LocalTensor<T2> y = yQueue_.template AllocTensor<T2>(); | 335 | LocalTensor<T2> y = yQueue_.template AllocTensor<T2>(); |
| 336 | - __local_mem__ T2* yLocal = (__local_mem__ T2*)y.GetPhyAddr(); | 336 | + __ubuf__ T2* yLocal = (__ubuf__ T2*)y.GetPhyAddr(); |
| 337 | 337 | ||
| 338 | LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>(); | 338 | LocalTensor<float> xMaxTensor = xMaxBuf_.Get<float>(); |
| 339 | - __local_mem__ float* xMaxLocal = (__local_mem__ float*)xMaxTensor.GetPhyAddr(); | 339 | + __ubuf__ float* xMaxLocal = (__ubuf__ float*)xMaxTensor.GetPhyAddr(); |
| 340 | 340 | ||
| 341 | LocalTensor<float> xSumTensor = xSumBuf_.Get<float>(); | 341 | LocalTensor<float> xSumTensor = xSumBuf_.Get<float>(); |
| 342 | - __local_mem__ float* xSumLocal = (__local_mem__ float*)xSumTensor.GetPhyAddr(); | 342 | + __ubuf__ float* xSumLocal = (__ubuf__ float*)xSumTensor.GetPhyAddr(); |
| 343 | 343 | ||
| 344 | uint32_t tileA0Len = tilingData_->tileA0Len; | 344 | uint32_t tileA0Len = tilingData_->tileA0Len; |
| 345 | uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen); | 345 | uint16_t curTileRLenVl = static_cast<uint16_t>(curTileRLen); |
| @@ -355,8 +355,8 @@ private: | |||
| 355 | 355 | ||
| 356 | for (uint16_t k = 0; k < loopA0Num; k++) { | 356 | for (uint16_t k = 0; k < loopA0Num; k++) { |
| 357 | pregMask = UpdateMask<float>(sreg); | 357 | pregMask = UpdateMask<float>(sreg); |
| 358 | - DataCopy<float, LoadDist::DIST_NORM>(sumReg, (__local_mem__ float*)xSumLocal + k * VL_FP32); | 358 | + LoadAlign<float, LoadDist::DIST_NORM>(sumReg, (__ubuf__ float*)xSumLocal + k * VL_FP32); |
| 359 | - DataCopy<float, LoadDist::DIST_NORM>(maxReg, (__local_mem__ float*)xMaxLocal + k * VL_FP32); | 359 | + LoadAlign<float, LoadDist::DIST_NORM>(maxReg, (__ubuf__ float*)xMaxLocal + k * VL_FP32); |
| 360 | Log(sumReg, sumReg, pregMask); | 360 | Log(sumReg, sumReg, pregMask); |
| 361 | for (uint16_t i = 0; i < curTileRLenVl; i++) { | 361 | for (uint16_t i = 0; i < curTileRLenVl; i++) { |
| 362 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; | 362 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; |
| @@ -367,11 +367,11 @@ private: | |||
| 367 | 367 | ||
| 368 | // copy out | 368 | // copy out |
| 369 | if constexpr (IsSameType<T2, float>::value) { | 369 | if constexpr (IsSameType<T2, float>::value) { |
| 370 | - DataCopy(((__local_mem__ float*)yLocal) + xOffset, yReg, pregMask); | 370 | + StoreAlign(((__ubuf__ float*)yLocal) + xOffset, yReg, pregMask); |
| 371 | } else { // fp16、bf16 | 371 | } else { // fp16、bf16 |
| 372 | RegTensor<T2> xFp16; | 372 | RegTensor<T2> xFp16; |
| 373 | Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask); | 373 | Cast<T2, float, castTraitFp32ToFp16>(xFp16, yReg, pregMask); |
| 374 | - DataCopy<T2, StoreDist::DIST_PACK_B32>(((__local_mem__ T2*)yLocal) + xOffset, xFp16, pregMask); | 374 | + StoreAlign<T2, StoreDist::DIST_PACK_B32>(((__ubuf__ T2*)yLocal) + xOffset, xFp16, pregMask); |
| 375 | } | 375 | } |
| 376 | } | 376 | } |
| 377 | } | 377 | } |
| @@ -382,14 +382,13 @@ private: | |||
| 382 | xQueue_.FreeTensor<T1>(x); | 382 | xQueue_.FreeTensor<T1>(x); |
| 383 | } | 383 | } |
| 384 | 384 | ||
| 385 | - __aicore__ inline void LoadTensorForDtypeT1(__local_mem__ T1* src, RegTensor<float>& dst, MaskReg& preg, | 385 | + __aicore__ inline void LoadTensorForDtypeT1(__ubuf__ T1* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset) |
| 386 | - uint32_t offset) | ||
| 387 | { | 386 | { |
| 388 | if constexpr (IsSameType<T1, float>::value) { | 387 | if constexpr (IsSameType<T1, float>::value) { |
| 389 | - DataCopy<float, LoadDist::DIST_NORM>(dst, (__local_mem__ float*)src + offset); | 388 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, (__ubuf__ float*)src + offset); |
| 390 | } else { // fp16、bf16 | 389 | } else { // fp16、bf16 |
| 391 | RegTensor<T1> xFp16; | 390 | RegTensor<T1> xFp16; |
| 392 | - DataCopy<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ T1*)src + offset)); | 391 | + LoadAlign<T1, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ T1*)src + offset)); |
| 393 | Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg); | 392 | Cast<float, T1, castTraitFp16ToFp32>(dst, xFp16, preg); |
| 394 | } | 393 | } |
| 395 | } | 394 | } |
| @@ -55,9 +55,9 @@ struct CalcInput : public ElemwiseBinaryOP<U, U, T> { | |||
| 55 | // gt = maximizeFactor * gt | 55 | // gt = maximizeFactor * gt |
| 56 | for (uint16_t loop = 0; loop < (uint16_t)repeatTimes; loop++) { | 56 | for (uint16_t loop = 0; loop < (uint16_t)repeatTimes; loop++) { |
| 57 | pregUp = MicroAPI::UpdateMask<U>(totalLen); | 57 | pregUp = MicroAPI::UpdateMask<U>(totalLen); |
| 58 | - MicroAPI::DataCopy<U, MicroAPI::PostLiteral::POST_MODE_UPDATE>(regGrad, gradAddr, (int32_t)oneRepeat); | 58 | + MicroAPI::LoadAlign<U, MicroAPI::PostLiteral::POST_MODE_UPDATE>(regGrad, gradAddr, (int32_t)oneRepeat); |
| 59 | - MicroAPI::DataCopy<U, MicroAPI::PostLiteral::POST_MODE_UPDATE>(gradOutAddr, regGrad, (int32_t)oneRepeat, | 59 | + MicroAPI::StoreAlign<U, MicroAPI::PostLiteral::POST_MODE_UPDATE>(gradOutAddr, regGrad, |
| 60 | - pregUp); | 60 | + (int32_t)oneRepeat, pregUp); |
| 61 | } | 61 | } |
| 62 | } | 62 | } |
| 63 | 63 | ||
| @@ -124,4 +124,4 @@ struct KDLGLogTargetFalse { | |||
| 124 | }; | 124 | }; |
| 125 | 125 | ||
| 126 | } // namespace KlDivLossGrad | 126 | } // namespace KlDivLossGrad |
| 127 | -#endif // ASCENDC_KL_DIV_LOSS_GRAD_DAG_H_ | 127 | +#endif // ASCENDC_KL_DIV_LOSS_GRAD_DAG_H_ |
| @@ -49,42 +49,40 @@ protected: | |||
| 49 | __aicore__ inline void UpdateCache(const int64_t basicBlockIdx, const LocalTensor<float>& cacheDbeta, | 49 | __aicore__ inline void UpdateCache(const int64_t basicBlockIdx, const LocalTensor<float>& cacheDbeta, |
| 50 | const LocalTensor<float>& cacheDgamma); | 50 | const LocalTensor<float>& cacheDgamma); |
| 51 | __aicore__ inline void LoadMeanRstd(int32_t taskIdx); | 51 | __aicore__ inline void LoadMeanRstd(int32_t taskIdx); |
| 52 | - __aicore__ inline void LoadDataToUb( | 52 | + __aicore__ inline void LoadDataToUb(TQue<TPosition::VECIN, 1>& inQue, TBuf<TPosition::VECCALC>& tbuf, |
| 53 | - TQue<TPosition::VECIN, 1>& inQue, TBuf<TPosition::VECCALC>& tbuf, const GlobalTensor<U>& gm, | 53 | + const GlobalTensor<U>& gm, const int64_t offset, const uint32_t count); |
| 54 | - const int64_t offset, const uint32_t count); | 54 | + __aicore__ inline void CopyInDyAndX(const LocalTensor<T>& dyTensor, const LocalTensor<T>& xTensor, |
| 55 | - __aicore__ inline void CopyInDyAndX( | 55 | + const int64_t offset, const uint32_t burstLen); |
| 56 | - const LocalTensor<T>& dyTensor, const LocalTensor<T>& xTensor, const int64_t offset, const uint32_t burstLen); | 56 | + __aicore__ inline void StoreDxToGm(TQue<TPosition::VECOUT, 1>& outQue, const int64_t gmOffset, |
| 57 | - __aicore__ inline void StoreDxToGm( | 57 | + const uint32_t count); |
| 58 | - TQue<TPosition::VECOUT, 1>& outQue, const int64_t gmOffset, const uint32_t count); | ||
| 59 | template <typename OUT> | 58 | template <typename OUT> |
| 60 | - __aicore__ inline void StoreDgamma( | 59 | + __aicore__ inline void StoreDgamma(int64_t channelIdx, const GlobalTensor<OUT>& dgammaOut, |
| 61 | - int64_t channelIdx, const GlobalTensor<OUT>& dgammaOut, const LocalTensor<float>& dsTensor, | 60 | + const LocalTensor<float>& dsTensor, const LocalTensor<float>& dbetaTensor, |
| 62 | - const LocalTensor<float>& dbetaTensor, const float mean, const float rstd); | 61 | + const float mean, const float rstd); |
| 63 | - __aicore__ inline void Fp32DgammaDbeta2GM( | 62 | + __aicore__ inline void Fp32DgammaDbeta2GM(int64_t channelIdx, GlobalTensor<float>& dgammaOut, |
| 64 | - int64_t channelIdx, GlobalTensor<float>& dgammaOut, const LocalTensor<float>& dsTensor, | 63 | + const LocalTensor<float>& dsTensor, GlobalTensor<float>& dbetaOut, |
| 65 | - GlobalTensor<float>& dbetaOut, const LocalTensor<float>& dbetaTensor, const float mean, const float rstd); | 64 | + const LocalTensor<float>& dbetaTensor, const float mean, |
| 66 | - __aicore__ inline void NonFp32DgammaDbeta2GM( | 65 | + const float rstd); |
| 67 | - int64_t channelIdx, const LocalTensor<float>& dsTensor, const LocalTensor<float>& dbetaTensor, | 66 | + __aicore__ inline void NonFp32DgammaDbeta2GM(int64_t channelIdx, const LocalTensor<float>& dsTensor, |
| 68 | - const float mean, const float rstd); | 67 | + const LocalTensor<float>& dbetaTensor, const float mean, |
| 69 | - __aicore__ inline void StoreDgammaDbeta( | 68 | + const float rstd); |
| 70 | - const int32_t taskIdx, const LocalTensor<float>& dsTensor, const LocalTensor<float>& dbetaTensor, | 69 | + __aicore__ inline void StoreDgammaDbeta(const int32_t taskIdx, const LocalTensor<float>& dsTensor, |
| 71 | - const float mean, const float rstd); | 70 | + const LocalTensor<float>& dbetaTensor, const float mean, const float rstd); |
| 72 | - __aicore__ inline void ComputeSum1Sum2( | 71 | + __aicore__ inline void ComputeSum1Sum2(const LocalTensor<float>& dbetaTensor, const LocalTensor<float>& dsTensor, |
| 73 | - const LocalTensor<float>& dbetaTensor, const LocalTensor<float>& dsTensor, | 72 | + const LocalTensor<float>& gammaTensor, float& sum1, float& sum2); |
| 74 | - const LocalTensor<float>& gammaTensor, float& sum1, float& sum2); | 73 | + __aicore__ inline void VFComputeSum1Sum2(const LocalTensor<float>& dbetaTensor, const LocalTensor<float>& dsTensor, |
| 75 | - __aicore__ inline void VFComputeSum1Sum2( | 74 | + const LocalTensor<float>& gammaTensor, float& sum1, float& sum2); |
| 76 | - const LocalTensor<float>& dbetaTensor, const LocalTensor<float>& dsTensor, | 75 | + __aicore__ inline void VFComputeBinaryFoldSum1Sum2(const LocalTensor<float>& dbetaTensor, |
| 77 | - const LocalTensor<float>& gammaTensor, float& sum1, float& sum2); | 76 | + const LocalTensor<float>& dsTensor, |
| 78 | - __aicore__ inline void VFComputeBinaryFoldSum1Sum2( | 77 | + const LocalTensor<float>& gammaTensor, float& sum1, float& sum2); |
| 79 | - const LocalTensor<float>& dbetaTensor, const LocalTensor<float>& dsTensor, | 78 | + __aicore__ inline void VFDbetaDgammaBinaryFoldCommon(const LocalTensor<T>& x, const LocalTensor<T>& dy, |
| 80 | - const LocalTensor<float>& gammaTensor, float& sum1, float& sum2); | 79 | + const LocalTensor<float>& dbeta, |
| 81 | - __aicore__ inline void VFDbetaDgammaBinaryFoldCommon( | 80 | + const LocalTensor<float>& dgamma, const uint32_t outputOffset, |
| 82 | - const LocalTensor<T>& x, const LocalTensor<T>& dy, const LocalTensor<float>& dbeta, | 81 | + const uint32_t curCNum); |
| 83 | - const LocalTensor<float>& dgamma, const uint32_t outputOffset, const uint32_t curCNum); | 82 | + __aicore__ inline void VFComputeMode1DxCommon(const LocalTensor<T>& dstTensor, const LocalTensor<T>& xTensor, |
| 84 | - __aicore__ inline void VFComputeMode1DxCommon( | 83 | + const LocalTensor<T>& dyTensor, const LocalTensor<float>& gammaTensor, |
| 85 | - const LocalTensor<T>& dstTensor, const LocalTensor<T>& xTensor, const LocalTensor<T>& dyTensor, | 84 | + const float C2, const float C3, const uint32_t gammaOffset, |
| 86 | - const LocalTensor<float>& gammaTensor, const float C2, const float C3, const uint32_t gammaOffset, | 85 | + const uint32_t curCNum); |
| 87 | - const uint32_t curCNum); | ||
| 88 | __aicore__ inline void InitStage2Mode2Buffer(); | 86 | __aicore__ inline void InitStage2Mode2Buffer(); |
| 89 | __aicore__ inline void InitStage2Mode1Buffer(); | 87 | __aicore__ inline void InitStage2Mode1Buffer(); |
| 90 | __aicore__ inline void stage2Mode1Process(int64_t cOffset, uint32_t currentCNum); | 88 | __aicore__ inline void stage2Mode1Process(int64_t cOffset, uint32_t currentCNum); |
| @@ -102,21 +100,21 @@ protected: | |||
| 102 | TBuf<TPosition::VECCALC>& outTbuf, TBuf<TPosition::VECCALC>& tempbuf, | 100 | TBuf<TPosition::VECCALC>& outTbuf, TBuf<TPosition::VECCALC>& tempbuf, |
| 103 | const GlobalTensor<float>& workspace, GlobalTensor<U>& gmOut, | 101 | const GlobalTensor<float>& workspace, GlobalTensor<U>& gmOut, |
| 104 | int64_t cOffset, uint32_t currentCNum); | 102 | int64_t cOffset, uint32_t currentCNum); |
| 105 | - __aicore__ inline void reduceNMode1LessThan2(__local_mem__ float* inUbAddr, __local_mem__ float* calUbAddr, | 103 | + __aicore__ inline void reduceNMode1LessThan2(__ubuf__ float* inUbAddr, __ubuf__ float* calUbAddr, |
| 106 | uint32_t currentCNum); | 104 | uint32_t currentCNum); |
| 107 | - __aicore__ inline void reduceNMode1LessThan4(__local_mem__ float* inUbAddr, __local_mem__ float* calUbAddr, | 105 | + __aicore__ inline void reduceNMode1LessThan4(__ubuf__ float* inUbAddr, __ubuf__ float* calUbAddr, |
| 108 | uint32_t currentCNum); | 106 | uint32_t currentCNum); |
| 109 | - __aicore__ inline void reduceNMode1LessThan8(__local_mem__ float* inUbAddr, __local_mem__ float* calUbAddr, | 107 | + __aicore__ inline void reduceNMode1LessThan8(__ubuf__ float* inUbAddr, __ubuf__ float* calUbAddr, |
| 110 | uint32_t currentCNum); | 108 | uint32_t currentCNum); |
| 111 | - __aicore__ inline void reduceNMode1MoreThan8(__local_mem__ float* inUbAddr, __local_mem__ float* tempUbAddr, | 109 | + __aicore__ inline void reduceNMode1MoreThan8(__ubuf__ float* inUbAddr, __ubuf__ float* tempUbAddr, |
| 112 | - __local_mem__ float* calUbAddr, uint32_t currentCNum); | 110 | + __ubuf__ float* calUbAddr, uint32_t currentCNum); |
| 113 | - __aicore__ inline void TwoRowAddWithTail(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, | 111 | + __aicore__ inline void TwoRowAddWithTail(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 114 | uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, | 112 | uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, |
| 115 | RegTensor<float>& rem, RegTensor<float>& nextRow, | 113 | RegTensor<float>& rem, RegTensor<float>& nextRow, |
| 116 | RegTensor<float>& remNextRow); | 114 | RegTensor<float>& remNextRow); |
| 117 | - __aicore__ inline void TwoRowAdd(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, uint32_t offset1, | 115 | + __aicore__ inline void TwoRowAdd(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, uint32_t offset1, |
| 118 | uint32_t offset2, RegTensor<float>& nextRow); | 116 | uint32_t offset2, RegTensor<float>& nextRow); |
| 119 | - __aicore__ inline void FlodSumDgammaVF(__local_mem__ float* src1, __local_mem__ float* src2, int32_t num); | 117 | + __aicore__ inline void FlodSumDgammaVF(__ubuf__ float* src1, __ubuf__ float* src2, int32_t num); |
| 120 | __aicore__ inline void CopyInStage2Mode2(TQue<QuePosition::VECIN, 1>& inQueue, const GlobalTensor<float>& inGm, | 118 | __aicore__ inline void CopyInStage2Mode2(TQue<QuePosition::VECIN, 1>& inQueue, const GlobalTensor<float>& inGm, |
| 121 | int64_t offset, int64_t currentNNum, int64_t currentCNum); | 119 | int64_t offset, int64_t currentNNum, int64_t currentCNum); |
| 122 | __aicore__ inline void ProcessStage2Mode2Once(int64_t ubTimes, int64_t currentCNum, | 120 | __aicore__ inline void ProcessStage2Mode2Once(int64_t ubTimes, int64_t currentCNum, |
| @@ -449,16 +447,16 @@ __aicore__ inline void GroupNormGradBase<T, U>::CustomReduceSum(const LocalTenso | |||
| 449 | RegTensor<float> vlSum; | 447 | RegTensor<float> vlSum; |
| 450 | 448 | ||
| 451 | MaskReg pregAll = CreateMask<float, MaskPattern::ALL>(); | 449 | MaskReg pregAll = CreateMask<float, MaskPattern::ALL>(); |
| 452 | - DataCopy(x1, src + 0 * VecLen_); | 450 | + LoadAlign(x1, src + 0 * VecLen_); |
| 453 | - DataCopy(x2, src + 1 * VecLen_); | 451 | + LoadAlign(x2, src + 1 * VecLen_); |
| 454 | - DataCopy(x3, src + 2 * VecLen_); | 452 | + LoadAlign(x3, src + 2 * VecLen_); |
| 455 | - DataCopy(x4, src + 3 * VecLen_); | 453 | + LoadAlign(x4, src + 3 * VecLen_); |
| 456 | Add(sum1, x1, x3, pregAll); | 454 | Add(sum1, x1, x3, pregAll); |
| 457 | Add(sum2, x2, x4, pregAll); | 455 | Add(sum2, x2, x4, pregAll); |
| 458 | Add(sum12, sum1, sum2, pregAll); | 456 | Add(sum12, sum1, sum2, pregAll); |
| 459 | - ReduceSum(vlSum, sum12, pregAll); | 457 | + Reduce<ReduceType::SUM>(vlSum, sum12, pregAll); |
| 460 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); | 458 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); |
| 461 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dst + idx, vlSum, pregMerge); | 459 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(dst + idx, vlSum, pregMerge); |
| 462 | } | 460 | } |
| 463 | } | 461 | } |
| 464 | 462 | ||
| @@ -565,9 +563,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::StoreDxToGm(TQue<TPosition::VECO | |||
| 565 | */ | 563 | */ |
| 566 | template <typename T, typename U> | 564 | template <typename T, typename U> |
| 567 | template <typename OUT> | 565 | template <typename OUT> |
| 568 | -__aicore__ inline void GroupNormGradBase<T, U>::StoreDgamma( | 566 | +__aicore__ inline void GroupNormGradBase<T, U>::StoreDgamma(int64_t channelIdx, const GlobalTensor<OUT>& dgammaOut, |
| 569 | - int64_t channelIdx, const GlobalTensor<OUT>& dgammaOut, const LocalTensor<float>& dsTensor, | 567 | + const LocalTensor<float>& dsTensor, |
| 570 | - const LocalTensor<float>& dbetaTensor, const float mean, const float rstd) | 568 | + const LocalTensor<float>& dbetaTensor, const float mean, |
| 569 | + const float rstd) | ||
| 571 | { | 570 | { |
| 572 | auto rstdScalar = rstd; | 571 | auto rstdScalar = rstd; |
| 573 | auto negMeanScalar = -mean; | 572 | auto negMeanScalar = -mean; |
| @@ -587,8 +586,8 @@ __aicore__ inline void GroupNormGradBase<T, U>::StoreDgamma( | |||
| 587 | uint32_t sregvl = (uint32_t)this->VecLen_; | 586 | uint32_t sregvl = (uint32_t)this->VecLen_; |
| 588 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { | 587 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { |
| 589 | preg = UpdateMask<float>(sreg); | 588 | preg = UpdateMask<float>(sreg); |
| 590 | - DataCopy(vregDs, ubDs + i * sregvl); | 589 | + LoadAlign(vregDs, ubDs + i * sregvl); |
| 591 | - DataCopy(vregDbeta, ubDbeta + i * sregvl); | 590 | + LoadAlign(vregDbeta, ubDbeta + i * sregvl); |
| 592 | Axpy(vregDs, vregDbeta, negMeanScalar, preg); | 591 | Axpy(vregDs, vregDbeta, negMeanScalar, preg); |
| 593 | Muls(vregDs, vregDs, rstdScalar, preg); | 592 | Muls(vregDs, vregDs, rstdScalar, preg); |
| 594 | StoreOneTensorForDtypeT<OUT>(ubDgamma, vregDs, preg, i * sregvl); | 593 | StoreOneTensorForDtypeT<OUT>(ubDgamma, vregDs, preg, i * sregvl); |
| @@ -742,10 +741,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeBinaryFoldSum1Sum2(cons | |||
| 742 | Mul(vregDgammaQ, vregGammaQ, vregDgammaQ, pregMain); | 741 | Mul(vregDgammaQ, vregGammaQ, vregDgammaQ, pregMain); |
| 743 | MulAddDst(vregDbetaQ, vregDbetaR, vregGammaR, pregLoop); | 742 | MulAddDst(vregDbetaQ, vregDbetaR, vregGammaR, pregLoop); |
| 744 | MulAddDst(vregDgammaQ, vregDgammaR, vregGammaR, pregLoop); | 743 | MulAddDst(vregDgammaQ, vregDgammaR, vregGammaR, pregLoop); |
| 745 | - ReduceSum(vregSumDgamma, vregDgammaQ, pregLoop); | 744 | + Reduce<ReduceType::SUM>(vregSumDgamma, vregDgammaQ, pregLoop); |
| 746 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + i, vregSumDgamma, pregMerge); | 745 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + i, vregSumDgamma, pregMerge); |
| 747 | - ReduceSum(vregSumDbeta, vregDbetaQ, pregLoop); | 746 | + Reduce<ReduceType::SUM>(vregSumDbeta, vregDbetaQ, pregLoop); |
| 748 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + i, vregSumDbeta, pregMerge); | 747 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + i, vregSumDbeta, pregMerge); |
| 749 | } | 748 | } |
| 750 | // step2: the tail (last 64 or less than 64) blocks reduce to 1. | 749 | // step2: the tail (last 64 or less than 64) blocks reduce to 1. |
| 751 | { | 750 | { |
| @@ -760,14 +759,14 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeBinaryFoldSum1Sum2(cons | |||
| 760 | Mul(vregDgammaQ, vregGammaQ, vregDgammaQ, pregMain); | 759 | Mul(vregDgammaQ, vregGammaQ, vregDgammaQ, pregMain); |
| 761 | MulDstAdd(vregDbetaR, vregGammaR, vregDbetaQ, pregLoop); | 760 | MulDstAdd(vregDbetaR, vregGammaR, vregDbetaQ, pregLoop); |
| 762 | MulDstAdd(vregDgammaR, vregGammaR, vregDgammaQ, pregLoop); | 761 | MulDstAdd(vregDgammaR, vregGammaR, vregDgammaQ, pregLoop); |
| 763 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDbetaQ, vregDbetaR, pregLoop); | 762 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDbetaQ, vregDbetaR, pregLoop); |
| 764 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDgammaQ, vregDgammaR, pregLoop); | 763 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDgammaQ, vregDgammaR, pregLoop); |
| 765 | - ReduceSum(vregSumDgamma, vregDgammaQ, pregMain); | 764 | + Reduce<ReduceType::SUM>(vregSumDgamma, vregDgammaQ, pregMain); |
| 766 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderGeneral, vregSumDgamma, | 765 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderGeneral, vregSumDgamma, |
| 767 | - pregMerge); | 766 | + pregMerge); |
| 768 | - ReduceSum(vregSumDbeta, vregDbetaQ, pregMain); | 767 | + Reduce<ReduceType::SUM>(vregSumDbeta, vregDbetaQ, pregMain); |
| 769 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderGeneral, vregSumDbeta, | 768 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderGeneral, vregSumDbeta, |
| 770 | - pregMerge); | 769 | + pregMerge); |
| 771 | } | 770 | } |
| 772 | // step3: non-overlapping portions of the first half reduce by 64, this part always 64 align | 771 | // step3: non-overlapping portions of the first half reduce by 64, this part always 64 align |
| 773 | for (uint16_t i = 0; i < static_cast<uint16_t>(quotientLoop - remainderLoop); i++) { | 772 | for (uint16_t i = 0; i < static_cast<uint16_t>(quotientLoop - remainderLoop); i++) { |
| @@ -776,12 +775,12 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeBinaryFoldSum1Sum2(cons | |||
| 776 | LoadOneTensorForDtypeT<float>(ubGamma, vregGamma, pregMain, ((i + remainderLoop) * sregvl)); | 775 | LoadOneTensorForDtypeT<float>(ubGamma, vregGamma, pregMain, ((i + remainderLoop) * sregvl)); |
| 777 | Mul(vregDbeta, vregGamma, vregDbeta, pregMain); | 776 | Mul(vregDbeta, vregGamma, vregDbeta, pregMain); |
| 778 | Mul(vregDgamma, vregGamma, vregDgamma, pregMain); | 777 | Mul(vregDgamma, vregGamma, vregDgamma, pregMain); |
| 779 | - ReduceSum(vregSumDgamma, vregDgamma, pregMain); | 778 | + Reduce<ReduceType::SUM>(vregSumDgamma, vregDgamma, pregMain); |
| 780 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderLoop + i, vregSumDgamma, | 779 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderLoop + i, vregSumDgamma, |
| 781 | - pregMerge); | 780 | + pregMerge); |
| 782 | - ReduceSum(vregSumDbeta, vregDbeta, pregMain); | 781 | + Reduce<ReduceType::SUM>(vregSumDbeta, vregDbeta, pregMain); |
| 783 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderLoop + i, vregSumDbeta, | 782 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderLoop + i, vregSumDbeta, |
| 784 | - pregMerge); | 783 | + pregMerge); |
| 785 | } | 784 | } |
| 786 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 785 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 787 | // step4: binary folding reduce calculation | 786 | // step4: binary folding reduce calculation |
| @@ -790,14 +789,15 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeBinaryFoldSum1Sum2(cons | |||
| 790 | for (uint16_t i = 0; i < binaryCGKLoop; i++) { | 789 | for (uint16_t i = 0; i < binaryCGKLoop; i++) { |
| 791 | curBinaryCGLoop = curBinaryCGLoop / 2; | 790 | curBinaryCGLoop = curBinaryCGLoop / 2; |
| 792 | for (uint16_t j = 0; j < curBinaryCGLoop; j++) { | 791 | for (uint16_t j = 0; j < curBinaryCGLoop; j++) { |
| 793 | - DataCopy(vregDgammaQ, ((__ubuf__ float*)ubBinaryDgamma + j * sregvl)); | 792 | + LoadAlign(vregDgammaQ, ((__ubuf__ float*)ubBinaryDgamma + j * sregvl)); |
| 794 | - DataCopy(vregDgammaR, ((__ubuf__ float*)ubBinaryDgamma + (j + curBinaryCGLoop) * sregvl)); | 793 | + LoadAlign(vregDgammaR, ((__ubuf__ float*)ubBinaryDgamma + (j + curBinaryCGLoop) * sregvl)); |
| 795 | Add(vregDgammaQ, vregDgammaQ, vregDgammaR, pregMain); | 794 | Add(vregDgammaQ, vregDgammaQ, vregDgammaR, pregMain); |
| 796 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + j * sregvl, vregDgammaQ, pregMain); | 795 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + j * sregvl, vregDgammaQ, |
| 797 | - DataCopy(vregDbetaQ, ((__ubuf__ float*)ubBinaryDbeta + j * sregvl)); | 796 | + pregMain); |
| 798 | - DataCopy(vregDbetaR, ((__ubuf__ float*)ubBinaryDbeta + (j + curBinaryCGLoop) * sregvl)); | 797 | + LoadAlign(vregDbetaQ, ((__ubuf__ float*)ubBinaryDbeta + j * sregvl)); |
| 798 | + LoadAlign(vregDbetaR, ((__ubuf__ float*)ubBinaryDbeta + (j + curBinaryCGLoop) * sregvl)); | ||
| 799 | Add(vregDbetaQ, vregDbetaQ, vregDbetaR, pregMain); | 799 | Add(vregDbetaQ, vregDbetaQ, vregDbetaR, pregMain); |
| 800 | - DataCopy(ubBinaryDbeta + j * sregvl, vregDbetaQ, pregMain); | 800 | + StoreAlign(ubBinaryDbeta + j * sregvl, vregDbetaQ, pregMain); |
| 801 | } | 801 | } |
| 802 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 802 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 803 | } | 803 | } |
| @@ -805,12 +805,12 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeBinaryFoldSum1Sum2(cons | |||
| 805 | { | 805 | { |
| 806 | uint32_t sreg2 = binaryCGLastNum; | 806 | uint32_t sreg2 = binaryCGLastNum; |
| 807 | MaskReg pregLoop = UpdateMask<float>(sreg2); | 807 | MaskReg pregLoop = UpdateMask<float>(sreg2); |
| 808 | - DataCopy(vregDgamma, ((__ubuf__ float*)ubBinaryDgamma)); | 808 | + LoadAlign(vregDgamma, ((__ubuf__ float*)ubBinaryDgamma)); |
| 809 | - ReduceSum(vregDgamma, vregDgamma, pregLoop); | 809 | + Reduce<ReduceType::SUM>(vregDgamma, vregDgamma, pregLoop); |
| 810 | - DataCopy(outDgamma, vregDgamma, pregMerge); | 810 | + StoreAlign(outDgamma, vregDgamma, pregMerge); |
| 811 | - DataCopy(vregDbeta, ((__ubuf__ float*)ubBinaryDbeta)); | 811 | + LoadAlign(vregDbeta, ((__ubuf__ float*)ubBinaryDbeta)); |
| 812 | - ReduceSum(vregDbeta, vregDbeta, pregLoop); | 812 | + Reduce<ReduceType::SUM>(vregDbeta, vregDbeta, pregLoop); |
| 813 | - DataCopy(outDbeta, vregDbeta, pregMerge); | 813 | + StoreAlign(outDbeta, vregDbeta, pregMerge); |
| 814 | } | 814 | } |
| 815 | } | 815 | } |
| 816 | outQueDgamma_.FreeTensor(binaryDbetaTensor); | 816 | outQueDgamma_.FreeTensor(binaryDbetaTensor); |
| @@ -822,8 +822,8 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeBinaryFoldSum1Sum2(cons | |||
| 822 | } | 822 | } |
| 823 | 823 | ||
| 824 | /* | 824 | /* |
| 825 | - sum1 = ReduceSum(dgamma * gamma) / D * HxW | 825 | + sum1 = Reduce<ReduceType::SUM>(dgamma * gamma) / D * HxW |
| 826 | - sum2 = ReduceSum(dbeta * gamma) / D * HxW | 826 | + sum2 = Reduce<ReduceType::SUM>(dbeta * gamma) / D * HxW |
| 827 | */ | 827 | */ |
| 828 | template <typename T, typename U> | 828 | template <typename T, typename U> |
| 829 | __aicore__ inline void GroupNormGradBase<T, U>::VFComputeSum1Sum2(const LocalTensor<float>& dbetaTensor, | 829 | __aicore__ inline void GroupNormGradBase<T, U>::VFComputeSum1Sum2(const LocalTensor<float>& dbetaTensor, |
| @@ -857,19 +857,19 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeSum1Sum2(const LocalTen | |||
| 857 | Duplicate(vregSumDs, 0, pregAll); | 857 | Duplicate(vregSumDs, 0, pregAll); |
| 858 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { | 858 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { |
| 859 | MaskReg preg = UpdateMask<float>(sreg); | 859 | MaskReg preg = UpdateMask<float>(sreg); |
| 860 | - DataCopy(vregGamma, ubGamma + i * sregvl); | 860 | + LoadAlign(vregGamma, ubGamma + i * sregvl); |
| 861 | - DataCopy(vregDbeta, ubDbeta + i * sregvl); | 861 | + LoadAlign(vregDbeta, ubDbeta + i * sregvl); |
| 862 | - DataCopy(vregDs, ubDs + i * sregvl); | 862 | + LoadAlign(vregDs, ubDs + i * sregvl); |
| 863 | MulDstAdd(vregDbeta, vregGamma, vregSumDbeta, preg); | 863 | MulDstAdd(vregDbeta, vregGamma, vregSumDbeta, preg); |
| 864 | MulDstAdd(vregDs, vregGamma, vregSumDs, preg); | 864 | MulDstAdd(vregDs, vregGamma, vregSumDs, preg); |
| 865 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregSumDbeta, vregDbeta, preg); | 865 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregSumDbeta, vregDbeta, preg); |
| 866 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregSumDs, vregDs, preg); | 866 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregSumDs, vregDs, preg); |
| 867 | } | 867 | } |
| 868 | - ReduceSum(vregSumDbeta, vregSumDbeta, pregAll); | 868 | + Reduce<ReduceType::SUM>(vregSumDbeta, vregSumDbeta, pregAll); |
| 869 | - ReduceSum(vregSumDs, vregSumDs, pregAll); | 869 | + Reduce<ReduceType::SUM>(vregSumDs, vregSumDs, pregAll); |
| 870 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); | 870 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); |
| 871 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(outDbeta, vregSumDbeta, pregMerge); | 871 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(outDbeta, vregSumDbeta, pregMerge); |
| 872 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(outDs, vregSumDs, pregMerge); | 872 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(outDs, vregSumDs, pregMerge); |
| 873 | } | 873 | } |
| 874 | TEventID eventIDVtoS1 = GetTPipePtr()->FetchEventID(HardEvent::V_S); | 874 | TEventID eventIDVtoS1 = GetTPipePtr()->FetchEventID(HardEvent::V_S); |
| 875 | SetFlag<HardEvent::V_S>(eventIDVtoS1); | 875 | SetFlag<HardEvent::V_S>(eventIDVtoS1); |
| @@ -985,8 +985,8 @@ __aicore__ inline void GroupNormGradBase<T, U>::stage2Mode1B32Compute( | |||
| 985 | reduceNMode1Wsp2Ub(inQue, workspace, cOffset, currentCNum); | 985 | reduceNMode1Wsp2Ub(inQue, workspace, cOffset, currentCNum); |
| 986 | LocalTensor<float> inUb = inQue.template DeQue<float>(); | 986 | LocalTensor<float> inUb = inQue.template DeQue<float>(); |
| 987 | LocalTensor<float> calUb = calQue.AllocTensor<float>(); | 987 | LocalTensor<float> calUb = calQue.AllocTensor<float>(); |
| 988 | - __local_mem__ float* inUbAddr = (__local_mem__ float*)inUb.GetPhyAddr(); | 988 | + __ubuf__ float* inUbAddr = (__ubuf__ float*)inUb.GetPhyAddr(); |
| 989 | - __local_mem__ float* calUbAddr = (__local_mem__ float*)calUb.GetPhyAddr(); | 989 | + __ubuf__ float* calUbAddr = (__ubuf__ float*)calUb.GetPhyAddr(); |
| 990 | if (reduceNCnt_ <= SCALE_COEF_TWO) { | 990 | if (reduceNCnt_ <= SCALE_COEF_TWO) { |
| 991 | reduceNMode1LessThan2(inUbAddr, calUbAddr, currentCNum); | 991 | reduceNMode1LessThan2(inUbAddr, calUbAddr, currentCNum); |
| 992 | } else if (reduceNCnt_ <= SCALE_COEF_FOUR) { | 992 | } else if (reduceNCnt_ <= SCALE_COEF_FOUR) { |
| @@ -995,7 +995,7 @@ __aicore__ inline void GroupNormGradBase<T, U>::stage2Mode1B32Compute( | |||
| 995 | reduceNMode1LessThan8(inUbAddr, calUbAddr, currentCNum); | 995 | reduceNMode1LessThan8(inUbAddr, calUbAddr, currentCNum); |
| 996 | } else { | 996 | } else { |
| 997 | LocalTensor<float> tempUb = tempTbuf.AllocTensor<float>(); | 997 | LocalTensor<float> tempUb = tempTbuf.AllocTensor<float>(); |
| 998 | - __local_mem__ float* tempUbAddr = (__local_mem__ float*)tempUb.GetPhyAddr(); | 998 | + __ubuf__ float* tempUbAddr = (__ubuf__ float*)tempUb.GetPhyAddr(); |
| 999 | reduceNMode1MoreThan8(inUbAddr, tempUbAddr, calUbAddr, currentCNum); | 999 | reduceNMode1MoreThan8(inUbAddr, tempUbAddr, calUbAddr, currentCNum); |
| 1000 | } | 1000 | } |
| 1001 | inQue.FreeTensor(inUb); | 1001 | inQue.FreeTensor(inUb); |
| @@ -1015,10 +1015,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::stage2Mode1B16Compute( | |||
| 1015 | reduceNMode1Wsp2Ub(inQue, workspace, cOffset, currentCNum); | 1015 | reduceNMode1Wsp2Ub(inQue, workspace, cOffset, currentCNum); |
| 1016 | LocalTensor<float> inUb = inQue.template DeQue<float>(); | 1016 | LocalTensor<float> inUb = inQue.template DeQue<float>(); |
| 1017 | LocalTensor<U> calUb = calQue.AllocTensor<U>(); | 1017 | LocalTensor<U> calUb = calQue.AllocTensor<U>(); |
| 1018 | - __local_mem__ float* inUbAddr = (__local_mem__ float*)inUb.GetPhyAddr(); | 1018 | + __ubuf__ float* inUbAddr = (__ubuf__ float*)inUb.GetPhyAddr(); |
| 1019 | - __local_mem__ U* calUbAddr = (__local_mem__ U*)calUb.GetPhyAddr(); | 1019 | + __ubuf__ U* calUbAddr = (__ubuf__ U*)calUb.GetPhyAddr(); |
| 1020 | LocalTensor<float> outTbufUb = outTbuf.Get<float>(); | 1020 | LocalTensor<float> outTbufUb = outTbuf.Get<float>(); |
| 1021 | - __local_mem__ float* outTbufUbAddr = (__local_mem__ float*)outTbufUb.GetPhyAddr(); | 1021 | + __ubuf__ float* outTbufUbAddr = (__ubuf__ float*)outTbufUb.GetPhyAddr(); |
| 1022 | if (reduceNCnt_ <= SCALE_COEF_TWO) { | 1022 | if (reduceNCnt_ <= SCALE_COEF_TWO) { |
| 1023 | reduceNMode1LessThan2(inUbAddr, outTbufUbAddr, currentCNum); | 1023 | reduceNMode1LessThan2(inUbAddr, outTbufUbAddr, currentCNum); |
| 1024 | } else if (reduceNCnt_ <= SCALE_COEF_FOUR) { | 1024 | } else if (reduceNCnt_ <= SCALE_COEF_FOUR) { |
| @@ -1027,10 +1027,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::stage2Mode1B16Compute( | |||
| 1027 | reduceNMode1LessThan8(inUbAddr, outTbufUbAddr, currentCNum); | 1027 | reduceNMode1LessThan8(inUbAddr, outTbufUbAddr, currentCNum); |
| 1028 | } else { | 1028 | } else { |
| 1029 | LocalTensor<float> tempUb = tempTbuf.AllocTensor<float>(); | 1029 | LocalTensor<float> tempUb = tempTbuf.AllocTensor<float>(); |
| 1030 | - __local_mem__ float* tempUbAddr = (__local_mem__ float*)tempUb.GetPhyAddr(); | 1030 | + __ubuf__ float* tempUbAddr = (__ubuf__ float*)tempUb.GetPhyAddr(); |
| 1031 | reduceNMode1MoreThan8(inUbAddr, tempUbAddr, outTbufUbAddr, currentCNum); | 1031 | reduceNMode1MoreThan8(inUbAddr, tempUbAddr, outTbufUbAddr, currentCNum); |
| 1032 | } | 1032 | } |
| 1033 | - VFCastFloat2T<U>((__local_mem__ U*)calUbAddr, (__local_mem__ float*)outTbufUbAddr, currentCNum, this->VecLen_); | 1033 | + VFCastFloat2T<U>((__ubuf__ U*)calUbAddr, (__ubuf__ float*)outTbufUbAddr, currentCNum, this->VecLen_); |
| 1034 | inQue.FreeTensor(inUb); | 1034 | inQue.FreeTensor(inUb); |
| 1035 | calQue.EnQue(calUb); | 1035 | calQue.EnQue(calUb); |
| 1036 | LocalTensor<U> calOutUb = calQue.template DeQue<U>(); | 1036 | LocalTensor<U> calOutUb = calQue.template DeQue<U>(); |
| @@ -1065,7 +1065,7 @@ __aicore__ inline void GroupNormGradBase<T, U>::ProcessStage2Mode2Once(int64_t u | |||
| 1065 | CopyInStage2Mode2(inQueDgammaChannel_, workspace[baseOffset], sum1Offset, this->nFactorStage2Mode2, | 1065 | CopyInStage2Mode2(inQueDgammaChannel_, workspace[baseOffset], sum1Offset, this->nFactorStage2Mode2, |
| 1066 | currentCNum); | 1066 | currentCNum); |
| 1067 | LocalTensor<float> sum1Ub = inQueDgammaChannel_.DeQue<float>(); | 1067 | LocalTensor<float> sum1Ub = inQueDgammaChannel_.DeQue<float>(); |
| 1068 | - __local_mem__ float* sum1Local = (__local_mem__ float*)sum1Ub.GetPhyAddr(); | 1068 | + __ubuf__ float* sum1Local = (__ubuf__ float*)sum1Ub.GetPhyAddr(); |
| 1069 | 1069 | ||
| 1070 | int64_t sum2Offset = ubTimes * this->cFactorStage2Mode2 + | 1070 | int64_t sum2Offset = ubTimes * this->cFactorStage2Mode2 + |
| 1071 | (basicBlockIdx + this->nLoop) * this->nFactorStage2Mode2 * this->C_; | 1071 | (basicBlockIdx + this->nLoop) * this->nFactorStage2Mode2 * this->C_; |
| @@ -1073,13 +1073,13 @@ __aicore__ inline void GroupNormGradBase<T, U>::ProcessStage2Mode2Once(int64_t u | |||
| 1073 | CopyInStage2Mode2(inQueDgammaChannel_, workspace[baseOffset], sum2Offset, this->nFactorStage2Mode2, | 1073 | CopyInStage2Mode2(inQueDgammaChannel_, workspace[baseOffset], sum2Offset, this->nFactorStage2Mode2, |
| 1074 | currentCNum); | 1074 | currentCNum); |
| 1075 | LocalTensor<float> sum2Ub = inQueDgammaChannel_.DeQue<float>(); | 1075 | LocalTensor<float> sum2Ub = inQueDgammaChannel_.DeQue<float>(); |
| 1076 | - __local_mem__ float* sum2Local = (__local_mem__ float*)sum2Ub.GetPhyAddr(); | 1076 | + __ubuf__ float* sum2Local = (__ubuf__ float*)sum2Ub.GetPhyAddr(); |
| 1077 | FlodSumDgammaVF(sum1Local, sum2Local, this->nFactorStage2Mode2 * currentCAlignNum); | 1077 | FlodSumDgammaVF(sum1Local, sum2Local, this->nFactorStage2Mode2 * currentCAlignNum); |
| 1078 | inQueDgammaChannel_.FreeTensor(sum2Ub); | 1078 | inQueDgammaChannel_.FreeTensor(sum2Ub); |
| 1079 | } else if ((basicBlockIdx == this->nMainFlodCount) && (this->nTail > 0)) { | 1079 | } else if ((basicBlockIdx == this->nMainFlodCount) && (this->nTail > 0)) { |
| 1080 | CopyInStage2Mode2(inQueDgammaChannel_, workspace[baseOffset], sum2Offset, this->nTail, currentCNum); | 1080 | CopyInStage2Mode2(inQueDgammaChannel_, workspace[baseOffset], sum2Offset, this->nTail, currentCNum); |
| 1081 | LocalTensor<float> sum2Ub = inQueDgammaChannel_.DeQue<float>(); | 1081 | LocalTensor<float> sum2Ub = inQueDgammaChannel_.DeQue<float>(); |
| 1082 | - __local_mem__ float* sum2Local = (__local_mem__ float*)sum2Ub.GetPhyAddr(); | 1082 | + __ubuf__ float* sum2Local = (__ubuf__ float*)sum2Ub.GetPhyAddr(); |
| 1083 | FlodSumDgammaVF(sum1Local, sum2Local, this->nTail * currentCAlignNum); | 1083 | FlodSumDgammaVF(sum1Local, sum2Local, this->nTail * currentCAlignNum); |
| 1084 | inQueDgammaChannel_.FreeTensor(sum2Ub); | 1084 | inQueDgammaChannel_.FreeTensor(sum2Ub); |
| 1085 | } | 1085 | } |
| @@ -1098,11 +1098,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::ProcessStage2Mode2Once(int64_t u | |||
| 1098 | Aligned(static_cast<int64_t>(this->cFactorStage2Mode2), | 1098 | Aligned(static_cast<int64_t>(this->cFactorStage2Mode2), |
| 1099 | static_cast<int64_t>(GetUbBlockSize() / sizeof(float)))); | 1099 | static_cast<int64_t>(GetUbBlockSize() / sizeof(float)))); |
| 1100 | } else { | 1100 | } else { |
| 1101 | - __local_mem__ U* outDgammaLocal = (__local_mem__ U*)outDgammaTensor.GetPhyAddr(); | 1101 | + __ubuf__ U* outDgammaLocal = (__ubuf__ U*)outDgammaTensor.GetPhyAddr(); |
| 1102 | - __local_mem__ float* | 1102 | + __ubuf__ float* sumDgammaResultCacheLocal = (__ubuf__ float*)sumDgammaAllCacheUb[this->resultCacheId * |
| 1103 | - sumDgammaResultCacheLocal = (__local_mem__ float*) | 1103 | + this->cFactorStage2Mode2] |
| 1104 | - sumDgammaAllCacheUb[this->resultCacheId * this->cFactorStage2Mode2] | 1104 | + .GetPhyAddr(); |
| 1105 | - .GetPhyAddr(); | ||
| 1106 | VFCastFloat2T(outDgammaLocal, sumDgammaResultCacheLocal, this->cFactorStage2Mode2, GetVRegSize()); | 1105 | VFCastFloat2T(outDgammaLocal, sumDgammaResultCacheLocal, this->cFactorStage2Mode2, GetVRegSize()); |
| 1107 | } | 1106 | } |
| 1108 | DataCopyExtParams copyOutParams; | 1107 | DataCopyExtParams copyOutParams; |
| @@ -1117,9 +1116,8 @@ __aicore__ inline void GroupNormGradBase<T, U>::ProcessStage2Mode2Once(int64_t u | |||
| 1117 | } | 1116 | } |
| 1118 | 1117 | ||
| 1119 | template <typename T, typename U> | 1118 | template <typename T, typename U> |
| 1120 | -__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan2(__local_mem__ float* inUbAddr, | 1119 | +__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan2(__ubuf__ float* inUbAddr, |
| 1121 | - __local_mem__ float* calUbAddr, | 1120 | + __ubuf__ float* calUbAddr, uint32_t currentCNum) |
| 1122 | - uint32_t currentCNum) | ||
| 1123 | { | 1121 | { |
| 1124 | uint32_t rStride = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / sizeof(float); | 1122 | uint32_t rStride = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / sizeof(float); |
| 1125 | uint16_t rLoopCount = static_cast<uint16_t>(reduceNCnt_); | 1123 | uint16_t rLoopCount = static_cast<uint16_t>(reduceNCnt_); |
| @@ -1136,18 +1134,17 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan2(__local_me | |||
| 1136 | pregLoop = UpdateMask<float>(sreg0); | 1134 | pregLoop = UpdateMask<float>(sreg0); |
| 1137 | Duplicate(sum, 0.0, pregLoop); | 1135 | Duplicate(sum, 0.0, pregLoop); |
| 1138 | for (uint16_t i = 0; i < rLoopCount; i++) { | 1136 | for (uint16_t i = 0; i < rLoopCount; i++) { |
| 1139 | - DataCopy(inld, ((__local_mem__ float*)inUbAddr + i * rStride + k * VecLen_)); | 1137 | + LoadAlign(inld, ((__ubuf__ float*)inUbAddr + i * rStride + k * VecLen_)); |
| 1140 | Add(sum, sum, inld, pregLoop); | 1138 | Add(sum, sum, inld, pregLoop); |
| 1141 | } | 1139 | } |
| 1142 | - DataCopy(((__local_mem__ float*)calUbAddr + k * VecLen_), sum, pregLoop); | 1140 | + StoreAlign(((__ubuf__ float*)calUbAddr + k * VecLen_), sum, pregLoop); |
| 1143 | } | 1141 | } |
| 1144 | } | 1142 | } |
| 1145 | } | 1143 | } |
| 1146 | 1144 | ||
| 1147 | template <typename T, typename U> | 1145 | template <typename T, typename U> |
| 1148 | -__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan4(__local_mem__ float* inUbAddr, | 1146 | +__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan4(__ubuf__ float* inUbAddr, |
| 1149 | - __local_mem__ float* calUbAddr, | 1147 | + __ubuf__ float* calUbAddr, uint32_t currentCNum) |
| 1150 | - uint32_t currentCNum) | ||
| 1151 | { | 1148 | { |
| 1152 | uint32_t currentANumAlign = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / | 1149 | uint32_t currentANumAlign = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / |
| 1153 | sizeof(float); | 1150 | sizeof(float); |
| @@ -1175,19 +1172,18 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan4(__local_me | |||
| 1175 | for (uint16_t k = 0; k < aLoopCount; k++) { | 1172 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 1176 | pregLoop = UpdateMask<float>(sreg0); | 1173 | pregLoop = UpdateMask<float>(sreg0); |
| 1177 | uint32_t aLoopOffset = k * VecLen_; | 1174 | uint32_t aLoopOffset = k * VecLen_; |
| 1178 | - DataCopy(((__local_mem__ float*)inUbAddr + validNumInXUb + aLoopOffset), zero, pregLoop); | 1175 | + StoreAlign(((__ubuf__ float*)inUbAddr + validNumInXUb + aLoopOffset), zero, pregLoop); |
| 1179 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 1176 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 1180 | TwoRowAddWithTail(x1, inUbAddr, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 1177 | TwoRowAddWithTail(x1, inUbAddr, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 1181 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow); | 1178 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow); |
| 1182 | - DataCopy(((__local_mem__ float*)calUbAddr + aLoopOffset), x1, pregLoop); | 1179 | + StoreAlign(((__ubuf__ float*)calUbAddr + aLoopOffset), x1, pregLoop); |
| 1183 | } | 1180 | } |
| 1184 | } | 1181 | } |
| 1185 | } | 1182 | } |
| 1186 | 1183 | ||
| 1187 | template <typename T, typename U> | 1184 | template <typename T, typename U> |
| 1188 | -__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan8(__local_mem__ float* inUbAddr, | 1185 | +__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan8(__ubuf__ float* inUbAddr, |
| 1189 | - __local_mem__ float* calUbAddr, | 1186 | + __ubuf__ float* calUbAddr, uint32_t currentCNum) |
| 1190 | - uint32_t currentCNum) | ||
| 1191 | { | 1187 | { |
| 1192 | uint32_t currentANumAlign = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / | 1188 | uint32_t currentANumAlign = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / |
| 1193 | sizeof(float); | 1189 | sizeof(float); |
| @@ -1220,7 +1216,7 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan8(__local_me | |||
| 1220 | for (uint16_t k = 0; k < aLoopCount; k++) { | 1216 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 1221 | pregLoop = UpdateMask<float>(sreg0); | 1217 | pregLoop = UpdateMask<float>(sreg0); |
| 1222 | uint32_t aLoopOffset = k * VecLen_; | 1218 | uint32_t aLoopOffset = k * VecLen_; |
| 1223 | - DataCopy(((__local_mem__ float*)inUbAddr + validNumInXUb + aLoopOffset), zero, pregLoop); | 1219 | + StoreAlign(((__ubuf__ float*)inUbAddr + validNumInXUb + aLoopOffset), zero, pregLoop); |
| 1224 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 1220 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 1225 | TwoRowAddWithTail(x1, inUbAddr, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 1221 | TwoRowAddWithTail(x1, inUbAddr, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 1226 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow); | 1222 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, remNextRow); |
| @@ -1228,16 +1224,15 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1LessThan8(__local_me | |||
| 1228 | remainderTailOffset2 + aLoopOffset, ROW_THREE_OFFSET * aLength + aLoopOffset, | 1224 | remainderTailOffset2 + aLoopOffset, ROW_THREE_OFFSET * aLength + aLoopOffset, |
| 1229 | remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow); | 1225 | remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow); |
| 1230 | Add(x1, x1, x2, pregLoop); | 1226 | Add(x1, x1, x2, pregLoop); |
| 1231 | - DataCopy(((__local_mem__ float*)calUbAddr + aLoopOffset), x1, pregLoop); | 1227 | + StoreAlign(((__ubuf__ float*)calUbAddr + aLoopOffset), x1, pregLoop); |
| 1232 | } | 1228 | } |
| 1233 | } | 1229 | } |
| 1234 | } | 1230 | } |
| 1235 | 1231 | ||
| 1236 | template <typename T, typename U> | 1232 | template <typename T, typename U> |
| 1237 | -__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1MoreThan8(__local_mem__ float* inUbAddr, | 1233 | +__aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1MoreThan8(__ubuf__ float* inUbAddr, |
| 1238 | - __local_mem__ float* tempUbAddr, | 1234 | + __ubuf__ float* tempUbAddr, |
| 1239 | - __local_mem__ float* calUbAddr, | 1235 | + __ubuf__ float* calUbAddr, uint32_t currentCNum) |
| 1240 | - uint32_t currentCNum) | ||
| 1241 | { | 1236 | { |
| 1242 | uint32_t currentANumAlign = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / | 1237 | uint32_t currentANumAlign = (((currentCNum * sizeof(float) + BLOCK_BYTES - 1) / BLOCK_BYTES) * BLOCK_BYTES) / |
| 1243 | sizeof(float); | 1238 | sizeof(float); |
| @@ -1299,7 +1294,7 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1MoreThan8(__local_me | |||
| 1299 | for (uint16_t k = 0; k < aLoopCount; k++) { | 1294 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 1300 | pregLoop = UpdateMask<float>(sreg0); | 1295 | pregLoop = UpdateMask<float>(sreg0); |
| 1301 | uint32_t aLoopOffset = k * VecLen_; | 1296 | uint32_t aLoopOffset = k * VecLen_; |
| 1302 | - DataCopy(((__local_mem__ float*)inUbAddr + validNumInXUb + aLoopOffset), zero, pregLoop); | 1297 | + StoreAlign(((__ubuf__ float*)inUbAddr + validNumInXUb + aLoopOffset), zero, pregLoop); |
| 1303 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 1298 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 1304 | // 前半部分与后半部分中，都为8行的部分 | 1299 | // 前半部分与后半部分中，都为8行的部分 |
| 1305 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainderLoopCount - 1); i++) { | 1300 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainderLoopCount - 1); i++) { |
| @@ -1319,7 +1314,7 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1MoreThan8(__local_me | |||
| 1319 | remOffset + rowSevenOffset * aLength, rem, nextRow, remNextRow); | 1314 | remOffset + rowSevenOffset * aLength, rem, nextRow, remNextRow); |
| 1320 | Add(x3, x3, x4, pregLoop); | 1315 | Add(x3, x3, x4, pregLoop); |
| 1321 | Add(x1, x1, x3, pregLoop); | 1316 | Add(x1, x1, x3, pregLoop); |
| 1322 | - DataCopy(((__local_mem__ float*)tempUbAddr + i * aLength + aLoopOffset), x1, pregLoop); | 1317 | + StoreAlign(((__ubuf__ float*)tempUbAddr + i * aLength + aLoopOffset), x1, pregLoop); |
| 1323 | } | 1318 | } |
| 1324 | // 前半部分为8行，后半部分可能不足8行 | 1319 | // 前半部分为8行，后半部分可能不足8行 |
| 1325 | { | 1320 | { |
| @@ -1341,8 +1336,8 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1MoreThan8(__local_me | |||
| 1341 | remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow); | 1336 | remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow); |
| 1342 | Add(x3, x3, x4, pregLoop); | 1337 | Add(x3, x3, x4, pregLoop); |
| 1343 | Add(x1, x1, x3, pregLoop); | 1338 | Add(x1, x1, x3, pregLoop); |
| 1344 | - DataCopy(((__local_mem__ float*)tempUbAddr + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, | 1339 | + StoreAlign(((__ubuf__ float*)tempUbAddr + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, |
| 1345 | - pregLoop); | 1340 | + pregLoop); |
| 1346 | } | 1341 | } |
| 1347 | // 剩余的前半部分，一次for循环，处理8行 | 1342 | // 剩余的前半部分，一次for循环，处理8行 |
| 1348 | for (uint16_t i = 0; i < quotientLoopCount; i++) { | 1343 | for (uint16_t i = 0; i < quotientLoopCount; i++) { |
| @@ -1357,51 +1352,49 @@ __aicore__ inline void GroupNormGradBase<T, U>::reduceNMode1MoreThan8(__local_me | |||
| 1357 | baseOffset + rowSevenOffset * aLength, nextRow); | 1352 | baseOffset + rowSevenOffset * aLength, nextRow); |
| 1358 | Add(x3, x3, x4, pregLoop); | 1353 | Add(x3, x3, x4, pregLoop); |
| 1359 | Add(x1, x1, x3, pregLoop); | 1354 | Add(x1, x1, x3, pregLoop); |
| 1360 | - DataCopy(((__local_mem__ float*)tempUbAddr + (remainderLoopCount + i) * aLength + aLoopOffset), x1, | 1355 | + StoreAlign(((__ubuf__ float*)tempUbAddr + (remainderLoopCount + i) * aLength + aLoopOffset), x1, |
| 1361 | - pregLoop); | 1356 | + pregLoop); |
| 1362 | } | 1357 | } |
| 1363 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 1358 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 1364 | uint16_t curBinaryAddInnerLoop = binaryAddInnerLoop; | 1359 | uint16_t curBinaryAddInnerLoop = binaryAddInnerLoop; |
| 1365 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { | 1360 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { |
| 1366 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / 2; | 1361 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / 2; |
| 1367 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { | 1362 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { |
| 1368 | - DataCopy(x1, ((__local_mem__ float*)tempUbAddr) + j * aLength + aLoopOffset); | 1363 | + LoadAlign(x1, ((__ubuf__ float*)tempUbAddr) + j * aLength + aLoopOffset); |
| 1369 | - DataCopy(x2, | 1364 | + LoadAlign(x2, ((__ubuf__ float*)tempUbAddr) + (j + curBinaryAddInnerLoop) * aLength + aLoopOffset); |
| 1370 | - ((__local_mem__ float*)tempUbAddr) + (j + curBinaryAddInnerLoop) * aLength + aLoopOffset); | ||
| 1371 | Add(x1, x1, x2, pregLoop); | 1365 | Add(x1, x1, x2, pregLoop); |
| 1372 | - DataCopy(((__local_mem__ float*)tempUbAddr + j * aLength + aLoopOffset), x1, pregLoop); | 1366 | + StoreAlign(((__ubuf__ float*)tempUbAddr + j * aLength + aLoopOffset), x1, pregLoop); |
| 1373 | } | 1367 | } |
| 1374 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 1368 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 1375 | } | 1369 | } |
| 1376 | - DataCopy(x1, ((__local_mem__ float*)tempUbAddr) + aLoopOffset); | 1370 | + LoadAlign(x1, ((__ubuf__ float*)tempUbAddr) + aLoopOffset); |
| 1377 | - DataCopy(((__local_mem__ float*)calUbAddr + aLoopOffset), x1, pregLoop); | 1371 | + StoreAlign(((__ubuf__ float*)calUbAddr + aLoopOffset), x1, pregLoop); |
| 1378 | } | 1372 | } |
| 1379 | } | 1373 | } |
| 1380 | } | 1374 | } |
| 1381 | 1375 | ||
| 1382 | template <typename T, typename U> | 1376 | template <typename T, typename U> |
| 1383 | -__aicore__ inline void GroupNormGradBase<T, U>::TwoRowAddWithTail(RegTensor<float>& dst, __local_mem__ float* input, | 1377 | +__aicore__ inline void GroupNormGradBase<T, U>::TwoRowAddWithTail(RegTensor<float>& dst, __ubuf__ float* input, |
| 1384 | MaskReg& preg, uint32_t offset1, uint32_t offset2, | 1378 | MaskReg& preg, uint32_t offset1, uint32_t offset2, |
| 1385 | uint32_t offset3, uint32_t offset4, | 1379 | uint32_t offset3, uint32_t offset4, |
| 1386 | RegTensor<float>& rem, RegTensor<float>& nextRow, | 1380 | RegTensor<float>& rem, RegTensor<float>& nextRow, |
| 1387 | RegTensor<float>& remNextRow) | 1381 | RegTensor<float>& remNextRow) |
| 1388 | { | 1382 | { |
| 1389 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 1383 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 1390 | - DataCopy(rem, ((__local_mem__ float*)(input) + (offset2))); | 1384 | + LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2))); |
| 1391 | Add(dst, dst, rem, preg); | 1385 | Add(dst, dst, rem, preg); |
| 1392 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset3))); | 1386 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset3))); |
| 1393 | - DataCopy(remNextRow, ((__local_mem__ float*)(input) + (offset4))); | 1387 | + LoadAlign(remNextRow, ((__ubuf__ float*)(input) + (offset4))); |
| 1394 | Add(nextRow, nextRow, remNextRow, preg); | 1388 | Add(nextRow, nextRow, remNextRow, preg); |
| 1395 | Add(dst, dst, nextRow, preg); | 1389 | Add(dst, dst, nextRow, preg); |
| 1396 | } | 1390 | } |
| 1397 | 1391 | ||
| 1398 | template <typename T, typename U> | 1392 | template <typename T, typename U> |
| 1399 | -__aicore__ inline void GroupNormGradBase<T, U>::TwoRowAdd(RegTensor<float>& dst, __local_mem__ float* input, | 1393 | +__aicore__ inline void GroupNormGradBase<T, U>::TwoRowAdd(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 1400 | - MaskReg& preg, uint32_t offset1, uint32_t offset2, | 1394 | + uint32_t offset1, uint32_t offset2, RegTensor<float>& nextRow) |
| 1401 | - RegTensor<float>& nextRow) | ||
| 1402 | { | 1395 | { |
| 1403 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 1396 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 1404 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset2))); | 1397 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset2))); |
| 1405 | Add(dst, dst, nextRow, preg); | 1398 | Add(dst, dst, nextRow, preg); |
| 1406 | } | 1399 | } |
| 1407 | 1400 | ||
| @@ -1422,8 +1415,7 @@ __aicore__ inline void GroupNormGradBase<T, U>::CopyInStage2Mode2(TQue<QuePositi | |||
| 1422 | } | 1415 | } |
| 1423 | 1416 | ||
| 1424 | template <typename T, typename U> | 1417 | template <typename T, typename U> |
| 1425 | -__aicore__ inline void GroupNormGradBase<T, U>::FlodSumDgammaVF(__local_mem__ float* src1, __local_mem__ float* src2, | 1418 | +__aicore__ inline void GroupNormGradBase<T, U>::FlodSumDgammaVF(__ubuf__ float* src1, __ubuf__ float* src2, int32_t num) |
| 1426 | - int32_t num) | ||
| 1427 | { | 1419 | { |
| 1428 | uint16_t vlFp32 = GetVRegSize() / sizeof(float); | 1420 | uint16_t vlFp32 = GetVRegSize() / sizeof(float); |
| 1429 | __VEC_SCOPE__ | 1421 | __VEC_SCOPE__ |
| @@ -1438,10 +1430,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::FlodSumDgammaVF(__local_mem__ fl | |||
| 1438 | mask = UpdateMask<float>(width); | 1430 | mask = UpdateMask<float>(width); |
| 1439 | auto src1Addr = src1 + i * vlFp32; | 1431 | auto src1Addr = src1 + i * vlFp32; |
| 1440 | auto src2Addr = src2 + i * vlFp32; | 1432 | auto src2Addr = src2 + i * vlFp32; |
| 1441 | - DataCopy<float, LoadDist::DIST_NORM>(src1Reg, src1Addr); | 1433 | + LoadAlign<float, LoadDist::DIST_NORM>(src1Reg, src1Addr); |
| 1442 | - DataCopy<float, LoadDist::DIST_NORM>(src2Reg, src2Addr); | 1434 | + LoadAlign<float, LoadDist::DIST_NORM>(src2Reg, src2Addr); |
| 1443 | Add(src1Reg, src1Reg, src2Reg, mask); | 1435 | Add(src1Reg, src1Reg, src2Reg, mask); |
| 1444 | - DataCopy<float, StoreDist::DIST_NORM>(src1Addr, src1Reg, mask); | 1436 | + StoreAlign<float, StoreDist::DIST_NORM>(src1Addr, src1Reg, mask); |
| 1445 | } | 1437 | } |
| 1446 | } | 1438 | } |
| 1447 | } | 1439 | } |
| @@ -1480,10 +1472,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFDbetaDgammaBinaryFoldCommon( | |||
| 1480 | 1472 | ||
| 1481 | __VEC_SCOPE__ | 1473 | __VEC_SCOPE__ |
| 1482 | { | 1474 | { |
| 1483 | - UnalignReg uSrcX; | 1475 | + UnalignRegForLoad uSrcX; |
| 1484 | - UnalignReg uSrcXR; | 1476 | + UnalignRegForLoad uSrcXR; |
| 1485 | - UnalignReg uSrcDy; | 1477 | + UnalignRegForLoad uSrcDy; |
| 1486 | - UnalignReg uSrcDyR; | 1478 | + UnalignRegForLoad uSrcDyR; |
| 1487 | RegTensor<float> vregDbeta; | 1479 | RegTensor<float> vregDbeta; |
| 1488 | RegTensor<float> vregDgamma; | 1480 | RegTensor<float> vregDgamma; |
| 1489 | RegTensor<float> vregX; | 1481 | RegTensor<float> vregX; |
| @@ -1505,10 +1497,10 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFDbetaDgammaBinaryFoldCommon( | |||
| 1505 | curUbDyR = ubDyR + ubOffSet; | 1497 | curUbDyR = ubDyR + ubOffSet; |
| 1506 | Duplicate(vregDbeta, 0, pregMain); | 1498 | Duplicate(vregDbeta, 0, pregMain); |
| 1507 | Duplicate(vregDgamma, 0, pregMain); | 1499 | Duplicate(vregDgamma, 0, pregMain); |
| 1508 | - DataCopyUnAlignPre(uSrcX, curUbX); | 1500 | + LoadUnAlignPre(uSrcX, curUbX); |
| 1509 | - DataCopyUnAlignPre(uSrcXR, curUbXR); | 1501 | + LoadUnAlignPre(uSrcXR, curUbXR); |
| 1510 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 1502 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 1511 | - DataCopyUnAlignPre(uSrcDyR, curUbDyR); | 1503 | + LoadUnAlignPre(uSrcDyR, curUbDyR); |
| 1512 | for (uint16_t i = 0; i < remainderGeneral; i++) { | 1504 | for (uint16_t i = 0; i < remainderGeneral; i++) { |
| 1513 | MaskReg pregLoop = UpdateMask<float>(sreg0); | 1505 | MaskReg pregLoop = UpdateMask<float>(sreg0); |
| 1514 | LoadUnAlignOneTensor<T>(curUbX, vregXQ, uSrcX, pregMain, sregvl); | 1506 | LoadUnAlignOneTensor<T>(curUbX, vregXQ, uSrcX, pregMain, sregvl); |
| @@ -1518,17 +1510,17 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFDbetaDgammaBinaryFoldCommon( | |||
| 1518 | Mul(vregXQ, vregXQ, vregDyQ, pregMain); | 1510 | Mul(vregXQ, vregXQ, vregDyQ, pregMain); |
| 1519 | MulAddDst(vregXQ, vregXR, vregDyR, pregLoop); | 1511 | MulAddDst(vregXQ, vregXR, vregDyR, pregLoop); |
| 1520 | Add(vregDyQ, vregDyQ, vregDyR, pregLoop); | 1512 | Add(vregDyQ, vregDyQ, vregDyR, pregLoop); |
| 1521 | - ReduceSum(vregDgamma, vregXQ, pregLoop); | 1513 | + Reduce<ReduceType::SUM>(vregDgamma, vregXQ, pregLoop); |
| 1522 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + i, vregDgamma, pregMerge); | 1514 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + i, vregDgamma, pregMerge); |
| 1523 | - ReduceSum(vregDbeta, vregDyQ, pregLoop); | 1515 | + Reduce<ReduceType::SUM>(vregDbeta, vregDyQ, pregLoop); |
| 1524 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + i, vregDbeta, pregMerge); | 1516 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + i, vregDbeta, pregMerge); |
| 1525 | } | 1517 | } |
| 1526 | { | 1518 | { |
| 1527 | MaskReg pregLoop = UpdateMask<float>(sreg0); | 1519 | MaskReg pregLoop = UpdateMask<float>(sreg0); |
| 1528 | - DataCopyUnAlignPre(uSrcX, curUbX); | 1520 | + LoadUnAlignPre(uSrcX, curUbX); |
| 1529 | - DataCopyUnAlignPre(uSrcXR, curUbXR); | 1521 | + LoadUnAlignPre(uSrcXR, curUbXR); |
| 1530 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 1522 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 1531 | - DataCopyUnAlignPre(uSrcDyR, curUbDyR); | 1523 | + LoadUnAlignPre(uSrcDyR, curUbDyR); |
| 1532 | LoadUnAlignOneTensor<T>(curUbX, vregXQ, uSrcX, pregMain, sregvl); | 1524 | LoadUnAlignOneTensor<T>(curUbX, vregXQ, uSrcX, pregMain, sregvl); |
| 1533 | LoadUnAlignOneTensor<T>(curUbXR, vregXR, uSrcXR, pregLoop, sregvl); | 1525 | LoadUnAlignOneTensor<T>(curUbXR, vregXR, uSrcXR, pregLoop, sregvl); |
| 1534 | LoadUnAlignOneTensor<T>(curUbDy, vregDyQ, uSrcDy, pregMain, sregvl); | 1526 | LoadUnAlignOneTensor<T>(curUbDy, vregDyQ, uSrcDy, pregMain, sregvl); |
| @@ -1536,27 +1528,27 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFDbetaDgammaBinaryFoldCommon( | |||
| 1536 | Mul(vregXQ, vregXQ, vregDyQ, pregMain); | 1528 | Mul(vregXQ, vregXQ, vregDyQ, pregMain); |
| 1537 | MulDstAdd(vregXR, vregDyR, vregXQ, pregLoop); | 1529 | MulDstAdd(vregXR, vregDyR, vregXQ, pregLoop); |
| 1538 | Add(tempDy, vregDyQ, vregDyR, pregLoop); | 1530 | Add(tempDy, vregDyQ, vregDyR, pregLoop); |
| 1539 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregXQ, vregXR, pregLoop); | 1531 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregXQ, vregXR, pregLoop); |
| 1540 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDyQ, tempDy, pregLoop); | 1532 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDyQ, tempDy, pregLoop); |
| 1541 | - ReduceSum(vregDgamma, vregXQ, pregMain); | 1533 | + Reduce<ReduceType::SUM>(vregDgamma, vregXQ, pregMain); |
| 1542 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderGeneral, vregDgamma, | 1534 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderGeneral, vregDgamma, |
| 1543 | - pregMerge); | 1535 | + pregMerge); |
| 1544 | - ReduceSum(vregDbeta, vregDyQ, pregMain); | 1536 | + Reduce<ReduceType::SUM>(vregDbeta, vregDyQ, pregMain); |
| 1545 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderGeneral, vregDbeta, | 1537 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderGeneral, vregDbeta, |
| 1546 | - pregMerge); | 1538 | + pregMerge); |
| 1547 | } | 1539 | } |
| 1548 | - DataCopyUnAlignPre(uSrcX, curUbX); | 1540 | + LoadUnAlignPre(uSrcX, curUbX); |
| 1549 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 1541 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 1550 | for (uint16_t i = 0; i < static_cast<uint16_t>(quotientLoop - remainderLoop); i++) { | 1542 | for (uint16_t i = 0; i < static_cast<uint16_t>(quotientLoop - remainderLoop); i++) { |
| 1551 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, pregMain, sregvl); | 1543 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, pregMain, sregvl); |
| 1552 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, pregMain, sregvl); | 1544 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, pregMain, sregvl); |
| 1553 | Mul(vregX, vregX, vregDy, pregMain); | 1545 | Mul(vregX, vregX, vregDy, pregMain); |
| 1554 | - ReduceSum(vregDgamma, vregX, pregMain); | 1546 | + Reduce<ReduceType::SUM>(vregDgamma, vregX, pregMain); |
| 1555 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderLoop + i, vregDgamma, | 1547 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + remainderLoop + i, vregDgamma, |
| 1556 | - pregMerge); | 1548 | + pregMerge); |
| 1557 | - ReduceSum(vregDbeta, vregDy, pregMain); | 1549 | + Reduce<ReduceType::SUM>(vregDbeta, vregDy, pregMain); |
| 1558 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderLoop + i, vregDbeta, | 1550 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + remainderLoop + i, vregDbeta, |
| 1559 | - pregMerge); | 1551 | + pregMerge); |
| 1560 | } | 1552 | } |
| 1561 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 1553 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 1562 | pregMain = CreateMask<float, MaskPattern::ALL>(); | 1554 | pregMain = CreateMask<float, MaskPattern::ALL>(); |
| @@ -1564,28 +1556,28 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFDbetaDgammaBinaryFoldCommon( | |||
| 1564 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { | 1556 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { |
| 1565 | curBinaryAddLoop = curBinaryAddLoop / 2; | 1557 | curBinaryAddLoop = curBinaryAddLoop / 2; |
| 1566 | for (uint16_t j = 0; j < curBinaryAddLoop; j++) { | 1558 | for (uint16_t j = 0; j < curBinaryAddLoop; j++) { |
| 1567 | - DataCopy(vregXQ, ((__ubuf__ float*)ubBinaryDgamma + j * sregvl)); | 1559 | + LoadAlign(vregXQ, ((__ubuf__ float*)ubBinaryDgamma + j * sregvl)); |
| 1568 | - DataCopy(vregXR, ((__ubuf__ float*)ubBinaryDgamma + (j + curBinaryAddLoop) * sregvl)); | 1560 | + LoadAlign(vregXR, ((__ubuf__ float*)ubBinaryDgamma + (j + curBinaryAddLoop) * sregvl)); |
| 1569 | Add(vregXQ, vregXQ, vregXR, pregMain); | 1561 | Add(vregXQ, vregXQ, vregXR, pregMain); |
| 1570 | - DataCopy(ubBinaryDgamma + j * sregvl, vregXQ, pregMain); | 1562 | + StoreAlign(ubBinaryDgamma + j * sregvl, vregXQ, pregMain); |
| 1571 | - DataCopy(vregDyQ, ((__ubuf__ float*)ubBinaryDbeta + j * sregvl)); | 1563 | + LoadAlign(vregDyQ, ((__ubuf__ float*)ubBinaryDbeta + j * sregvl)); |
| 1572 | - DataCopy(vregDyR, ((__ubuf__ float*)ubBinaryDbeta + (j + curBinaryAddLoop) * sregvl)); | 1564 | + LoadAlign(vregDyR, ((__ubuf__ float*)ubBinaryDbeta + (j + curBinaryAddLoop) * sregvl)); |
| 1573 | Add(vregDyQ, vregDyQ, vregDyR, pregMain); | 1565 | Add(vregDyQ, vregDyQ, vregDyR, pregMain); |
| 1574 | - DataCopy(ubBinaryDbeta + j * sregvl, vregDyQ, pregMain); | 1566 | + StoreAlign(ubBinaryDbeta + j * sregvl, vregDyQ, pregMain); |
| 1575 | } | 1567 | } |
| 1576 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 1568 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 1577 | } | 1569 | } |
| 1578 | { | 1570 | { |
| 1579 | uint32_t sreg2 = binaryAddLastNum; | 1571 | uint32_t sreg2 = binaryAddLastNum; |
| 1580 | MaskReg pregLoop = UpdateMask<float>(sreg2); | 1572 | MaskReg pregLoop = UpdateMask<float>(sreg2); |
| 1581 | - DataCopy(vregDgamma, ((__ubuf__ float*)ubBinaryDgamma)); | 1573 | + LoadAlign(vregDgamma, ((__ubuf__ float*)ubBinaryDgamma)); |
| 1582 | - ReduceSum(vregDgamma, vregDgamma, pregLoop); | 1574 | + Reduce<ReduceType::SUM>(vregDgamma, vregDgamma, pregLoop); |
| 1583 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDgamma + outputOffset + cgIdx, vregDgamma, | 1575 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDgamma + outputOffset + cgIdx, vregDgamma, |
| 1584 | - pregMerge); | 1576 | + pregMerge); |
| 1585 | - DataCopy(vregDbeta, ((__ubuf__ float*)ubBinaryDbeta)); | 1577 | + LoadAlign(vregDbeta, ((__ubuf__ float*)ubBinaryDbeta)); |
| 1586 | - ReduceSum(vregDbeta, vregDbeta, pregLoop); | 1578 | + Reduce<ReduceType::SUM>(vregDbeta, vregDbeta, pregLoop); |
| 1587 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + outputOffset + cgIdx, vregDbeta, | 1579 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + outputOffset + cgIdx, vregDbeta, |
| 1588 | - pregMerge); | 1580 | + pregMerge); |
| 1589 | } | 1581 | } |
| 1590 | } | 1582 | } |
| 1591 | } | 1583 | } |
| @@ -1613,9 +1605,9 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeMode1DxCommon( | |||
| 1613 | 1605 | ||
| 1614 | __VEC_SCOPE__ | 1606 | __VEC_SCOPE__ |
| 1615 | { | 1607 | { |
| 1616 | - UnalignReg uSrcX; | 1608 | + UnalignRegForLoad uSrcX; |
| 1617 | - UnalignReg uSrcDy; | 1609 | + UnalignRegForLoad uSrcDy; |
| 1618 | - UnalignReg uValue; | 1610 | + UnalignRegForStore uValue; |
| 1619 | RegTensor<float> vregX; | 1611 | RegTensor<float> vregX; |
| 1620 | RegTensor<float> vregDy; | 1612 | RegTensor<float> vregDy; |
| 1621 | RegTensor<float> vregGamma; | 1613 | RegTensor<float> vregGamma; |
| @@ -1627,9 +1619,9 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeMode1DxCommon( | |||
| 1627 | curUbDy = ubDy + ubOffSet; | 1619 | curUbDy = ubDy + ubOffSet; |
| 1628 | curUbDst = ubDst + ubOffSet; | 1620 | curUbDst = ubDst + ubOffSet; |
| 1629 | uint32_t dataLen = loopCnt * sregvl; | 1621 | uint32_t dataLen = loopCnt * sregvl; |
| 1630 | - DataCopy<float, LoadDist::DIST_BRC_B32>(vregGamma, ubGamma + gammaOffset + idx); | 1622 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(vregGamma, ubGamma + gammaOffset + idx); |
| 1631 | - DataCopyUnAlignPre(uSrcX, curUbX); | 1623 | + LoadUnAlignPre(uSrcX, curUbX); |
| 1632 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 1624 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 1633 | Muls(vregGamma, vregGamma, rstdScalar, pregAll); | 1625 | Muls(vregGamma, vregGamma, rstdScalar, pregAll); |
| 1634 | for (uint16_t i = 0; i < loopCnt; ++i) { | 1626 | for (uint16_t i = 0; i < loopCnt; ++i) { |
| 1635 | preg = UpdateMask<float>(dataLen); | 1627 | preg = UpdateMask<float>(dataLen); |
| @@ -1643,8 +1635,8 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeMode1DxCommon( | |||
| 1643 | { | 1635 | { |
| 1644 | uint32_t tail = tailNum; | 1636 | uint32_t tail = tailNum; |
| 1645 | preg = UpdateMask<float>(tail); | 1637 | preg = UpdateMask<float>(tail); |
| 1646 | - DataCopyUnAlignPre(uSrcX, curUbX); | 1638 | + LoadUnAlignPre(uSrcX, curUbX); |
| 1647 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 1639 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 1648 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, tailNum); | 1640 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, tailNum); |
| 1649 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, tailNum); | 1641 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, tailNum); |
| 1650 | Muls(vregX, vregX, C2, preg); | 1642 | Muls(vregX, vregX, C2, preg); |
| @@ -1652,7 +1644,7 @@ __aicore__ inline void GroupNormGradBase<T, U>::VFComputeMode1DxCommon( | |||
| 1652 | Adds(vregX, vregX, C3, preg); | 1644 | Adds(vregX, vregX, C3, preg); |
| 1653 | StoreUnAlignOneTensor<T>(curUbDst, vregX, uValue, preg, tailNum); | 1645 | StoreUnAlignOneTensor<T>(curUbDst, vregX, uValue, preg, tailNum); |
| 1654 | } | 1646 | } |
| 1655 | - DataCopyUnAlignPost(curUbDst, uValue, 0); | 1647 | + StoreUnAlignPost(curUbDst, uValue, 0); |
| 1656 | } | 1648 | } |
| 1657 | } | 1649 | } |
| 1658 | } | 1650 | } |
| @@ -22,16 +22,17 @@ | |||
| 22 | namespace GroupNormGrad { | 22 | namespace GroupNormGrad { |
| 23 | using namespace AscendC; | 23 | using namespace AscendC; |
| 24 | using AscendC::MicroAPI::CreateMask; | 24 | using AscendC::MicroAPI::CreateMask; |
| 25 | -using AscendC::MicroAPI::DataCopyUnAlignPost; | ||
| 26 | -using AscendC::MicroAPI::DataCopyUnAlignPre; | ||
| 27 | using AscendC::MicroAPI::LoadDist; | 25 | using AscendC::MicroAPI::LoadDist; |
| 26 | +using AscendC::MicroAPI::LoadUnAlignPre; | ||
| 28 | using AscendC::MicroAPI::LocalMemBar; | 27 | using AscendC::MicroAPI::LocalMemBar; |
| 29 | using AscendC::MicroAPI::MaskPattern; | 28 | using AscendC::MicroAPI::MaskPattern; |
| 30 | using AscendC::MicroAPI::MaskReg; | 29 | using AscendC::MicroAPI::MaskReg; |
| 31 | using AscendC::MicroAPI::MemType; | 30 | using AscendC::MicroAPI::MemType; |
| 32 | using AscendC::MicroAPI::RegTensor; | 31 | using AscendC::MicroAPI::RegTensor; |
| 33 | using AscendC::MicroAPI::StoreDist; | 32 | using AscendC::MicroAPI::StoreDist; |
| 34 | -using AscendC::MicroAPI::UnalignReg; | 33 | +using AscendC::MicroAPI::StoreUnAlignPost; |
| 34 | +using AscendC::MicroAPI::UnalignRegForLoad; | ||
| 35 | +using AscendC::MicroAPI::UnalignRegForStore; | ||
| 35 | using AscendC::MicroAPI::UpdateMask; | 36 | using AscendC::MicroAPI::UpdateMask; |
| 36 | using namespace NormCommon; | 37 | using namespace NormCommon; |
| 37 | using namespace NormCommon::NormCommonRegbase; | 38 | using namespace NormCommon::NormCommonRegbase; |
| @@ -65,87 +66,86 @@ constexpr static AscendC::MicroAPI::CastTrait castTraitB322B16 = { | |||
| 65 | }; | 66 | }; |
| 66 | 67 | ||
| 67 | template <typename T> | 68 | template <typename T> |
| 68 | -__aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T* src1, __local_mem__ T* src2, RegTensor<float>& dst1, | 69 | +__aicore__ inline void LoadTwoTensorForDtypeT(__ubuf__ T* src1, __ubuf__ T* src2, RegTensor<float>& dst1, |
| 69 | RegTensor<float>& dst2, MaskReg& dst1Preg, MaskReg& dst2Preg, | 70 | RegTensor<float>& dst2, MaskReg& dst1Preg, MaskReg& dst2Preg, |
| 70 | uint32_t src1Offset, uint32_t src2Offset) | 71 | uint32_t src1Offset, uint32_t src2Offset) |
| 71 | { | 72 | { |
| 72 | if constexpr (IsSameType<T, half>::value) { | 73 | if constexpr (IsSameType<T, half>::value) { |
| 73 | RegTensor<half> xFp16Q; | 74 | RegTensor<half> xFp16Q; |
| 74 | RegTensor<half> xFp16R; | 75 | RegTensor<half> xFp16R; |
| 75 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ half*)(src1) + (src1Offset))); | 76 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__ubuf__ half*)(src1) + (src1Offset))); |
| 76 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ half*)(src2) + (src2Offset))); | 77 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__ubuf__ half*)(src2) + (src2Offset))); |
| 77 | Cast<float, half, castTraitB162B32>(dst1, xFp16Q, dst1Preg); | 78 | Cast<float, half, castTraitB162B32>(dst1, xFp16Q, dst1Preg); |
| 78 | Cast<float, half, castTraitB162B32>(dst2, xFp16R, dst2Preg); | 79 | Cast<float, half, castTraitB162B32>(dst2, xFp16R, dst2Preg); |
| 79 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 80 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 80 | RegTensor<bfloat16_t> xFp16Q; | 81 | RegTensor<bfloat16_t> xFp16Q; |
| 81 | RegTensor<bfloat16_t> xFp16R; | 82 | RegTensor<bfloat16_t> xFp16R; |
| 82 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ bfloat16_t*)(src1) + (src1Offset))); | 83 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__ubuf__ bfloat16_t*)(src1) + (src1Offset))); |
| 83 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ bfloat16_t*)(src2) + (src2Offset))); | 84 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__ubuf__ bfloat16_t*)(src2) + (src2Offset))); |
| 84 | Cast<float, bfloat16_t, castTraitB162B32>(dst1, xFp16Q, dst1Preg); | 85 | Cast<float, bfloat16_t, castTraitB162B32>(dst1, xFp16Q, dst1Preg); |
| 85 | Cast<float, bfloat16_t, castTraitB162B32>(dst2, xFp16R, dst2Preg); | 86 | Cast<float, bfloat16_t, castTraitB162B32>(dst2, xFp16R, dst2Preg); |
| 86 | } else { | 87 | } else { |
| 87 | - DataCopy(dst1, ((__local_mem__ float*)(src1) + (src1Offset))); | 88 | + LoadAlign(dst1, ((__ubuf__ float*)(src1) + (src1Offset))); |
| 88 | - DataCopy(dst2, ((__local_mem__ float*)(src2) + (src2Offset))); | 89 | + LoadAlign(dst2, ((__ubuf__ float*)(src2) + (src2Offset))); |
| 89 | } | 90 | } |
| 90 | } | 91 | } |
| 91 | 92 | ||
| 92 | template <typename T> | 93 | template <typename T> |
| 93 | -__aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T* input, RegTensor<float>& dst, MaskReg& preg, | 94 | +__aicore__ inline void LoadOneTensorForDtypeT(__ubuf__ T* input, RegTensor<float>& dst, MaskReg& preg, uint32_t offset) |
| 94 | - uint32_t offset) | ||
| 95 | { | 95 | { |
| 96 | if constexpr (IsSameType<T, half>::value) { | 96 | if constexpr (IsSameType<T, half>::value) { |
| 97 | RegTensor<half> xFp16; | 97 | RegTensor<half> xFp16; |
| 98 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ half*)(input) + (offset))); | 98 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ half*)(input) + (offset))); |
| 99 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); | 99 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); |
| 100 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 100 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 101 | RegTensor<bfloat16_t> xBf16; | 101 | RegTensor<bfloat16_t> xBf16; |
| 102 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__local_mem__ bfloat16_t*)(input) + (offset))); | 102 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__ubuf__ bfloat16_t*)(input) + (offset))); |
| 103 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); | 103 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); |
| 104 | } else { | 104 | } else { |
| 105 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset))); | 105 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset))); |
| 106 | } | 106 | } |
| 107 | } | 107 | } |
| 108 | 108 | ||
| 109 | template <typename T> | 109 | template <typename T> |
| 110 | -__aicore__ inline void LoadUnAlignOneTensor(__local_mem__ T*& input, RegTensor<float>& dst, UnalignReg& uSrc, | 110 | +__aicore__ inline void LoadUnAlignOneTensor(__ubuf__ T*& input, RegTensor<float>& dst, UnalignRegForLoad& uSrc, |
| 111 | MaskReg& preg, uint32_t postUpdateStride) | 111 | MaskReg& preg, uint32_t postUpdateStride) |
| 112 | { | 112 | { |
| 113 | if constexpr (IsSameType<T, half>::value) { | 113 | if constexpr (IsSameType<T, half>::value) { |
| 114 | RegTensor<half> xFp16; | 114 | RegTensor<half> xFp16; |
| 115 | RegTensor<half> xFp16UnPack; | 115 | RegTensor<half> xFp16UnPack; |
| 116 | - DataCopyUnAlign(xFp16, uSrc, input, postUpdateStride); | 116 | + LoadUnAlign(xFp16, uSrc, input, postUpdateStride); |
| 117 | UnPack((RegTensor<uint32_t>&)xFp16UnPack, (RegTensor<uint16_t>&)xFp16); | 117 | UnPack((RegTensor<uint32_t>&)xFp16UnPack, (RegTensor<uint16_t>&)xFp16); |
| 118 | Cast<float, half, castTraitB162B32>(dst, xFp16UnPack, preg); | 118 | Cast<float, half, castTraitB162B32>(dst, xFp16UnPack, preg); |
| 119 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 119 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 120 | RegTensor<bfloat16_t> xBf16; | 120 | RegTensor<bfloat16_t> xBf16; |
| 121 | RegTensor<bfloat16_t> xBf16UnPack; | 121 | RegTensor<bfloat16_t> xBf16UnPack; |
| 122 | - DataCopyUnAlign(xBf16, uSrc, input, postUpdateStride); | 122 | + LoadUnAlign(xBf16, uSrc, input, postUpdateStride); |
| 123 | UnPack((RegTensor<uint32_t>&)xBf16UnPack, (RegTensor<uint16_t>&)xBf16); | 123 | UnPack((RegTensor<uint32_t>&)xBf16UnPack, (RegTensor<uint16_t>&)xBf16); |
| 124 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16UnPack, preg); | 124 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16UnPack, preg); |
| 125 | } else { | 125 | } else { |
| 126 | - DataCopyUnAlign(dst, uSrc, input, postUpdateStride); | 126 | + LoadUnAlign(dst, uSrc, input, postUpdateStride); |
| 127 | } | 127 | } |
| 128 | } | 128 | } |
| 129 | 129 | ||
| 130 | template <typename T> | 130 | template <typename T> |
| 131 | -__aicore__ inline void StoreOneTensorForDtypeT(__local_mem__ T* output, RegTensor<float>& src, MaskReg& preg, | 131 | +__aicore__ inline void StoreOneTensorForDtypeT(__ubuf__ T* output, RegTensor<float>& src, MaskReg& preg, |
| 132 | uint32_t offset) | 132 | uint32_t offset) |
| 133 | { | 133 | { |
| 134 | if constexpr (IsSameType<T, half>::value) { | 134 | if constexpr (IsSameType<T, half>::value) { |
| 135 | RegTensor<half> xFp16; | 135 | RegTensor<half> xFp16; |
| 136 | Cast<half, float, castTraitB322B16>(xFp16, src, preg); | 136 | Cast<half, float, castTraitB322B16>(xFp16, src, preg); |
| 137 | - DataCopy<half, StoreDist::DIST_PACK_B32>(((__local_mem__ half*)(output) + offset), xFp16, preg); | 137 | + StoreAlign<half, StoreDist::DIST_PACK_B32>(((__ubuf__ half*)(output) + offset), xFp16, preg); |
| 138 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 138 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 139 | RegTensor<bfloat16_t> xBf16; | 139 | RegTensor<bfloat16_t> xBf16; |
| 140 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg); | 140 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg); |
| 141 | - DataCopy<bfloat16_t, StoreDist::DIST_PACK_B32>(output + offset, xBf16, preg); | 141 | + StoreAlign<bfloat16_t, StoreDist::DIST_PACK_B32>(output + offset, xBf16, preg); |
| 142 | } else { | 142 | } else { |
| 143 | - DataCopy(output + offset, src, preg); | 143 | + StoreAlign(output + offset, src, preg); |
| 144 | } | 144 | } |
| 145 | } | 145 | } |
| 146 | 146 | ||
| 147 | template <typename T> | 147 | template <typename T> |
| 148 | -__aicore__ inline void StoreUnAlignOneTensor(__local_mem__ T*& output, RegTensor<float>& src, UnalignReg& uValue, | 148 | +__aicore__ inline void StoreUnAlignOneTensor(__ubuf__ T*& output, RegTensor<float>& src, UnalignRegForStore& uValue, |
| 149 | MaskReg& preg, uint32_t postUpdateStride) | 149 | MaskReg& preg, uint32_t postUpdateStride) |
| 150 | { | 150 | { |
| 151 | if constexpr (IsSameType<T, half>::value) { | 151 | if constexpr (IsSameType<T, half>::value) { |
| @@ -153,15 +153,15 @@ __aicore__ inline void StoreUnAlignOneTensor(__local_mem__ T*& output, RegTensor | |||
| 153 | RegTensor<half> xFp16Pack; | 153 | RegTensor<half> xFp16Pack; |
| 154 | Cast<half, float, castTraitB322B16>(xFp16, src, preg); | 154 | Cast<half, float, castTraitB322B16>(xFp16, src, preg); |
| 155 | Pack((RegTensor<uint16_t>&)xFp16Pack, (RegTensor<uint32_t>&)xFp16); | 155 | Pack((RegTensor<uint16_t>&)xFp16Pack, (RegTensor<uint32_t>&)xFp16); |
| 156 | - DataCopyUnAlign(output, xFp16Pack, uValue, postUpdateStride); | 156 | + StoreUnAlign(output, xFp16Pack, uValue, postUpdateStride); |
| 157 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 157 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 158 | RegTensor<bfloat16_t> xBf16; | 158 | RegTensor<bfloat16_t> xBf16; |
| 159 | RegTensor<bfloat16_t> xBf16Pack; | 159 | RegTensor<bfloat16_t> xBf16Pack; |
| 160 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg); | 160 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, src, preg); |
| 161 | Pack((RegTensor<uint16_t>&)xBf16Pack, (RegTensor<uint32_t>&)xBf16); | 161 | Pack((RegTensor<uint16_t>&)xBf16Pack, (RegTensor<uint32_t>&)xBf16); |
| 162 | - DataCopyUnAlign(output, xBf16Pack, uValue, postUpdateStride); | 162 | + StoreUnAlign(output, xBf16Pack, uValue, postUpdateStride); |
| 163 | } else { | 163 | } else { |
| 164 | - DataCopyUnAlign(output, src, uValue, postUpdateStride); | 164 | + StoreUnAlign(output, src, uValue, postUpdateStride); |
| 165 | } | 165 | } |
| 166 | } | 166 | } |
| 167 | 167 | ||
| @@ -183,15 +183,15 @@ __aicore__ inline void VFCastFloat2T(const __ubuf__ T* ubAddrOut, const __ubuf__ | |||
| 183 | if constexpr (IsSameType<T, half>::value) { | 183 | if constexpr (IsSameType<T, half>::value) { |
| 184 | RegTensor<half> vregB16; | 184 | RegTensor<half> vregB16; |
| 185 | RegTensor<float> vregF32; | 185 | RegTensor<float> vregF32; |
| 186 | - DataCopy(vregF32, srcAddr + i * sregvl); | 186 | + LoadAlign(vregF32, srcAddr + i * sregvl); |
| 187 | Cast<half, float, castTraitB322B16>(vregB16, vregF32, preg); | 187 | Cast<half, float, castTraitB322B16>(vregB16, vregF32, preg); |
| 188 | - DataCopy<half, StoreDist::DIST_PACK_B32>(dstAddr + i * sregvl, vregB16, preg); | 188 | + StoreAlign<half, StoreDist::DIST_PACK_B32>(dstAddr + i * sregvl, vregB16, preg); |
| 189 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 189 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 190 | RegTensor<bfloat16_t> vregBF16; | 190 | RegTensor<bfloat16_t> vregBF16; |
| 191 | RegTensor<float> vregF32; | 191 | RegTensor<float> vregF32; |
| 192 | - DataCopy(vregF32, srcAddr + i * sregvl); | 192 | + LoadAlign(vregF32, srcAddr + i * sregvl); |
| 193 | Cast<bfloat16_t, float, castTraitB322B16>(vregBF16, vregF32, preg); | 193 | Cast<bfloat16_t, float, castTraitB322B16>(vregBF16, vregF32, preg); |
| 194 | - DataCopy<bfloat16_t, StoreDist::DIST_PACK_B32>(dstAddr + i * sregvl, vregBF16, preg); | 194 | + StoreAlign<bfloat16_t, StoreDist::DIST_PACK_B32>(dstAddr + i * sregvl, vregBF16, preg); |
| 195 | } | 195 | } |
| 196 | } | 196 | } |
| 197 | } | 197 | } |
| @@ -215,15 +215,15 @@ __aicore__ inline void VFCastT2Float(const __ubuf__ float* ubAddrOut, const __ub | |||
| 215 | if constexpr (IsSameType<T, half>::value) { | 215 | if constexpr (IsSameType<T, half>::value) { |
| 216 | RegTensor<half> vregB16; | 216 | RegTensor<half> vregB16; |
| 217 | RegTensor<float> vregF32; | 217 | RegTensor<float> vregF32; |
| 218 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(vregB16, srcAddr + i * sregvl); | 218 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(vregB16, srcAddr + i * sregvl); |
| 219 | Cast<float, half, castTraitB162B32>(vregF32, vregB16, preg); | 219 | Cast<float, half, castTraitB162B32>(vregF32, vregB16, preg); |
| 220 | - DataCopy(dstAddr + i * sregvl, vregF32, preg); | 220 | + StoreAlign(dstAddr + i * sregvl, vregF32, preg); |
| 221 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 221 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 222 | RegTensor<bfloat16_t> vregBF16; | 222 | RegTensor<bfloat16_t> vregBF16; |
| 223 | RegTensor<float> vregF32; | 223 | RegTensor<float> vregF32; |
| 224 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(vregBF16, srcAddr + i * sregvl); | 224 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(vregBF16, srcAddr + i * sregvl); |
| 225 | Cast<float, bfloat16_t, castTraitB162B32>(vregF32, vregBF16, preg); | 225 | Cast<float, bfloat16_t, castTraitB162B32>(vregF32, vregBF16, preg); |
| 226 | - DataCopy(dstAddr + i * sregvl, vregF32, preg); | 226 | + StoreAlign(dstAddr + i * sregvl, vregF32, preg); |
| 227 | } | 227 | } |
| 228 | } | 228 | } |
| 229 | } | 229 | } |
| @@ -234,10 +234,10 @@ __aicore__ inline void VFCastT2Float(const __ubuf__ float* ubAddrOut, const __ub | |||
| 234 | dgamma = reduceSum(dy * x) | 234 | dgamma = reduceSum(dy * x) |
| 235 | */ | 235 | */ |
| 236 | template <typename T> | 236 | template <typename T> |
| 237 | -__aicore__ inline void VFComputeDbetaDs( | 237 | +__aicore__ inline void VFComputeDbetaDs(const LocalTensor<T>& x, const LocalTensor<T>& dy, |
| 238 | - const LocalTensor<T>& x, const LocalTensor<T>& dy, const LocalTensor<float>& dbeta, | 238 | + const LocalTensor<float>& dbeta, const LocalTensor<float>& dgamma, |
| 239 | - const LocalTensor<float>& dgamma, uint32_t eleNumPerC, uint32_t vecLen, uint32_t storeBaseOffset, | 239 | + uint32_t eleNumPerC, uint32_t vecLen, uint32_t storeBaseOffset, |
| 240 | - uint16_t loopCount) | 240 | + uint16_t loopCount) |
| 241 | { | 241 | { |
| 242 | __ubuf__ T* ubX = (__ubuf__ T*)x.GetPhyAddr(); | 242 | __ubuf__ T* ubX = (__ubuf__ T*)x.GetPhyAddr(); |
| 243 | __ubuf__ T* ubDy = (__ubuf__ T*)dy.GetPhyAddr(); | 243 | __ubuf__ T* ubDy = (__ubuf__ T*)dy.GetPhyAddr(); |
| @@ -249,8 +249,8 @@ __aicore__ inline void VFComputeDbetaDs( | |||
| 249 | 249 | ||
| 250 | __VEC_SCOPE__ | 250 | __VEC_SCOPE__ |
| 251 | { | 251 | { |
| 252 | - UnalignReg uSrcX; | 252 | + UnalignRegForLoad uSrcX; |
| 253 | - UnalignReg uSrcDy; | 253 | + UnalignRegForLoad uSrcDy; |
| 254 | RegTensor<float> vregDbeta; | 254 | RegTensor<float> vregDbeta; |
| 255 | RegTensor<float> vregDgamma; | 255 | RegTensor<float> vregDgamma; |
| 256 | RegTensor<float> tempDbeta; | 256 | RegTensor<float> tempDbeta; |
| @@ -266,27 +266,28 @@ __aicore__ inline void VFComputeDbetaDs( | |||
| 266 | curUbDy = ubDy + ubOffSet; | 266 | curUbDy = ubDy + ubOffSet; |
| 267 | Duplicate(vregDbeta, 0, pregAll); | 267 | Duplicate(vregDbeta, 0, pregAll); |
| 268 | Duplicate(vregDgamma, 0, pregAll); | 268 | Duplicate(vregDgamma, 0, pregAll); |
| 269 | - DataCopyUnAlignPre(uSrcX, curUbX); | 269 | + LoadUnAlignPre(uSrcX, curUbX); |
| 270 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 270 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 271 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { | 271 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { |
| 272 | preg = UpdateMask<float>(sreg); | 272 | preg = UpdateMask<float>(sreg); |
| 273 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, sregvl); | 273 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, sregvl); |
| 274 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, sregvl); | 274 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, sregvl); |
| 275 | MulDstAdd(vregX, vregDy, vregDgamma, preg); | 275 | MulDstAdd(vregX, vregDy, vregDgamma, preg); |
| 276 | Add(tempDbeta, vregDbeta, vregDy, preg); | 276 | Add(tempDbeta, vregDbeta, vregDy, preg); |
| 277 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDbeta, tempDbeta, preg); | 277 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDbeta, tempDbeta, preg); |
| 278 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDgamma, vregX, preg); | 278 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDgamma, vregX, preg); |
| 279 | } | 279 | } |
| 280 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); | 280 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); |
| 281 | - ReduceSum(vregDbeta, vregDbeta, pregAll); | 281 | + Reduce<ReduceType::SUM>(vregDbeta, vregDbeta, pregAll); |
| 282 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + storeBaseOffset + idx, vregDbeta, pregMerge); | 282 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + storeBaseOffset + idx, vregDbeta, pregMerge); |
| 283 | - ReduceSum(vregDgamma, vregDgamma, pregAll); | 283 | + Reduce<ReduceType::SUM>(vregDgamma, vregDgamma, pregAll); |
| 284 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDgamma + storeBaseOffset + idx, vregDgamma, pregMerge); | 284 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDgamma + storeBaseOffset + idx, vregDgamma, |
| 285 | + pregMerge); | ||
| 285 | } | 286 | } |
| 286 | } | 287 | } |
| 287 | } | 288 | } |
| 288 | 289 | ||
| 289 | -template<typename U> | 290 | +template <typename U> |
| 290 | __aicore__ inline void UpdateCacheStage2Mode2(const LocalTensor<U>& dstTensor, const LocalTensor<U>& srcTensor, | 291 | __aicore__ inline void UpdateCacheStage2Mode2(const LocalTensor<U>& dstTensor, const LocalTensor<U>& srcTensor, |
| 291 | const int64_t cacheId, const int64_t stride, const int64_t count) | 292 | const int64_t cacheId, const int64_t stride, const int64_t count) |
| 292 | { | 293 | { |
| @@ -294,9 +295,9 @@ __aicore__ inline void UpdateCacheStage2Mode2(const LocalTensor<U>& dstTensor, c | |||
| 294 | uint16_t innerLoopTimes = cacheId; | 295 | uint16_t innerLoopTimes = cacheId; |
| 295 | uint32_t outerLoopStride = GetVRegSize() / sizeof(U); | 296 | uint32_t outerLoopStride = GetVRegSize() / sizeof(U); |
| 296 | uint32_t innerLoopStride = stride; | 297 | uint32_t innerLoopStride = stride; |
| 297 | - __local_mem__ U* dst = (__local_mem__ U*)dstTensor.GetPhyAddr(); | 298 | + __ubuf__ U* dst = (__ubuf__ U*)dstTensor.GetPhyAddr(); |
| 298 | - __local_mem__ U* cache = (__local_mem__ U*)dstTensor.GetPhyAddr() + cacheId * stride; | 299 | + __ubuf__ U* cache = (__ubuf__ U*)dstTensor.GetPhyAddr() + cacheId * stride; |
| 299 | - __local_mem__ U* src = (__local_mem__ U*)srcTensor.GetPhyAddr(); | 300 | + __ubuf__ U* src = (__ubuf__ U*)srcTensor.GetPhyAddr(); |
| 300 | __VEC_SCOPE__ | 301 | __VEC_SCOPE__ |
| 301 | { | 302 | { |
| 302 | uint32_t sreg = static_cast<uint32_t>(count); | 303 | uint32_t sreg = static_cast<uint32_t>(count); |
| @@ -304,16 +305,16 @@ __aicore__ inline void UpdateCacheStage2Mode2(const LocalTensor<U>& dstTensor, c | |||
| 304 | MaskReg pMask; | 305 | MaskReg pMask; |
| 305 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 306 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 306 | pMask = UpdateMask<U>(sreg); | 307 | pMask = UpdateMask<U>(sreg); |
| 307 | - DataCopy(aReg, (__local_mem__ U*)src + i * outerLoopStride); | 308 | + LoadAlign(aReg, (__ubuf__ U*)src + i * outerLoopStride); |
| 308 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 309 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 309 | - DataCopy(bReg, (__local_mem__ U*)dst + i * outerLoopStride + j * innerLoopStride); | 310 | + LoadAlign(bReg, (__ubuf__ U*)dst + i * outerLoopStride + j * innerLoopStride); |
| 310 | Add<U, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 311 | Add<U, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 311 | } | 312 | } |
| 312 | - DataCopy((__local_mem__ U*)cache + i * outerLoopStride, aReg, pMask); | 313 | + StoreAlign((__ubuf__ U*)cache + i * outerLoopStride, aReg, pMask); |
| 313 | } | 314 | } |
| 314 | } | 315 | } |
| 315 | } | 316 | } |
| 316 | 317 | ||
| 317 | __aicore__ inline int64_t GetCacheId(const int64_t idx) { return ScalarGetCountOfValue<1>(idx ^ (idx + 1)) - 1; } | 318 | __aicore__ inline int64_t GetCacheId(const int64_t idx) { return ScalarGetCountOfValue<1>(idx ^ (idx + 1)) - 1; } |
| 318 | } // namespace GroupNormGrad | 319 | } // namespace GroupNormGrad |
| 319 | -#endif | 320 | +#endif |
| @@ -33,21 +33,20 @@ public: | |||
| 33 | 33 | ||
| 34 | private: | 34 | private: |
| 35 | __aicore__ inline void InitBuffer(const GroupNormGradRegBaseTilingData* tilingData); | 35 | __aicore__ inline void InitBuffer(const GroupNormGradRegBaseTilingData* tilingData); |
| 36 | - __aicore__ inline void Compute( | 36 | + __aicore__ inline void Compute(int32_t taskIdx, const LocalTensor<T>& xTensor, const LocalTensor<T>& dyTensor, |
| 37 | - int32_t taskIdx, const LocalTensor<T>& xTensor, const LocalTensor<T>& dyTensor, const LocalTensor<T>& dxTensor, | 37 | + const LocalTensor<T>& dxTensor, const float mean, const float rstd); |
| 38 | - const float mean, const float rstd); | 38 | + __aicore__ inline void VFMode0DbetaDsOneHw(const LocalTensor<T>& x, const LocalTensor<T>& dy, |
| 39 | - __aicore__ inline void VFMode0DbetaDsOneHw( | 39 | + const LocalTensor<float>& dbeta, const LocalTensor<float>& dgamma); |
| 40 | - const LocalTensor<T>& x, const LocalTensor<T>& dy, const LocalTensor<float>& dbeta, | 40 | + __aicore__ inline void ComputeMode0Dx(int32_t taskIdx, const LocalTensor<T>& xTensor, |
| 41 | - const LocalTensor<float>& dgamma); | 41 | + const LocalTensor<T>& dyTensor, const LocalTensor<T>& dxTensor, |
| 42 | - __aicore__ inline void ComputeMode0Dx( | 42 | + LocalTensor<float>& dbetaTensor, LocalTensor<float>& dsTensor, |
| 43 | - int32_t taskIdx, const LocalTensor<T>& xTensor, const LocalTensor<T>& dyTensor, const LocalTensor<T>& dxTensor, | 43 | + const float mean, const float rstd); |
| 44 | - LocalTensor<float>& dbetaTensor, LocalTensor<float>& dsTensor, const float mean, const float rstd); | 44 | + __aicore__ inline void VFComputeMode0DxOneHw(const LocalTensor<T>& dstTensor, const LocalTensor<T>& xTensor, |
| 45 | - __aicore__ inline void VFComputeMode0DxOneHw( | 45 | + const LocalTensor<T>& dyTensor, const LocalTensor<float>& gammaTensor, |
| 46 | - const LocalTensor<T>& dstTensor, const LocalTensor<T>& xTensor, const LocalTensor<T>& dyTensor, | 46 | + const float C2, const float C3, const float rstd); |
| 47 | - const LocalTensor<float>& gammaTensor, const float C2, const float C3, const float rstd); | 47 | + __aicore__ inline void VFComputeMode0Dx(const LocalTensor<T>& dstTensor, const LocalTensor<T>& xTensor, |
| 48 | - __aicore__ inline void VFComputeMode0Dx( | 48 | + const LocalTensor<T>& dyTensor, const LocalTensor<float>& gammaTensor, |
| 49 | - const LocalTensor<T>& dstTensor, const LocalTensor<T>& xTensor, const LocalTensor<T>& dyTensor, | 49 | + const float C2, const float C3, const float rstd); |
| 50 | - const LocalTensor<float>& gammaTensor, const float C2, const float C3, const float rstd); | ||
| 51 | __aicore__ inline void Stage1Process(); | 50 | __aicore__ inline void Stage1Process(); |
| 52 | }; | 51 | }; |
| 53 | 52 | ||
| @@ -166,7 +165,7 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::Compute(int32_t taskIdx, co | |||
| 166 | VFMode0DbetaDsOneHw(xTensor, dyTensor, dbetaTensor, dsTensor); | 165 | VFMode0DbetaDsOneHw(xTensor, dyTensor, dbetaTensor, dsTensor); |
| 167 | } else if (this->eleNumPerC_ <= this->VecLen_) { | 166 | } else if (this->eleNumPerC_ <= this->VecLen_) { |
| 168 | VFComputeDbetaDs<T>(xTensor, dyTensor, dbetaTensor, dsTensor, this->eleNumPerC_, this->VecLen_, 0, | 167 | VFComputeDbetaDs<T>(xTensor, dyTensor, dbetaTensor, dsTensor, this->eleNumPerC_, this->VecLen_, 0, |
| 169 | - static_cast<uint16_t>(this->C_G_)); | 168 | + static_cast<uint16_t>(this->C_G_)); |
| 170 | } else { | 169 | } else { |
| 171 | this->VFDbetaDgammaBinaryFoldCommon(xTensor, dyTensor, dbetaTensor, dsTensor, 0, this->C_G_); | 170 | this->VFDbetaDgammaBinaryFoldCommon(xTensor, dyTensor, dbetaTensor, dsTensor, 0, this->C_G_); |
| 172 | } | 171 | } |
| @@ -203,14 +202,14 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFMode0DbetaDsOneHw(const L | |||
| 203 | 202 | ||
| 204 | __VEC_SCOPE__ | 203 | __VEC_SCOPE__ |
| 205 | { | 204 | { |
| 206 | - UnalignReg uSrcX; | 205 | + UnalignRegForLoad uSrcX; |
| 207 | - UnalignReg uSrcDy; | 206 | + UnalignRegForLoad uSrcDy; |
| 208 | - UnalignReg uDbeta; | 207 | + UnalignRegForStore uDbeta; |
| 209 | - UnalignReg uDgamma; | 208 | + UnalignRegForStore uDgamma; |
| 210 | RegTensor<float> vregX; | 209 | RegTensor<float> vregX; |
| 211 | RegTensor<float> vregDy; | 210 | RegTensor<float> vregDy; |
| 212 | - DataCopyUnAlignPre(uSrcX, ubX); | 211 | + LoadUnAlignPre(uSrcX, ubX); |
| 213 | - DataCopyUnAlignPre(uSrcDy, ubDy); | 212 | + LoadUnAlignPre(uSrcDy, ubDy); |
| 214 | uint32_t sreg = (uint32_t)this->C_G_; | 213 | uint32_t sreg = (uint32_t)this->C_G_; |
| 215 | for (uint16_t i = 0; i < loopCnt; i++) { | 214 | for (uint16_t i = 0; i < loopCnt; i++) { |
| 216 | MaskReg preg = UpdateMask<float>(sreg); | 215 | MaskReg preg = UpdateMask<float>(sreg); |
| @@ -223,16 +222,16 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFMode0DbetaDsOneHw(const L | |||
| 223 | { | 222 | { |
| 224 | uint32_t tail = tailNum; | 223 | uint32_t tail = tailNum; |
| 225 | MaskReg preg = UpdateMask<float>(tail); | 224 | MaskReg preg = UpdateMask<float>(tail); |
| 226 | - DataCopyUnAlignPre(uSrcX, ubX); | 225 | + LoadUnAlignPre(uSrcX, ubX); |
| 227 | - DataCopyUnAlignPre(uSrcDy, ubDy); | 226 | + LoadUnAlignPre(uSrcDy, ubDy); |
| 228 | LoadUnAlignOneTensor<T>(ubX, vregX, uSrcX, preg, tailNum); | 227 | LoadUnAlignOneTensor<T>(ubX, vregX, uSrcX, preg, tailNum); |
| 229 | LoadUnAlignOneTensor<T>(ubDy, vregDy, uSrcDy, preg, tailNum); | 228 | LoadUnAlignOneTensor<T>(ubDy, vregDy, uSrcDy, preg, tailNum); |
| 230 | Mul(vregX, vregX, vregDy, preg); | 229 | Mul(vregX, vregX, vregDy, preg); |
| 231 | StoreUnAlignOneTensor(ubDbeta, vregDy, uDbeta, preg, tailNum); | 230 | StoreUnAlignOneTensor(ubDbeta, vregDy, uDbeta, preg, tailNum); |
| 232 | StoreUnAlignOneTensor(ubDgamma, vregX, uDgamma, preg, tailNum); | 231 | StoreUnAlignOneTensor(ubDgamma, vregX, uDgamma, preg, tailNum); |
| 233 | } | 232 | } |
| 234 | - DataCopyUnAlignPost(ubDbeta, uDbeta, 0); | 233 | + StoreUnAlignPost(ubDbeta, uDbeta, 0); |
| 235 | - DataCopyUnAlignPost(ubDgamma, uDgamma, 0); | 234 | + StoreUnAlignPost(ubDgamma, uDgamma, 0); |
| 236 | } | 235 | } |
| 237 | } | 236 | } |
| 238 | 237 | ||
| @@ -284,14 +283,14 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFComputeMode0DxOneHw( | |||
| 284 | 283 | ||
| 285 | __VEC_SCOPE__ | 284 | __VEC_SCOPE__ |
| 286 | { | 285 | { |
| 287 | - UnalignReg uSrcX; | 286 | + UnalignRegForLoad uSrcX; |
| 288 | - UnalignReg uSrcDy; | 287 | + UnalignRegForLoad uSrcDy; |
| 289 | - UnalignReg uValue; | 288 | + UnalignRegForStore uValue; |
| 290 | RegTensor<float> vregX; | 289 | RegTensor<float> vregX; |
| 291 | RegTensor<float> vregDy; | 290 | RegTensor<float> vregDy; |
| 292 | RegTensor<float> vregGamma; | 291 | RegTensor<float> vregGamma; |
| 293 | - DataCopyUnAlignPre(uSrcX, ubX); | 292 | + LoadUnAlignPre(uSrcX, ubX); |
| 294 | - DataCopyUnAlignPre(uSrcDy, ubDy); | 293 | + LoadUnAlignPre(uSrcDy, ubDy); |
| 295 | for (uint16_t i = 0; i < loopCnt; ++i) { | 294 | for (uint16_t i = 0; i < loopCnt; ++i) { |
| 296 | uint32_t dataLen = (uint32_t)loopCnt * sregvl; | 295 | uint32_t dataLen = (uint32_t)loopCnt * sregvl; |
| 297 | MaskReg preg = UpdateMask<float>(dataLen); | 296 | MaskReg preg = UpdateMask<float>(dataLen); |
| @@ -307,8 +306,8 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFComputeMode0DxOneHw( | |||
| 307 | { | 306 | { |
| 308 | uint32_t tail = tailNum; | 307 | uint32_t tail = tailNum; |
| 309 | MaskReg preg = UpdateMask<float>(tail); | 308 | MaskReg preg = UpdateMask<float>(tail); |
| 310 | - DataCopyUnAlignPre(uSrcX, ubX); | 309 | + LoadUnAlignPre(uSrcX, ubX); |
| 311 | - DataCopyUnAlignPre(uSrcDy, ubDy); | 310 | + LoadUnAlignPre(uSrcDy, ubDy); |
| 312 | LoadUnAlignOneTensor<T>(ubX, vregX, uSrcX, preg, (uint32_t)tailNum); | 311 | LoadUnAlignOneTensor<T>(ubX, vregX, uSrcX, preg, (uint32_t)tailNum); |
| 313 | LoadUnAlignOneTensor<T>(ubDy, vregDy, uSrcDy, preg, (uint32_t)tailNum); | 312 | LoadUnAlignOneTensor<T>(ubDy, vregDy, uSrcDy, preg, (uint32_t)tailNum); |
| 314 | LoadOneTensorForDtypeT<float>(ubGamma, vregGamma, preg, loopCnt * sregvl); | 313 | LoadOneTensorForDtypeT<float>(ubGamma, vregGamma, preg, loopCnt * sregvl); |
| @@ -318,7 +317,7 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFComputeMode0DxOneHw( | |||
| 318 | Adds(vregX, vregX, C3, preg); | 317 | Adds(vregX, vregX, C3, preg); |
| 319 | StoreUnAlignOneTensor<T>(ubDst, vregX, uValue, preg, tailNum); | 318 | StoreUnAlignOneTensor<T>(ubDst, vregX, uValue, preg, tailNum); |
| 320 | } | 319 | } |
| 321 | - DataCopyUnAlignPost(ubDst, uValue, 0); | 320 | + StoreUnAlignPost(ubDst, uValue, 0); |
| 322 | } | 321 | } |
| 323 | } | 322 | } |
| 324 | 323 | ||
| @@ -350,9 +349,9 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFComputeMode0Dx(const Loca | |||
| 350 | 349 | ||
| 351 | __VEC_SCOPE__ | 350 | __VEC_SCOPE__ |
| 352 | { | 351 | { |
| 353 | - UnalignReg uSrcX; | 352 | + UnalignRegForLoad uSrcX; |
| 354 | - UnalignReg uSrcDy; | 353 | + UnalignRegForLoad uSrcDy; |
| 355 | - UnalignReg uValue; | 354 | + UnalignRegForStore uValue; |
| 356 | RegTensor<float> vregX; | 355 | RegTensor<float> vregX; |
| 357 | RegTensor<float> vregDy; | 356 | RegTensor<float> vregDy; |
| 358 | RegTensor<float> vregGamma; | 357 | RegTensor<float> vregGamma; |
| @@ -364,9 +363,9 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFComputeMode0Dx(const Loca | |||
| 364 | curUbDy = ubDy + ubOffSet; | 363 | curUbDy = ubDy + ubOffSet; |
| 365 | curUbDst = ubDst + ubOffSet; | 364 | curUbDst = ubDst + ubOffSet; |
| 366 | uint32_t dataLen = (uint32_t)loopCnt * sregvl; | 365 | uint32_t dataLen = (uint32_t)loopCnt * sregvl; |
| 367 | - DataCopy<float, LoadDist::DIST_BRC_B32>(vregGamma, ubGamma + cgIdx); | 366 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(vregGamma, ubGamma + cgIdx); |
| 368 | - DataCopyUnAlignPre(uSrcX, curUbX); | 367 | + LoadUnAlignPre(uSrcX, curUbX); |
| 369 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 368 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 370 | Muls(vregGamma, vregGamma, rstdScalar, pregAll); | 369 | Muls(vregGamma, vregGamma, rstdScalar, pregAll); |
| 371 | for (uint16_t i = 0; i < (uint16_t)loopCnt; ++i) { | 370 | for (uint16_t i = 0; i < (uint16_t)loopCnt; ++i) { |
| 372 | preg = UpdateMask<float>(dataLen); | 371 | preg = UpdateMask<float>(dataLen); |
| @@ -380,8 +379,8 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFComputeMode0Dx(const Loca | |||
| 380 | { | 379 | { |
| 381 | uint32_t tail = tailNum; | 380 | uint32_t tail = tailNum; |
| 382 | preg = UpdateMask<float>(tail); | 381 | preg = UpdateMask<float>(tail); |
| 383 | - DataCopyUnAlignPre(uSrcX, curUbX); | 382 | + LoadUnAlignPre(uSrcX, curUbX); |
| 384 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 383 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 385 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, tailNum); | 384 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, tailNum); |
| 386 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, tailNum); | 385 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, tailNum); |
| 387 | Muls(vregX, vregX, C2, preg); | 386 | Muls(vregX, vregX, C2, preg); |
| @@ -389,7 +388,7 @@ __aicore__ inline void GroupNormGradGFullLoad<T, U>::VFComputeMode0Dx(const Loca | |||
| 389 | Adds(vregX, vregX, C3, preg); | 388 | Adds(vregX, vregX, C3, preg); |
| 390 | StoreUnAlignOneTensor<T>(curUbDst, vregX, uValue, preg, tailNum); | 389 | StoreUnAlignOneTensor<T>(curUbDst, vregX, uValue, preg, tailNum); |
| 391 | } | 390 | } |
| 392 | - DataCopyUnAlignPost(curUbDst, uValue, 0); | 391 | + StoreUnAlignPost(curUbDst, uValue, 0); |
| 393 | } | 392 | } |
| 394 | } | 393 | } |
| 395 | } | 394 | } |
| @@ -209,10 +209,10 @@ __aicore__ inline void GroupNormGradReCompute<T, U>::VFMode2DbetaDs( | |||
| 209 | Mul(vregXF, vregXF, vregDyF, pregLoop); | 209 | Mul(vregXF, vregXF, vregDyF, pregLoop); |
| 210 | MulDstAdd(vregXM, vregDyM, vregXF, pregMain); | 210 | MulDstAdd(vregXM, vregDyM, vregXF, pregMain); |
| 211 | Add(vregDyM, vregDyM, vregDyF, pregMain); | 211 | Add(vregDyM, vregDyM, vregDyF, pregMain); |
| 212 | - ReduceSum(vregDgamma, vregXM, pregLoop); | 212 | + Reduce<ReduceType::SUM>(vregDgamma, vregXM, pregLoop); |
| 213 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + i, vregDgamma, pregMerge); | 213 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + i, vregDgamma, pregMerge); |
| 214 | - ReduceSum(vregDbeta, vregDyM, pregLoop); | 214 | + Reduce<ReduceType::SUM>(vregDbeta, vregDyM, pregLoop); |
| 215 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + i, vregDbeta, pregMerge); | 215 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + i, vregDbeta, pregMerge); |
| 216 | } | 216 | } |
| 217 | // step2: reduce the fold tail (last 64 or less than 64) blocks reduce to 1. | 217 | // step2: reduce the fold tail (last 64 or less than 64) blocks reduce to 1. |
| 218 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainerLoopTimes); i++) { | 218 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainerLoopTimes); i++) { |
| @@ -224,13 +224,14 @@ __aicore__ inline void GroupNormGradReCompute<T, U>::VFMode2DbetaDs( | |||
| 224 | Mul(vregXM, vregXM, vregDyM, pregMain); | 224 | Mul(vregXM, vregXM, vregDyM, pregMain); |
| 225 | MulDstAdd(vregXF, vregDyF, vregXM, pregLoop); | 225 | MulDstAdd(vregXF, vregDyF, vregXM, pregLoop); |
| 226 | Add(tempDy, vregDyM, vregDyF, pregLoop); | 226 | Add(tempDy, vregDyM, vregDyF, pregLoop); |
| 227 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregXM, vregXF, pregLoop); | 227 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregXM, vregXF, pregLoop); |
| 228 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDyM, tempDy, pregLoop); | 228 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDyM, tempDy, pregLoop); |
| 229 | - ReduceSum(vregDgamma, vregXM, pregMain); | 229 | + Reduce<ReduceType::SUM>(vregDgamma, vregXM, pregMain); |
| 230 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + overLapLoopTimes, vregDgamma, | 230 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDgamma + overLapLoopTimes, vregDgamma, |
| 231 | - pregMerge); | 231 | + pregMerge); |
| 232 | - ReduceSum(vregDbeta, vregDyM, pregMain); | 232 | + Reduce<ReduceType::SUM>(vregDbeta, vregDyM, pregMain); |
| 233 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + overLapLoopTimes, vregDbeta, pregMerge); | 233 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + overLapLoopTimes, vregDbeta, |
| 234 | + pregMerge); | ||
| 234 | } | 235 | } |
| 235 | // step3: non-overlapping portions of the first half reduce by 64, this part always 64 align | 236 | // step3: non-overlapping portions of the first half reduce by 64, this part always 64 align |
| 236 | uint32_t sreg2 = unFoldAddNum; | 237 | uint32_t sreg2 = unFoldAddNum; |
| @@ -239,12 +240,12 @@ __aicore__ inline void GroupNormGradReCompute<T, U>::VFMode2DbetaDs( | |||
| 239 | LoadOneTensorForDtypeT<T>(ubXMain, vregXM, pregLoop, (i + overLapLoopTimes + remainerLoopTimes) * sregvl); | 240 | LoadOneTensorForDtypeT<T>(ubXMain, vregXM, pregLoop, (i + overLapLoopTimes + remainerLoopTimes) * sregvl); |
| 240 | LoadOneTensorForDtypeT<T>(ubDyMain, vregDyM, pregLoop, (i + overLapLoopTimes + remainerLoopTimes) * sregvl); | 241 | LoadOneTensorForDtypeT<T>(ubDyMain, vregDyM, pregLoop, (i + overLapLoopTimes + remainerLoopTimes) * sregvl); |
| 241 | Mul(vregXM, vregXM, vregDyM, pregLoop); | 242 | Mul(vregXM, vregXM, vregDyM, pregLoop); |
| 242 | - ReduceSum(vregDgamma, vregXM, pregLoop); | 243 | + Reduce<ReduceType::SUM>(vregDgamma, vregXM, pregLoop); |
| 243 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 244 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 244 | ubBinaryDgamma + i + overLapLoopTimes + remainerLoopTimes, vregDgamma, pregMerge); | 245 | ubBinaryDgamma + i + overLapLoopTimes + remainerLoopTimes, vregDgamma, pregMerge); |
| 245 | - ReduceSum(vregDbeta, vregDyM, pregLoop); | 246 | + Reduce<ReduceType::SUM>(vregDbeta, vregDyM, pregLoop); |
| 246 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubBinaryDbeta + i + overLapLoopTimes + remainerLoopTimes, | 247 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 247 | - vregDbeta, pregMerge); | 248 | + ubBinaryDbeta + i + overLapLoopTimes + remainerLoopTimes, vregDbeta, pregMerge); |
| 248 | } | 249 | } |
| 249 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 250 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 250 | // step4: binary folding reduce calculation | 251 | // step4: binary folding reduce calculation |
| @@ -253,15 +254,15 @@ __aicore__ inline void GroupNormGradReCompute<T, U>::VFMode2DbetaDs( | |||
| 253 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { | 254 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { |
| 254 | curBinaryAddLoop = curBinaryAddLoop / 2; | 255 | curBinaryAddLoop = curBinaryAddLoop / 2; |
| 255 | for (uint16_t j = 0; j < curBinaryAddLoop; j++) { | 256 | for (uint16_t j = 0; j < curBinaryAddLoop; j++) { |
| 256 | - DataCopy(vregXM, ((__ubuf__ float*)ubBinaryDgamma + j * sregvl)); | 257 | + LoadAlign(vregXM, ((__ubuf__ float*)ubBinaryDgamma + j * sregvl)); |
| 257 | - DataCopy(vregXF, ((__ubuf__ float*)ubBinaryDgamma + (j + curBinaryAddLoop) * sregvl)); | 258 | + LoadAlign(vregXF, ((__ubuf__ float*)ubBinaryDgamma + (j + curBinaryAddLoop) * sregvl)); |
| 258 | Add(vregXM, vregXM, vregXF, pregMain); | 259 | Add(vregXM, vregXM, vregXF, pregMain); |
| 259 | - DataCopy(((__ubuf__ float*)ubBinaryDgamma + j * sregvl), vregXM, pregMain); | 260 | + StoreAlign(((__ubuf__ float*)ubBinaryDgamma + j * sregvl), vregXM, pregMain); |
| 260 | 261 | ||
| 261 | - DataCopy(vregDyM, ((__ubuf__ float*)ubBinaryDbeta + j * sregvl)); | 262 | + LoadAlign(vregDyM, ((__ubuf__ float*)ubBinaryDbeta + j * sregvl)); |
| 262 | - DataCopy(vregDyF, ((__ubuf__ float*)ubBinaryDbeta + (j + curBinaryAddLoop) * sregvl)); | 263 | + LoadAlign(vregDyF, ((__ubuf__ float*)ubBinaryDbeta + (j + curBinaryAddLoop) * sregvl)); |
| 263 | Add(vregDyM, vregDyM, vregDyF, pregMain); | 264 | Add(vregDyM, vregDyM, vregDyF, pregMain); |
| 264 | - DataCopy(((__ubuf__ float*)ubBinaryDbeta + j * sregvl), vregDyM, pregMain); | 265 | + StoreAlign(((__ubuf__ float*)ubBinaryDbeta + j * sregvl), vregDyM, pregMain); |
| 265 | } | 266 | } |
| 266 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 267 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 267 | } | 268 | } |
| @@ -270,21 +271,21 @@ __aicore__ inline void GroupNormGradReCompute<T, U>::VFMode2DbetaDs( | |||
| 270 | uint32_t sreg3 = binaryUbLastNum; | 271 | uint32_t sreg3 = binaryUbLastNum; |
| 271 | uint32_t pos = loopIdx & 0xFF; | 272 | uint32_t pos = loopIdx & 0xFF; |
| 272 | MaskReg pregLoop = UpdateMask<float>(sreg3); | 273 | MaskReg pregLoop = UpdateMask<float>(sreg3); |
| 273 | - DataCopy(vregDgamma, ((__ubuf__ float*)ubBinaryDgamma)); | 274 | + LoadAlign(vregDgamma, ((__ubuf__ float*)ubBinaryDgamma)); |
| 274 | - ReduceSum(vregDgamma, vregDgamma, pregLoop); | 275 | + Reduce<ReduceType::SUM>(vregDgamma, vregDgamma, pregLoop); |
| 275 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDgamma + pos, vregDgamma, pregMerge); | 276 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDgamma + pos, vregDgamma, pregMerge); |
| 276 | - DataCopy(vregDbeta, ((__ubuf__ float*)ubBinaryDbeta)); | 277 | + LoadAlign(vregDbeta, ((__ubuf__ float*)ubBinaryDbeta)); |
| 277 | - ReduceSum(vregDbeta, vregDbeta, pregLoop); | 278 | + Reduce<ReduceType::SUM>(vregDbeta, vregDbeta, pregLoop); |
| 278 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + pos, vregDbeta, pregMerge); | 279 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + pos, vregDbeta, pregMerge); |
| 279 | } | 280 | } |
| 280 | } // end VF | 281 | } // end VF |
| 281 | this->outQueDx_.FreeTensor(binaryDbetaTensor); | 282 | this->outQueDx_.FreeTensor(binaryDbetaTensor); |
| 282 | } | 283 | } |
| 283 | 284 | ||
| 284 | /* | 285 | /* |
| 285 | - dbeta = ReduceSum(dy) | 286 | + dbeta = Reduce<ReduceType::SUM>(dy) |
| 286 | temp = xHat * rstd - mean * rstd | 287 | temp = xHat * rstd - mean * rstd |
| 287 | - dgamma = ReduceSum(dy * temp) | 288 | + dgamma = Reduce<ReduceType::SUM>(dy * temp) |
| 288 | */ | 289 | */ |
| 289 | template <typename T, typename U> | 290 | template <typename T, typename U> |
| 290 | __aicore__ inline void GroupNormGradReCompute<T, U>::Mode2DbetaDs(const LocalTensor<float>& dbeta, | 291 | __aicore__ inline void GroupNormGradReCompute<T, U>::Mode2DbetaDs(const LocalTensor<float>& dbeta, |
| @@ -467,4 +468,4 @@ __aicore__ inline void GroupNormGradReCompute<T, U>::VFComputeMode2Dx(const Loca | |||
| 467 | } | 468 | } |
| 468 | } | 469 | } |
| 469 | } // namespace GroupNormGrad | 470 | } // namespace GroupNormGrad |
| 470 | -#endif | 471 | +#endif |
| @@ -301,8 +301,8 @@ __aicore__ inline void GroupNormGradSmallNGCFullLoad<T, U>::VFMode0DbetaDs( | |||
| 301 | 301 | ||
| 302 | __VEC_SCOPE__ | 302 | __VEC_SCOPE__ |
| 303 | { | 303 | { |
| 304 | - UnalignReg uSrcX; | 304 | + UnalignRegForLoad uSrcX; |
| 305 | - UnalignReg uSrcDy; | 305 | + UnalignRegForLoad uSrcDy; |
| 306 | RegTensor<float> vregDbeta; | 306 | RegTensor<float> vregDbeta; |
| 307 | RegTensor<float> vregDs; | 307 | RegTensor<float> vregDs; |
| 308 | RegTensor<float> tempDbeta; | 308 | RegTensor<float> tempDbeta; |
| @@ -319,22 +319,22 @@ __aicore__ inline void GroupNormGradSmallNGCFullLoad<T, U>::VFMode0DbetaDs( | |||
| 319 | curUbDy = ubDy + ubOffSet; | 319 | curUbDy = ubDy + ubOffSet; |
| 320 | Duplicate(vregDbeta, 0, pregAll); | 320 | Duplicate(vregDbeta, 0, pregAll); |
| 321 | Duplicate(vregDs, 0, pregAll); | 321 | Duplicate(vregDs, 0, pregAll); |
| 322 | - DataCopyUnAlignPre(uSrcX, curUbX); | 322 | + LoadUnAlignPre(uSrcX, curUbX); |
| 323 | - DataCopyUnAlignPre(uSrcDy, curUbDy); | 323 | + LoadUnAlignPre(uSrcDy, curUbDy); |
| 324 | 324 | ||
| 325 | preg = UpdateMask<float>(sreg); | 325 | preg = UpdateMask<float>(sreg); |
| 326 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, sregvl); | 326 | LoadUnAlignOneTensor<T>(curUbX, vregX, uSrcX, preg, sregvl); |
| 327 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, sregvl); | 327 | LoadUnAlignOneTensor<T>(curUbDy, vregDy, uSrcDy, preg, sregvl); |
| 328 | MulDstAdd(vregX, vregDy, vregDs, preg); | 328 | MulDstAdd(vregX, vregDy, vregDs, preg); |
| 329 | Add(tempDbeta, vregDbeta, vregDy, preg); | 329 | Add(tempDbeta, vregDbeta, vregDy, preg); |
| 330 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDbeta, tempDbeta, preg); | 330 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDbeta, tempDbeta, preg); |
| 331 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDs, vregX, preg); | 331 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(vregDs, vregX, preg); |
| 332 | 332 | ||
| 333 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); | 333 | MaskReg pregMerge = CreateMask<float, MaskPattern::VL1>(); |
| 334 | - ReduceSum(vregDbeta, vregDbeta, pregAll); | 334 | + Reduce<ReduceType::SUM>(vregDbeta, vregDbeta, pregAll); |
| 335 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + cIdxOffSet + idx, vregDbeta, pregMerge); | 335 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDbeta + cIdxOffSet + idx, vregDbeta, pregMerge); |
| 336 | - ReduceSum(vregDs, vregDs, pregAll); | 336 | + Reduce<ReduceType::SUM>(vregDs, vregDs, pregAll); |
| 337 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDs + cIdxOffSet + idx, vregDs, pregMerge); | 337 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(ubDs + cIdxOffSet + idx, vregDs, pregMerge); |
| 338 | } | 338 | } |
| 339 | } | 339 | } |
| 340 | } | 340 | } |
| @@ -366,12 +366,12 @@ __aicore__ inline void GroupNormGradSmallNGCFullLoad<T, U>::VFComputeStage1Ds(co | |||
| 366 | uint32_t sregvl = (uint32_t)this->VecLen_; | 366 | uint32_t sregvl = (uint32_t)this->VecLen_; |
| 367 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { | 367 | for (uint16_t i = 0; i < (uint16_t)repeatTimes; ++i) { |
| 368 | preg = UpdateMask<float>(sreg); | 368 | preg = UpdateMask<float>(sreg); |
| 369 | - DataCopy(vregDs, ubDs + i * sregvl); | 369 | + LoadAlign(vregDs, ubDs + i * sregvl); |
| 370 | - DataCopy(vregDb, ubDb + i * sregvl); | 370 | + LoadAlign(vregDb, ubDb + i * sregvl); |
| 371 | Muls(vregDb, vregDb, meanScalar, preg); | 371 | Muls(vregDb, vregDb, meanScalar, preg); |
| 372 | Sub(vregDs, vregDs, vregDb, preg); | 372 | Sub(vregDs, vregDs, vregDb, preg); |
| 373 | Muls(vregDs, vregDs, rstdScalar, preg); | 373 | Muls(vregDs, vregDs, rstdScalar, preg); |
| 374 | - DataCopy(ubDst + i * sregvl, vregDs, preg); | 374 | + StoreAlign(ubDst + i * sregvl, vregDs, preg); |
| 375 | } | 375 | } |
| 376 | } | 376 | } |
| 377 | } | 377 | } |
| @@ -397,4 +397,4 @@ __aicore__ inline void GroupNormGradSmallNGCFullLoad<T, U>::ComputeStage2() | |||
| 397 | } | 397 | } |
| 398 | } | 398 | } |
| 399 | } // namespace GroupNormGrad | 399 | } // namespace GroupNormGrad |
| 400 | -#endif | 400 | +#endif |
| @@ -21,7 +21,8 @@ using namespace AscendC; | |||
| 21 | using namespace AscendC::MicroAPI; | 21 | using namespace AscendC::MicroAPI; |
| 22 | using AscendC::MicroAPI::MaskReg; | 22 | using AscendC::MicroAPI::MaskReg; |
| 23 | using AscendC::MicroAPI::RegTensor; | 23 | using AscendC::MicroAPI::RegTensor; |
| 24 | -using AscendC::MicroAPI::UnalignReg; | 24 | +using AscendC::MicroAPI::UnalignRegForLoad; |
| 25 | +using AscendC::MicroAPI::UnalignRegForStore; | ||
| 25 | static constexpr int32_t BLOCK_SIZE = 32; | 26 | static constexpr int32_t BLOCK_SIZE = 32; |
| 26 | static constexpr int32_t FOUR_BUF = 4; | 27 | static constexpr int32_t FOUR_BUF = 4; |
| 27 | static constexpr int32_t FP32_ONE_REPEAT = 64; | 28 | static constexpr int32_t FP32_ONE_REPEAT = 64; |
| @@ -94,52 +95,49 @@ __aicore__ inline uint32_t RoundDown(uint32_t x) | |||
| 94 | } | 95 | } |
| 95 | 96 | ||
| 96 | template <typename T> | 97 | template <typename T> |
| 97 | -__aicore__ inline void LoadInputData(RegTensor<float>& dst, __local_mem__ T* src, MaskReg pregLoop, uint32_t srcOffset) | 98 | +__aicore__ inline void LoadInputData(RegTensor<float>& dst, __ubuf__ T* src, MaskReg pregLoop, uint32_t srcOffset) |
| 98 | { | 99 | { |
| 99 | if constexpr (IsSameType<T, float>::value) { | 100 | if constexpr (IsSameType<T, float>::value) { |
| 100 | - DataCopy(dst, src + srcOffset); | 101 | + LoadAlign(dst, src + srcOffset); |
| 101 | } else { | 102 | } else { |
| 102 | RegTensor<T> tmp; | 103 | RegTensor<T> tmp; |
| 103 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(tmp, src + srcOffset); | 104 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>(tmp, src + srcOffset); |
| 104 | Cast<float, T, castTraitB162B32Even>(dst, tmp, pregLoop); | 105 | Cast<float, T, castTraitB162B32Even>(dst, tmp, pregLoop); |
| 105 | } | 106 | } |
| 106 | } | 107 | } |
| 107 | 108 | ||
| 108 | template <typename T> | 109 | template <typename T> |
| 109 | -__aicore__ inline void LoadGammaAndBetaData(RegTensor<float>& gamma, RegTensor<float>& beta, | 110 | +__aicore__ inline void LoadGammaAndBetaData(RegTensor<float>& gamma, RegTensor<float>& beta, __ubuf__ T* gammaLocal, |
| 110 | - __local_mem__ T* gammaLocal, __local_mem__ T* betaLocal, MaskReg pregLoop, | 111 | + __ubuf__ T* betaLocal, MaskReg pregLoop, uint32_t srcOffset) |
| 111 | - uint32_t srcOffset) | ||
| 112 | { | 112 | { |
| 113 | if constexpr (IsSameType<T, float>::value) { | 113 | if constexpr (IsSameType<T, float>::value) { |
| 114 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(gamma, gammaLocal + srcOffset); | 114 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(gamma, gammaLocal + srcOffset); |
| 115 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(beta, betaLocal + srcOffset); | 115 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(beta, betaLocal + srcOffset); |
| 116 | } else { | 116 | } else { |
| 117 | RegTensor<T> gammaB16; | 117 | RegTensor<T> gammaB16; |
| 118 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_BRC_B16>(gammaB16, gammaLocal + srcOffset); | 118 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_BRC_B16>(gammaB16, gammaLocal + srcOffset); |
| 119 | Cast<float, T, castTraitB162B32Even>(gamma, gammaB16, pregLoop); | 119 | Cast<float, T, castTraitB162B32Even>(gamma, gammaB16, pregLoop); |
| 120 | RegTensor<T> betaB16; | 120 | RegTensor<T> betaB16; |
| 121 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_BRC_B16>(betaB16, betaLocal + srcOffset); | 121 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_BRC_B16>(betaB16, betaLocal + srcOffset); |
| 122 | Cast<float, T, castTraitB162B32Even>(beta, betaB16, pregLoop); | 122 | Cast<float, T, castTraitB162B32Even>(beta, betaB16, pregLoop); |
| 123 | } | 123 | } |
| 124 | } | 124 | } |
| 125 | 125 | ||
| 126 | template <typename T> | 126 | template <typename T> |
| 127 | -__aicore__ inline void StoreOutputData(__local_mem__ T* dst, RegTensor<float>& src, MaskReg pregLoop, | 127 | +__aicore__ inline void StoreOutputData(__ubuf__ T* dst, RegTensor<float>& src, MaskReg pregLoop, uint32_t dstOffset) |
| 128 | - uint32_t dstOffset) | ||
| 129 | { | 128 | { |
| 130 | if constexpr (IsSameType<T, float>::value) { | 129 | if constexpr (IsSameType<T, float>::value) { |
| 131 | - DataCopy(dst + dstOffset, src, pregLoop); | 130 | + StoreAlign(dst + dstOffset, src, pregLoop); |
| 132 | } else { | 131 | } else { |
| 133 | RegTensor<T> tmpB16; | 132 | RegTensor<T> tmpB16; |
| 134 | Cast<T, float, castTraitB322B16Even>(tmpB16, src, pregLoop); | 133 | Cast<T, float, castTraitB322B16Even>(tmpB16, src, pregLoop); |
| 135 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + dstOffset, tmpB16, pregLoop); | 134 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + dstOffset, tmpB16, pregLoop); |
| 136 | } | 135 | } |
| 137 | } | 136 | } |
| 138 | 137 | ||
| 139 | template <typename T> | 138 | template <typename T> |
| 140 | -__aicore__ inline void VFInnerWelfordParallelUpdateWithInit(__local_mem__ T* x1Local, __local_mem__ float* tmpMeanLocal, | 139 | +__aicore__ inline void VFInnerWelfordParallelUpdateWithInit(__ubuf__ T* x1Local, __ubuf__ float* tmpMeanLocal, |
| 141 | - __local_mem__ float* tmpVarLocal, uint64_t calLen, | 140 | + __ubuf__ float* tmpVarLocal, uint64_t calLen, float scale) |
| 142 | - float scale) | ||
| 143 | { | 141 | { |
| 144 | uint16_t loopCount = CeilDiv(calLen, VL_FP32); | 142 | uint16_t loopCount = CeilDiv(calLen, VL_FP32); |
| 145 | __VEC_SCOPE__ | 143 | __VEC_SCOPE__ |
| @@ -160,13 +158,13 @@ __aicore__ inline void VFInnerWelfordParallelUpdateWithInit(__local_mem__ T* x1L | |||
| 160 | Sub(delta1, x1, tmpMean, pregLoop); | 158 | Sub(delta1, x1, tmpMean, pregLoop); |
| 161 | Muls(delta2, delta1, scale, pregLoop); | 159 | Muls(delta2, delta1, scale, pregLoop); |
| 162 | Add(tmpMean, tmpMean, delta2, pregLoop); | 160 | Add(tmpMean, tmpMean, delta2, pregLoop); |
| 163 | - DataCopy(tmpMeanLocal + i * VL_FP32, tmpMean, pregLoop); | 161 | + StoreAlign(tmpMeanLocal + i * VL_FP32, tmpMean, pregLoop); |
| 164 | 162 | ||
| 165 | Duplicate(tmpVar, 0.0, pregLoop); | 163 | Duplicate(tmpVar, 0.0, pregLoop); |
| 166 | Sub(delta3, x1, tmpMean, pregLoop); | 164 | Sub(delta3, x1, tmpMean, pregLoop); |
| 167 | Mul(delat4, delta1, delta3, pregLoop); | 165 | Mul(delat4, delta1, delta3, pregLoop); |
| 168 | Add(tmpVar, tmpVar, delat4, pregLoop); | 166 | Add(tmpVar, tmpVar, delat4, pregLoop); |
| 169 | - DataCopy(tmpVarLocal + i * VL_FP32, tmpVar, pregLoop); | 167 | + StoreAlign(tmpVarLocal + i * VL_FP32, tmpVar, pregLoop); |
| 170 | } | 168 | } |
| 171 | } | 169 | } |
| 172 | } | 170 | } |
| @@ -181,8 +179,8 @@ __aicore__ inline void VFInnerWelfordParallelUpdateWithInit(__local_mem__ T* x1L | |||
| 181 | return count, mean, var | 179 | return count, mean, var |
| 182 | */ | 180 | */ |
| 183 | template <typename T> | 181 | template <typename T> |
| 184 | -__aicore__ inline void VFInnerWelfordParallelUpdate(__local_mem__ T* x1Local, __local_mem__ float* tmpMeanLocal, | 182 | +__aicore__ inline void VFInnerWelfordParallelUpdate(__ubuf__ T* x1Local, __ubuf__ float* tmpMeanLocal, |
| 185 | - __local_mem__ float* tmpVarLocal, uint64_t calLen, float scale) | 183 | + __ubuf__ float* tmpVarLocal, uint64_t calLen, float scale) |
| 186 | { | 184 | { |
| 187 | uint16_t loopCount = CeilDiv(calLen, VL_FP32); | 185 | uint16_t loopCount = CeilDiv(calLen, VL_FP32); |
| 188 | __VEC_SCOPE__ | 186 | __VEC_SCOPE__ |
| @@ -199,24 +197,24 @@ __aicore__ inline void VFInnerWelfordParallelUpdate(__local_mem__ T* x1Local, __ | |||
| 199 | for (uint16_t i = 0; i < loopCount; i++) { | 197 | for (uint16_t i = 0; i < loopCount; i++) { |
| 200 | pregLoop = UpdateMask<float>(sreg0); | 198 | pregLoop = UpdateMask<float>(sreg0); |
| 201 | LoadInputData<T>(x1, x1Local, pregLoop, i * VL_FP32); | 199 | LoadInputData<T>(x1, x1Local, pregLoop, i * VL_FP32); |
| 202 | - DataCopy(tmpMean, tmpMeanLocal + i * VL_FP32); | 200 | + LoadAlign(tmpMean, tmpMeanLocal + i * VL_FP32); |
| 203 | Sub(delta1, x1, tmpMean, pregLoop); | 201 | Sub(delta1, x1, tmpMean, pregLoop); |
| 204 | Muls(delta2, delta1, scale, pregLoop); | 202 | Muls(delta2, delta1, scale, pregLoop); |
| 205 | Add(tmpMean, tmpMean, delta2, pregLoop); | 203 | Add(tmpMean, tmpMean, delta2, pregLoop); |
| 206 | - DataCopy(tmpMeanLocal + i * VL_FP32, tmpMean, pregLoop); | 204 | + StoreAlign(tmpMeanLocal + i * VL_FP32, tmpMean, pregLoop); |
| 207 | 205 | ||
| 208 | - DataCopy(tmpVar, tmpVarLocal + i * VL_FP32); | 206 | + LoadAlign(tmpVar, tmpVarLocal + i * VL_FP32); |
| 209 | Sub(delta3, x1, tmpMean, pregLoop); | 207 | Sub(delta3, x1, tmpMean, pregLoop); |
| 210 | Mul(delat4, delta1, delta3, pregLoop); | 208 | Mul(delat4, delta1, delta3, pregLoop); |
| 211 | Add(tmpVar, tmpVar, delat4, pregLoop); | 209 | Add(tmpVar, tmpVar, delat4, pregLoop); |
| 212 | - DataCopy(tmpVarLocal + i * VL_FP32, tmpVar, pregLoop); | 210 | + StoreAlign(tmpVarLocal + i * VL_FP32, tmpVar, pregLoop); |
| 213 | } | 211 | } |
| 214 | } | 212 | } |
| 215 | } | 213 | } |
| 216 | 214 | ||
| 217 | template <typename T> | 215 | template <typename T> |
| 218 | -__aicore__ inline void VFWelfordParallelUpdate(__local_mem__ T* x1Local, __local_mem__ float* tmpMeanLocal, | 216 | +__aicore__ inline void VFWelfordParallelUpdate(__ubuf__ T* x1Local, __ubuf__ float* tmpMeanLocal, |
| 219 | - __local_mem__ float* tmpVarLocal, uint64_t curLoop, uint64_t calLen, | 217 | + __ubuf__ float* tmpVarLocal, uint64_t curLoop, uint64_t calLen, |
| 220 | float scale) | 218 | float scale) |
| 221 | { | 219 | { |
| 222 | if (curLoop == 0) { | 220 | if (curLoop == 0) { |
| @@ -237,10 +235,9 @@ __aicore__ inline void VFWelfordParallelUpdate(__local_mem__ T* x1Local, __local | |||
| 237 | welford采用二分累加计算mean和variance, 基本逻辑为: | 235 | welford采用二分累加计算mean和variance, 基本逻辑为: |
| 238 | 先将尾块折叠到整块上，整尾块vadd之后，做一次vcadd回刷到UB上，剩余整块直接vcadd回刷到UB上，最后对UB上的结果做完全二分对折 | 236 | 先将尾块折叠到整块上，整尾块vadd之后，做一次vcadd回刷到UB上，剩余整块直接vcadd回刷到UB上，最后对UB上的结果做完全二分对折 |
| 239 | */ | 237 | */ |
| 240 | -__aicore__ inline void VFWelfordParallelFinalizeAlign(__local_mem__ float* meanLocal, __local_mem__ float* rstdLocal, | 238 | +__aicore__ inline void VFWelfordParallelFinalizeAlign(__ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, |
| 241 | - __local_mem__ float* tmpMeanLocal, | 239 | + __ubuf__ float* tmpMeanLocal, __ubuf__ float* tmpVarLocal, |
| 242 | - __local_mem__ float* tmpVarLocal, | 240 | + __ubuf__ float* dichotomyAddLocal, uint32_t reduceCount, |
| 243 | - __local_mem__ float* dichotomyAddLocal, uint32_t reduceCount, | ||
| 244 | uint32_t dichotomyAddPower, uint32_t dichotomyAddK, | 241 | uint32_t dichotomyAddPower, uint32_t dichotomyAddK, |
| 245 | uint32_t dichotomyAddLastNum, uint32_t offset, float reduceScale, | 242 | uint32_t dichotomyAddLastNum, uint32_t offset, float reduceScale, |
| 246 | float scale, float cnt, float eps) | 243 | float scale, float cnt, float eps) |
| @@ -272,28 +269,28 @@ __aicore__ inline void VFWelfordParallelFinalizeAlign(__local_mem__ float* meanL | |||
| 272 | // PART1: 整尾块合并 | 269 | // PART1: 整尾块合并 |
| 273 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 270 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 274 | pregLoop = UpdateMask<float>(sreg0); | 271 | pregLoop = UpdateMask<float>(sreg0); |
| 275 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 272 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 276 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 273 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 277 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, scale, pregMain); | 274 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, scale, pregMain); |
| 278 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, scale, pregLoop); | 275 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, scale, pregLoop); |
| 279 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 276 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 280 | - ReduceSum(mean, sumMean, pregMain); | 277 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 281 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, | 278 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, |
| 282 | - pregMerge); | 279 | + pregMerge); |
| 283 | } | 280 | } |
| 284 | 281 | ||
| 285 | // PART2: 整块剩余部分vcadd回刷UB | 282 | // PART2: 整块剩余部分vcadd回刷UB |
| 286 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderLoopCount); | 283 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderLoopCount); |
| 287 | i++) { | 284 | i++) { |
| 288 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderLoopCount) * VL_FP32); | 285 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderLoopCount) * VL_FP32); |
| 289 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, scale, pregMain); | 286 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, scale, pregMain); |
| 290 | - ReduceSum(mean, dichotomyAddMeanL, pregMain); | 287 | + Reduce<ReduceType::SUM>(mean, dichotomyAddMeanL, pregMain); |
| 291 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 288 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 292 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, mean, pregMerge); | 289 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, mean, pregMerge); |
| 293 | } | 290 | } |
| 294 | 291 | ||
| 295 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 292 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 296 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); | 293 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); |
| 297 | 294 | ||
| 298 | Duplicate(one, float(1.0), pregMain); | 295 | Duplicate(one, float(1.0), pregMain); |
| 299 | Duplicate(mean, mean, pregMain); | 296 | Duplicate(mean, mean, pregMain); |
| @@ -301,46 +298,46 @@ __aicore__ inline void VFWelfordParallelFinalizeAlign(__local_mem__ float* meanL | |||
| 301 | // PART1: 整尾块合并 | 298 | // PART1: 整尾块合并 |
| 302 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 299 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 303 | pregLoop = UpdateMask<float>(sreg0); | 300 | pregLoop = UpdateMask<float>(sreg0); |
| 304 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 301 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 305 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 302 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 306 | Mul(deltaL, deltaL, deltaL, pregMain); | 303 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 307 | Muls(deltaL, deltaL, cnt, pregMain); | 304 | Muls(deltaL, deltaL, cnt, pregMain); |
| 308 | - DataCopy(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); | 305 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); |
| 309 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 306 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 310 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 307 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 311 | 308 | ||
| 312 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 309 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 313 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); | 310 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); |
| 314 | Mul(deltaR, deltaR, deltaR, pregLoop); | 311 | Mul(deltaR, deltaR, deltaR, pregLoop); |
| 315 | Muls(deltaR, deltaR, cnt, pregLoop); | 312 | Muls(deltaR, deltaR, cnt, pregLoop); |
| 316 | - DataCopy(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); | 313 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); |
| 317 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); | 314 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); |
| 318 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); | 315 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); |
| 319 | 316 | ||
| 320 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 317 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 321 | - ReduceSum(var, sumVar, pregMain); | 318 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 322 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, | 319 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, |
| 323 | - pregMerge); | 320 | + pregMerge); |
| 324 | } | 321 | } |
| 325 | 322 | ||
| 326 | // PART2: 整块剩余部分vcadd回刷UB | 323 | // PART2: 整块剩余部分vcadd回刷UB |
| 327 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderLoopCount); | 324 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderLoopCount); |
| 328 | i++) { | 325 | i++) { |
| 329 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderLoopCount) * VL_FP32); | 326 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderLoopCount) * VL_FP32); |
| 330 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 327 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 331 | Mul(deltaL, deltaL, deltaL, pregMain); | 328 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 332 | Muls(deltaL, deltaL, cnt, pregMain); | 329 | Muls(deltaL, deltaL, cnt, pregMain); |
| 333 | - DataCopy(dichotomyAddVarL, tmpVarLocal + (i + dichotomyAddReminderLoopCount) * VL_FP32); | 330 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + (i + dichotomyAddReminderLoopCount) * VL_FP32); |
| 334 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 331 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 335 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 332 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 336 | - ReduceSum(var, dichotomyAddVarL, pregMain); | 333 | + Reduce<ReduceType::SUM>(var, dichotomyAddVarL, pregMain); |
| 337 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 334 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 338 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, var, pregMerge); | 335 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, var, pregMerge); |
| 339 | } | 336 | } |
| 340 | 337 | ||
| 341 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 338 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 342 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); | 339 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); |
| 343 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); | 340 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); |
| 344 | } | 341 | } |
| 345 | } | 342 | } |
| 346 | 343 | ||
| @@ -368,10 +365,9 @@ __aicore__ inline void VFWelfordParallelFinalizeAlign(__local_mem__ float* meanL | |||
| 368 | 365 | ||
| 369 | // welford整块大于等于二分累加整块 | 366 | // welford整块大于等于二分累加整块 |
| 370 | __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation1( | 367 | __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation1( |
| 371 | - __local_mem__ float* meanLocal, __local_mem__ float* rstdLocal, __local_mem__ float* tmpMeanLocal, | 368 | + __ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, __ubuf__ float* tmpMeanLocal, __ubuf__ float* tmpVarLocal, |
| 372 | - __local_mem__ float* tmpVarLocal, __local_mem__ float* dichotomyAddLocal, uint32_t reduceCount, | 369 | + __ubuf__ float* dichotomyAddLocal, uint32_t reduceCount, uint32_t dichotomyAddPower, uint32_t dichotomyAddK, |
| 373 | - uint32_t dichotomyAddPower, uint32_t dichotomyAddK, uint32_t dichotomyAddLastNum, uint32_t offset, | 370 | + uint32_t dichotomyAddLastNum, uint32_t offset, uint32_t tailSize, float reduceScale, float cnt, float eps) |
| 374 | - uint32_t tailSize, float reduceScale, float cnt, float eps) | ||
| 375 | { | 371 | { |
| 376 | float tailCnt = cnt + float(1.0); | 372 | float tailCnt = cnt + float(1.0); |
| 377 | float coeff = tailCnt / cnt; | 373 | float coeff = tailCnt / cnt; |
| @@ -417,14 +413,14 @@ __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation1( | |||
| 417 | 413 | ||
| 418 | // 整块使用tailCountScale,尾块使用tailCountScale | 414 | // 整块使用tailCountScale,尾块使用tailCountScale |
| 419 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { | 415 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { |
| 420 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 416 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 421 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 417 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 422 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); | 418 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); |
| 423 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, tailCountScale, pregMain); | 419 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, tailCountScale, pregMain); |
| 424 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 420 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 425 | - ReduceSum(mean, sumMean, pregMain); | 421 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 426 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, | 422 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, |
| 427 | - pregMerge); | 423 | + pregMerge); |
| 428 | } | 424 | } |
| 429 | 425 | ||
| 430 | // 处理welford第一次非对齐点, 整块使用tailCountScale,尾块部分使用tailCountScale, 部分使用countScale | 426 | // 处理welford第一次非对齐点, 整块使用tailCountScale,尾块部分使用tailCountScale, 部分使用countScale |
| @@ -433,146 +429,147 @@ __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation1( | |||
| 433 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { | 429 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { |
| 434 | pregLoop = UpdateMask<float>(sreg0); | 430 | pregLoop = UpdateMask<float>(sreg0); |
| 435 | pregLoop1 = UpdateMask<float>(sreg1); | 431 | pregLoop1 = UpdateMask<float>(sreg1); |
| 436 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); | 432 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); |
| 437 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); | 433 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); |
| 438 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); | 434 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); |
| 439 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); | 435 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); |
| 440 | Muls(tmp, dichotomyAddMeanR, coeff, pregLoop1); | 436 | Muls(tmp, dichotomyAddMeanR, coeff, pregLoop1); |
| 441 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(dichotomyAddMeanR, tmp, pregLoop1); | 437 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(dichotomyAddMeanR, tmp, pregLoop1); |
| 442 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 438 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 443 | - ReduceSum(mean, sumMean, pregMain); | 439 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 444 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 440 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 445 | dichotomyAddLocal + i + welfordDiffLoopCount, mean, pregMerge); | 441 | dichotomyAddLocal + i + welfordDiffLoopCount, mean, pregMerge); |
| 446 | } | 442 | } |
| 447 | 443 | ||
| 448 | // 整块使用tailCountScale,尾块使用countScale | 444 | // 整块使用tailCountScale,尾块使用countScale |
| 449 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 445 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 450 | pregLoop = UpdateMask<float>(sreg0); | 446 | pregLoop = UpdateMask<float>(sreg0); |
| 451 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | 447 | + LoadAlign(dichotomyAddMeanL, |
| 452 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + | 448 | + tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); |
| 453 | - dichotomyAddPower); | 449 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + |
| 450 | + welfordDiffReminderAlign + dichotomyAddPower); | ||
| 454 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); | 451 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); |
| 455 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); | 452 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); |
| 456 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 453 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 457 | - ReduceSum(mean, sumMean, pregMain); | 454 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 458 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 455 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 459 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, mean, pregMerge); | 456 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, mean, pregMerge); |
| 460 | } | 457 | } |
| 461 | // PART2: 整块剩余部分vcadd回刷UB,使用tailCountScale | 458 | // PART2: 整块剩余部分vcadd回刷UB,使用tailCountScale |
| 462 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); | 459 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); |
| 463 | i++) { | 460 | i++) { |
| 464 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); | 461 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); |
| 465 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); | 462 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); |
| 466 | - ReduceSum(mean, dichotomyAddMeanL, pregMain); | 463 | + Reduce<ReduceType::SUM>(mean, dichotomyAddMeanL, pregMain); |
| 467 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 464 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 468 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, mean, pregMerge); | 465 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, mean, pregMerge); |
| 469 | } | 466 | } |
| 470 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 467 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 471 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); | 468 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); |
| 472 | 469 | ||
| 473 | // 计算rstd | 470 | // 计算rstd |
| 474 | Duplicate(one, float(1.0), pregMain); | 471 | Duplicate(one, float(1.0), pregMain); |
| 475 | Duplicate(mean, mean, pregMain); | 472 | Duplicate(mean, mean, pregMain); |
| 476 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { | 473 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { |
| 477 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 474 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 478 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 475 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 479 | Mul(deltaL, deltaL, deltaL, pregMain); | 476 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 480 | Muls(deltaL, deltaL, tailCnt, pregMain); | 477 | Muls(deltaL, deltaL, tailCnt, pregMain); |
| 481 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 478 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 482 | Sub(deltaR, dichotomyAddMeanR, mean, pregMain); | 479 | Sub(deltaR, dichotomyAddMeanR, mean, pregMain); |
| 483 | Mul(deltaR, deltaR, deltaR, pregMain); | 480 | Mul(deltaR, deltaR, deltaR, pregMain); |
| 484 | Muls(deltaR, deltaR, tailCnt, pregMain); | 481 | Muls(deltaR, deltaR, tailCnt, pregMain); |
| 485 | 482 | ||
| 486 | - DataCopy(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); | 483 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); |
| 487 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 484 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 488 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 485 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 489 | - DataCopy(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); | 486 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); |
| 490 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregMain); | 487 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregMain); |
| 491 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregMain); | 488 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregMain); |
| 492 | 489 | ||
| 493 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 490 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 494 | - ReduceSum(var, sumVar, pregMain); | 491 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 495 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, | 492 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, |
| 496 | - pregMerge); | 493 | + pregMerge); |
| 497 | } | 494 | } |
| 498 | sreg0 = dichotomyAddReminder - welfordDiffLoopCount * VL_FP32; | 495 | sreg0 = dichotomyAddReminder - welfordDiffLoopCount * VL_FP32; |
| 499 | sreg1 = welfordDiffReminder; | 496 | sreg1 = welfordDiffReminder; |
| 500 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { | 497 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { |
| 501 | pregLoop = UpdateMask<float>(sreg0); | 498 | pregLoop = UpdateMask<float>(sreg0); |
| 502 | pregLoop1 = UpdateMask<float>(sreg1); | 499 | pregLoop1 = UpdateMask<float>(sreg1); |
| 503 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); | 500 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); |
| 504 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 501 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 505 | Mul(deltaL, deltaL, deltaL, pregMain); | 502 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 506 | Muls(deltaL, deltaL, tailCnt, pregMain); | 503 | Muls(deltaL, deltaL, tailCnt, pregMain); |
| 507 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); | 504 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); |
| 508 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); | 505 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); |
| 509 | Mul(deltaR, deltaR, deltaR, pregLoop); | 506 | Mul(deltaR, deltaR, deltaR, pregLoop); |
| 510 | Muls(deltaR, deltaR, cnt, pregLoop); | 507 | Muls(deltaR, deltaR, cnt, pregLoop); |
| 511 | Muls(tmp, deltaR, coeff, pregLoop1); | 508 | Muls(tmp, deltaR, coeff, pregLoop1); |
| 512 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(deltaR, tmp, pregLoop1); | 509 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(deltaR, tmp, pregLoop1); |
| 513 | 510 | ||
| 514 | - DataCopy(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32); | 511 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32); |
| 515 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 512 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 516 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 513 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 517 | - DataCopy(dichotomyAddVarR, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); | 514 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); |
| 518 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); | 515 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); |
| 519 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); | 516 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); |
| 520 | 517 | ||
| 521 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 518 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 522 | - ReduceSum(var, sumVar, pregMain); | 519 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 523 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 520 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 524 | dichotomyAddLocal + i + welfordDiffLoopCount, var, pregMerge); | 521 | dichotomyAddLocal + i + welfordDiffLoopCount, var, pregMerge); |
| 525 | } | 522 | } |
| 526 | 523 | ||
| 527 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 524 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 528 | pregLoop = UpdateMask<float>(sreg0); | 525 | pregLoop = UpdateMask<float>(sreg0); |
| 529 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | 526 | + LoadAlign(dichotomyAddMeanL, |
| 527 | + tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | ||
| 530 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 528 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 531 | Mul(deltaL, deltaL, deltaL, pregMain); | 529 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 532 | Muls(deltaL, deltaL, tailCnt, pregMain); | 530 | Muls(deltaL, deltaL, tailCnt, pregMain); |
| 533 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + | 531 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + |
| 534 | - dichotomyAddPower); | 532 | + welfordDiffReminderAlign + dichotomyAddPower); |
| 535 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); | 533 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); |
| 536 | Mul(deltaR, deltaR, deltaR, pregLoop); | 534 | Mul(deltaR, deltaR, deltaR, pregLoop); |
| 537 | Muls(deltaR, deltaR, cnt, pregLoop); | 535 | Muls(deltaR, deltaR, cnt, pregLoop); |
| 538 | 536 | ||
| 539 | - DataCopy(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | 537 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); |
| 540 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 538 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 541 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 539 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 542 | - DataCopy(dichotomyAddVarR, | 540 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + |
| 543 | - tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + dichotomyAddPower); | 541 | + dichotomyAddPower); |
| 544 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); | 542 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); |
| 545 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); | 543 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); |
| 546 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 544 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 547 | - ReduceSum(var, sumVar, pregMain); | 545 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 548 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 546 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 549 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, var, pregMerge); | 547 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, var, pregMerge); |
| 550 | } | 548 | } |
| 551 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); | 549 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); |
| 552 | i++) { | 550 | i++) { |
| 553 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); | 551 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); |
| 554 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 552 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 555 | Mul(deltaL, deltaL, deltaL, pregMain); | 553 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 556 | Muls(deltaL, deltaL, tailCnt, pregMain); | 554 | Muls(deltaL, deltaL, tailCnt, pregMain); |
| 557 | - DataCopy(dichotomyAddVarL, tmpVarLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); | 555 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); |
| 558 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 556 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 559 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 557 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 560 | - ReduceSum(var, dichotomyAddVarL, pregMain); | 558 | + Reduce<ReduceType::SUM>(var, dichotomyAddVarL, pregMain); |
| 561 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 559 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 562 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, var, pregMerge); | 560 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, var, pregMerge); |
| 563 | } | 561 | } |
| 564 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 562 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 565 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); | 563 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); |
| 566 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); | 564 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); |
| 567 | } | 565 | } |
| 568 | } | 566 | } |
| 569 | 567 | ||
| 570 | // welford整块小于二分累加整块，并且小于等于二分累加尾块向上对齐 | 568 | // welford整块小于二分累加整块，并且小于等于二分累加尾块向上对齐 |
| 571 | __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation2( | 569 | __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation2( |
| 572 | - __local_mem__ float* meanLocal, __local_mem__ float* rstdLocal, __local_mem__ float* tmpMeanLocal, | 570 | + __ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, __ubuf__ float* tmpMeanLocal, __ubuf__ float* tmpVarLocal, |
| 573 | - __local_mem__ float* tmpVarLocal, __local_mem__ float* dichotomyAddLocal, uint32_t reduceCount, | 571 | + __ubuf__ float* dichotomyAddLocal, uint32_t reduceCount, uint32_t dichotomyAddPower, uint32_t dichotomyAddK, |
| 574 | - uint32_t dichotomyAddPower, uint32_t dichotomyAddK, uint32_t dichotomyAddLastNum, uint32_t offset, | 572 | + uint32_t dichotomyAddLastNum, uint32_t offset, uint32_t tailSize, float reduceScale, float cnt, float eps) |
| 575 | - uint32_t tailSize, float reduceScale, float cnt, float eps) | ||
| 576 | { | 573 | { |
| 577 | float tailCnt = cnt + float(1.0); | 574 | float tailCnt = cnt + float(1.0); |
| 578 | float coeff = tailCnt / cnt; | 575 | float coeff = tailCnt / cnt; |
| @@ -618,14 +615,14 @@ __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation2( | |||
| 618 | 615 | ||
| 619 | // 整块使用tailCountScale,尾块使用countScale | 616 | // 整块使用tailCountScale,尾块使用countScale |
| 620 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { | 617 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { |
| 621 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 618 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 622 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 619 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 623 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); | 620 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); |
| 624 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregMain); | 621 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregMain); |
| 625 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 622 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 626 | - ReduceSum(mean, sumMean, pregMain); | 623 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 627 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, | 624 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, |
| 628 | - pregMerge); | 625 | + pregMerge); |
| 629 | } | 626 | } |
| 630 | 627 | ||
| 631 | // 处理welford第一次非对齐点, 尾块使用countScale,整块部分使用tailCountScale, 部分使用countScale | 628 | // 处理welford第一次非对齐点, 尾块使用countScale,整块部分使用tailCountScale, 部分使用countScale |
| @@ -634,147 +631,148 @@ __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation2( | |||
| 634 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { | 631 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { |
| 635 | pregLoop = UpdateMask<float>(sreg0); | 632 | pregLoop = UpdateMask<float>(sreg0); |
| 636 | pregLoop1 = UpdateMask<float>(sreg1); | 633 | pregLoop1 = UpdateMask<float>(sreg1); |
| 637 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); | 634 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); |
| 638 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); | 635 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); |
| 639 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); | 636 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); |
| 640 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); | 637 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); |
| 641 | Muls(tmp, dichotomyAddMeanL, coeff, pregLoop1); | 638 | Muls(tmp, dichotomyAddMeanL, coeff, pregLoop1); |
| 642 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(dichotomyAddMeanL, tmp, pregLoop1); | 639 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(dichotomyAddMeanL, tmp, pregLoop1); |
| 643 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 640 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 644 | - ReduceSum(mean, sumMean, pregMain); | 641 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 645 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 642 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 646 | dichotomyAddLocal + i + welfordDiffLoopCount, mean, pregMerge); | 643 | dichotomyAddLocal + i + welfordDiffLoopCount, mean, pregMerge); |
| 647 | } | 644 | } |
| 648 | 645 | ||
| 649 | // 整块使用countScale,尾块使用countScale | 646 | // 整块使用countScale,尾块使用countScale |
| 650 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 647 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 651 | pregLoop = UpdateMask<float>(sreg0); | 648 | pregLoop = UpdateMask<float>(sreg0); |
| 652 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | 649 | + LoadAlign(dichotomyAddMeanL, |
| 653 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + | 650 | + tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); |
| 654 | - dichotomyAddPower); | 651 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + |
| 652 | + welfordDiffReminderAlign + dichotomyAddPower); | ||
| 655 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); | 653 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); |
| 656 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); | 654 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); |
| 657 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 655 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 658 | - ReduceSum(mean, sumMean, pregMain); | 656 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 659 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 657 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 660 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, mean, pregMerge); | 658 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, mean, pregMerge); |
| 661 | } | 659 | } |
| 662 | // PART2: 整块剩余部分vcadd回刷UB,使用countScale | 660 | // PART2: 整块剩余部分vcadd回刷UB,使用countScale |
| 663 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); | 661 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); |
| 664 | i++) { | 662 | i++) { |
| 665 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); | 663 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); |
| 666 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); | 664 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); |
| 667 | - ReduceSum(mean, dichotomyAddMeanL, pregMain); | 665 | + Reduce<ReduceType::SUM>(mean, dichotomyAddMeanL, pregMain); |
| 668 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 666 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 669 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, mean, pregMerge); | 667 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, mean, pregMerge); |
| 670 | } | 668 | } |
| 671 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 669 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 672 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); | 670 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); |
| 673 | 671 | ||
| 674 | // 计算rstd | 672 | // 计算rstd |
| 675 | Duplicate(one, float(1.0), pregMain); | 673 | Duplicate(one, float(1.0), pregMain); |
| 676 | Duplicate(mean, mean, pregMain); | 674 | Duplicate(mean, mean, pregMain); |
| 677 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { | 675 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { |
| 678 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 676 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 679 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 677 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 680 | Mul(deltaL, deltaL, deltaL, pregMain); | 678 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 681 | Muls(deltaL, deltaL, tailCnt, pregMain); | 679 | Muls(deltaL, deltaL, tailCnt, pregMain); |
| 682 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 680 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 683 | Sub(deltaR, dichotomyAddMeanR, mean, pregMain); | 681 | Sub(deltaR, dichotomyAddMeanR, mean, pregMain); |
| 684 | Mul(deltaR, deltaR, deltaR, pregMain); | 682 | Mul(deltaR, deltaR, deltaR, pregMain); |
| 685 | Muls(deltaR, deltaR, cnt, pregMain); | 683 | Muls(deltaR, deltaR, cnt, pregMain); |
| 686 | 684 | ||
| 687 | - DataCopy(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); | 685 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); |
| 688 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 686 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 689 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 687 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 690 | - DataCopy(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); | 688 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); |
| 691 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregMain); | 689 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregMain); |
| 692 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregMain); | 690 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregMain); |
| 693 | 691 | ||
| 694 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 692 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 695 | - ReduceSum(var, sumVar, pregMain); | 693 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 696 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, | 694 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, |
| 697 | - pregMerge); | 695 | + pregMerge); |
| 698 | } | 696 | } |
| 699 | sreg0 = dichotomyAddReminder - welfordDiffLoopCount * VL_FP32; | 697 | sreg0 = dichotomyAddReminder - welfordDiffLoopCount * VL_FP32; |
| 700 | sreg1 = welfordDiffReminder; | 698 | sreg1 = welfordDiffReminder; |
| 701 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { | 699 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { |
| 702 | pregLoop = UpdateMask<float>(sreg0); | 700 | pregLoop = UpdateMask<float>(sreg0); |
| 703 | pregLoop1 = UpdateMask<float>(sreg1); | 701 | pregLoop1 = UpdateMask<float>(sreg1); |
| 704 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); | 702 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32); |
| 705 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 703 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 706 | Mul(deltaL, deltaL, deltaL, pregMain); | 704 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 707 | Muls(deltaL, deltaL, cnt, pregMain); | 705 | Muls(deltaL, deltaL, cnt, pregMain); |
| 708 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); | 706 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); |
| 709 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); | 707 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); |
| 710 | Mul(deltaR, deltaR, deltaR, pregLoop); | 708 | Mul(deltaR, deltaR, deltaR, pregLoop); |
| 711 | Muls(deltaR, deltaR, cnt, pregLoop); | 709 | Muls(deltaR, deltaR, cnt, pregLoop); |
| 712 | Muls(tmp, deltaL, coeff, pregLoop1); | 710 | Muls(tmp, deltaL, coeff, pregLoop1); |
| 713 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(deltaL, tmp, pregLoop1); | 711 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(deltaL, tmp, pregLoop1); |
| 714 | 712 | ||
| 715 | - DataCopy(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32); | 713 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32); |
| 716 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 714 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 717 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 715 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 718 | - DataCopy(dichotomyAddVarR, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); | 716 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddPower); |
| 719 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); | 717 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); |
| 720 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); | 718 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); |
| 721 | 719 | ||
| 722 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 720 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 723 | - ReduceSum(var, sumVar, pregMain); | 721 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 724 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 722 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 725 | dichotomyAddLocal + i + welfordDiffLoopCount, var, pregMerge); | 723 | dichotomyAddLocal + i + welfordDiffLoopCount, var, pregMerge); |
| 726 | } | 724 | } |
| 727 | 725 | ||
| 728 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 726 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 729 | pregLoop = UpdateMask<float>(sreg0); | 727 | pregLoop = UpdateMask<float>(sreg0); |
| 730 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | 728 | + LoadAlign(dichotomyAddMeanL, |
| 729 | + tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | ||
| 731 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 730 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 732 | Mul(deltaL, deltaL, deltaL, pregMain); | 731 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 733 | Muls(deltaL, deltaL, cnt, pregMain); | 732 | Muls(deltaL, deltaL, cnt, pregMain); |
| 734 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + | 733 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + |
| 735 | - dichotomyAddPower); | 734 | + welfordDiffReminderAlign + dichotomyAddPower); |
| 736 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); | 735 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); |
| 737 | Mul(deltaR, deltaR, deltaR, pregLoop); | 736 | Mul(deltaR, deltaR, deltaR, pregLoop); |
| 738 | Muls(deltaR, deltaR, cnt, pregLoop); | 737 | Muls(deltaR, deltaR, cnt, pregLoop); |
| 739 | 738 | ||
| 740 | - DataCopy(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); | 739 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign); |
| 741 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 740 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 742 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 741 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 743 | - DataCopy(dichotomyAddVarR, | 742 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + |
| 744 | - tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + welfordDiffReminderAlign + dichotomyAddPower); | 743 | + dichotomyAddPower); |
| 745 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); | 744 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); |
| 746 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); | 745 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); |
| 747 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 746 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 748 | - ReduceSum(var, sumVar, pregMain); | 747 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 749 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 748 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 750 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, var, pregMerge); | 749 | dichotomyAddLocal + i + welfordDiffLoopCount + welfordReminderLoopCount, var, pregMerge); |
| 751 | } | 750 | } |
| 752 | 751 | ||
| 753 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); | 752 | for (uint16_t i = 0; i < static_cast<uint16_t>(dichotomyAddPowerLoopCount - dichotomyAddReminderRealLoopCount); |
| 754 | i++) { | 753 | i++) { |
| 755 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); | 754 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); |
| 756 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 755 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 757 | Mul(deltaL, deltaL, deltaL, pregMain); | 756 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 758 | Muls(deltaL, deltaL, cnt, pregMain); | 757 | Muls(deltaL, deltaL, cnt, pregMain); |
| 759 | - DataCopy(dichotomyAddVarL, tmpVarLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); | 758 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + (i + dichotomyAddReminderRealLoopCount) * VL_FP32); |
| 760 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 759 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 761 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 760 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 762 | - ReduceSum(var, dichotomyAddVarL, pregMain); | 761 | + Reduce<ReduceType::SUM>(var, dichotomyAddVarL, pregMain); |
| 763 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 762 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 764 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, var, pregMerge); | 763 | dichotomyAddLocal + dichotomyAddReminderRealLoopCount + i, var, pregMerge); |
| 765 | } | 764 | } |
| 766 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 765 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 767 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); | 766 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); |
| 768 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); | 767 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); |
| 769 | } | 768 | } |
| 770 | } | 769 | } |
| 771 | 770 | ||
| 772 | // 场景3：welford整块小于二分累加整块，并且大于二分累加尾块向上对齐 | 771 | // 场景3：welford整块小于二分累加整块，并且大于二分累加尾块向上对齐 |
| 773 | __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation3( | 772 | __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation3( |
| 774 | - __local_mem__ float* meanLocal, __local_mem__ float* rstdLocal, __local_mem__ float* tmpMeanLocal, | 773 | + __ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, __ubuf__ float* tmpMeanLocal, __ubuf__ float* tmpVarLocal, |
| 775 | - __local_mem__ float* tmpVarLocal, __local_mem__ float* dichotomyAddLocal, uint32_t reduceCount, | 774 | + __ubuf__ float* dichotomyAddLocal, uint32_t reduceCount, uint32_t dichotomyAddPower, uint32_t dichotomyAddK, |
| 776 | - uint32_t dichotomyAddPower, uint32_t dichotomyAddK, uint32_t dichotomyAddLastNum, uint32_t offset, | 775 | + uint32_t dichotomyAddLastNum, uint32_t offset, uint32_t tailSize, float reduceScale, float cnt, float eps) |
| 777 | - uint32_t tailSize, float reduceScale, float cnt, float eps) | ||
| 778 | { | 776 | { |
| 779 | float tailCnt = cnt + float(1.0); | 777 | float tailCnt = cnt + float(1.0); |
| 780 | float coeff = tailCnt / cnt; | 778 | float coeff = tailCnt / cnt; |
| @@ -822,50 +820,50 @@ __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation3( | |||
| 822 | // 整块使用tailCountScale, 尾块使用CountScale | 820 | // 整块使用tailCountScale, 尾块使用CountScale |
| 823 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 821 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 824 | pregLoop = UpdateMask<float>(sreg0); | 822 | pregLoop = UpdateMask<float>(sreg0); |
| 825 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 823 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 826 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 824 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 827 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); | 825 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); |
| 828 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); | 826 | Muls(dichotomyAddMeanR, dichotomyAddMeanR, countScale, pregLoop); |
| 829 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); | 827 | Add(sumMean, dichotomyAddMeanL, dichotomyAddMeanR, pregMain); |
| 830 | - ReduceSum(mean, sumMean, pregMain); | 828 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 831 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, | 829 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, mean, |
| 832 | - pregMerge); | 830 | + pregMerge); |
| 833 | } | 831 | } |
| 834 | 832 | ||
| 835 | // 剩余整块需要拆分成多部分 | 833 | // 剩余整块需要拆分成多部分 |
| 836 | // 整块剩余部分回刷UB，整块使用tailCountScale | 834 | // 整块剩余部分回刷UB，整块使用tailCountScale |
| 837 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { | 835 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { |
| 838 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddReminderRoundUp); | 836 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddReminderRoundUp); |
| 839 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); | 837 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, tailCountScale, pregMain); |
| 840 | - ReduceSum(mean, dichotomyAddMeanL, pregMain); | 838 | + Reduce<ReduceType::SUM>(mean, dichotomyAddMeanL, pregMain); |
| 841 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 839 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 842 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, mean, pregMerge); | 840 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, mean, pregMerge); |
| 843 | } | 841 | } |
| 844 | 842 | ||
| 845 | sreg0 = welfordDiffReminder; | 843 | sreg0 = welfordDiffReminder; |
| 846 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { | 844 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { |
| 847 | pregLoop = UpdateMask<float>(sreg0); | 845 | pregLoop = UpdateMask<float>(sreg0); |
| 848 | - DataCopy(dichotomyAddMeanL, | 846 | + LoadAlign(dichotomyAddMeanL, |
| 849 | - tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddReminderRoundUp); | 847 | + tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddReminderRoundUp); |
| 850 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); | 848 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); |
| 851 | Muls(tmp, dichotomyAddMeanL, coeff, pregLoop); | 849 | Muls(tmp, dichotomyAddMeanL, coeff, pregLoop); |
| 852 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(dichotomyAddMeanL, tmp, pregLoop); | 850 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(dichotomyAddMeanL, tmp, pregLoop); |
| 853 | - ReduceSum(mean, dichotomyAddMeanL, pregMain); | 851 | + Reduce<ReduceType::SUM>(mean, dichotomyAddMeanL, pregMain); |
| 854 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 852 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 855 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + i, mean, pregMerge); | 853 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + i, mean, pregMerge); |
| 856 | } | 854 | } |
| 857 | 855 | ||
| 858 | for (uint16_t i = 0; i < dichotomyAddPowerRemainLoopCount; i++) { | 856 | for (uint16_t i = 0; i < dichotomyAddPowerRemainLoopCount; i++) { |
| 859 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddPowerOffset); | 857 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddPowerOffset); |
| 860 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); | 858 | Muls(dichotomyAddMeanL, dichotomyAddMeanL, countScale, pregMain); |
| 861 | - ReduceSum(mean, dichotomyAddMeanL, pregMain); | 859 | + Reduce<ReduceType::SUM>(mean, dichotomyAddMeanL, pregMain); |
| 862 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 860 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 863 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + welfordReminderLoopCount + i, | 861 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + welfordReminderLoopCount + i, |
| 864 | mean, pregMerge); | 862 | mean, pregMerge); |
| 865 | } | 863 | } |
| 866 | 864 | ||
| 867 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 865 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 868 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); | 866 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + offset, mean, pregMerge); |
| 869 | 867 | ||
| 870 | // 计算rstd | 868 | // 计算rstd |
| 871 | Duplicate(one, float(1.0), pregMain); | 869 | Duplicate(one, float(1.0), pregMain); |
| @@ -874,85 +872,84 @@ __aicore__ inline void VFWelfordParallelFinalizeNonAlignSituation3( | |||
| 874 | sreg0 = dichotomyAddReminder; | 872 | sreg0 = dichotomyAddReminder; |
| 875 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { | 873 | for (uint16_t i = 0; i < dichotomyAddReminderLoopCount; i++) { |
| 876 | pregLoop = UpdateMask<float>(sreg0); | 874 | pregLoop = UpdateMask<float>(sreg0); |
| 877 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); | 875 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32); |
| 878 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 876 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 879 | Mul(deltaL, deltaL, deltaL, pregMain); | 877 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 880 | Muls(deltaL, deltaL, tailCnt, pregMain); | 878 | Muls(deltaL, deltaL, tailCnt, pregMain); |
| 881 | - DataCopy(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); | 879 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + i * VL_FP32); |
| 882 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 880 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 883 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 881 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 884 | 882 | ||
| 885 | - DataCopy(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); | 883 | + LoadAlign(dichotomyAddMeanR, tmpMeanLocal + i * VL_FP32 + dichotomyAddPower); |
| 886 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); | 884 | Sub(deltaR, dichotomyAddMeanR, mean, pregLoop); |
| 887 | Mul(deltaR, deltaR, deltaR, pregLoop); | 885 | Mul(deltaR, deltaR, deltaR, pregLoop); |
| 888 | Muls(deltaR, deltaR, cnt, pregLoop); | 886 | Muls(deltaR, deltaR, cnt, pregLoop); |
| 889 | - DataCopy(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); | 887 | + LoadAlign(dichotomyAddVarR, tmpVarLocal + i * VL_FP32 + dichotomyAddPower); |
| 890 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); | 888 | Add(dichotomyAddVarR, dichotomyAddVarR, deltaR, pregLoop); |
| 891 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); | 889 | Muls(dichotomyAddVarR, dichotomyAddVarR, reduceScale, pregLoop); |
| 892 | 890 | ||
| 893 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); | 891 | Add(sumVar, dichotomyAddVarL, dichotomyAddVarR, pregMain); |
| 894 | - ReduceSum(var, sumVar, pregMain); | 892 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 895 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, | 893 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + i, var, |
| 896 | - pregMerge); | 894 | + pregMerge); |
| 897 | } | 895 | } |
| 898 | 896 | ||
| 899 | // 整块剩余部分回刷UB，整块使用tailCountScale | 897 | // 整块剩余部分回刷UB，整块使用tailCountScale |
| 900 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { | 898 | for (uint16_t i = 0; i < welfordDiffLoopCount; i++) { |
| 901 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddReminderRoundUp); | 899 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddReminderRoundUp); |
| 902 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 900 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 903 | Mul(deltaL, deltaL, deltaL, pregMain); | 901 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 904 | Muls(deltaL, deltaL, tailCnt, pregMain); | 902 | Muls(deltaL, deltaL, tailCnt, pregMain); |
| 905 | - DataCopy(dichotomyAddVarL, tmpVarLocal + i * VL_FP32 + dichotomyAddReminderRoundUp); | 903 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + i * VL_FP32 + dichotomyAddReminderRoundUp); |
| 906 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 904 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 907 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 905 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 908 | - ReduceSum(var, dichotomyAddVarL, pregMain); | 906 | + Reduce<ReduceType::SUM>(var, dichotomyAddVarL, pregMain); |
| 909 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 907 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 910 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, var, pregMerge); | 908 | dichotomyAddLocal + dichotomyAddReminderLoopCount + i, var, pregMerge); |
| 911 | } | 909 | } |
| 912 | 910 | ||
| 913 | sreg0 = welfordDiffReminder; | 911 | sreg0 = welfordDiffReminder; |
| 914 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { | 912 | for (uint16_t i = 0; i < welfordReminderLoopCount; i++) { |
| 915 | pregLoop = UpdateMask<float>(sreg0); | 913 | pregLoop = UpdateMask<float>(sreg0); |
| 916 | - DataCopy(dichotomyAddMeanL, | 914 | + LoadAlign(dichotomyAddMeanL, |
| 917 | - tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddReminderRoundUp); | 915 | + tmpMeanLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddReminderRoundUp); |
| 918 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 916 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 919 | Mul(deltaL, deltaL, deltaL, pregMain); | 917 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 920 | Muls(deltaL, deltaL, cnt, pregMain); | 918 | Muls(deltaL, deltaL, cnt, pregMain); |
| 921 | Muls(tmp, deltaL, coeff, pregLoop); | 919 | Muls(tmp, deltaL, coeff, pregLoop); |
| 922 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(deltaL, tmp, pregLoop); | 920 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(deltaL, tmp, pregLoop); |
| 923 | - DataCopy(dichotomyAddVarL, | 921 | + LoadAlign(dichotomyAddVarL, |
| 924 | - tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddReminderRoundUp); | 922 | + tmpVarLocal + (i + welfordDiffLoopCount) * VL_FP32 + dichotomyAddReminderRoundUp); |
| 925 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 923 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 926 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 924 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 927 | - ReduceSum(var, dichotomyAddVarL, pregMain); | 925 | + Reduce<ReduceType::SUM>(var, dichotomyAddVarL, pregMain); |
| 928 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 926 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 929 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + i, var, pregMerge); | 927 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + i, var, pregMerge); |
| 930 | } | 928 | } |
| 931 | 929 | ||
| 932 | for (uint16_t i = 0; i < dichotomyAddPowerRemainLoopCount; i++) { | 930 | for (uint16_t i = 0; i < dichotomyAddPowerRemainLoopCount; i++) { |
| 933 | - DataCopy(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddPowerOffset); | 931 | + LoadAlign(dichotomyAddMeanL, tmpMeanLocal + i * VL_FP32 + dichotomyAddPowerOffset); |
| 934 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); | 932 | Sub(deltaL, dichotomyAddMeanL, mean, pregMain); |
| 935 | Mul(deltaL, deltaL, deltaL, pregMain); | 933 | Mul(deltaL, deltaL, deltaL, pregMain); |
| 936 | Muls(deltaL, deltaL, cnt, pregMain); | 934 | Muls(deltaL, deltaL, cnt, pregMain); |
| 937 | - DataCopy(dichotomyAddVarL, tmpVarLocal + i * VL_FP32 + dichotomyAddPowerOffset); | 935 | + LoadAlign(dichotomyAddVarL, tmpVarLocal + i * VL_FP32 + dichotomyAddPowerOffset); |
| 938 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); | 936 | Add(dichotomyAddVarL, dichotomyAddVarL, deltaL, pregMain); |
| 939 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); | 937 | Muls(dichotomyAddVarL, dichotomyAddVarL, reduceScale, pregMain); |
| 940 | - ReduceSum(var, dichotomyAddVarL, pregMain); | 938 | + Reduce<ReduceType::SUM>(var, dichotomyAddVarL, pregMain); |
| 941 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 939 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 942 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + welfordReminderLoopCount + i, | 940 | dichotomyAddLocal + dichotomyAddReminderLoopCount + welfordDiffLoopCount + welfordReminderLoopCount + i, |
| 943 | var, pregMerge); | 941 | var, pregMerge); |
| 944 | } | 942 | } |
| 945 | 943 | ||
| 946 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 944 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 947 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); | 945 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); |
| 948 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); | 946 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + offset, rstd, pregMerge); |
| 949 | } | 947 | } |
| 950 | } | 948 | } |
| 951 | 949 | ||
| 952 | -__aicore__ inline void VFWelfordParallelFinalizeNonAlign(__local_mem__ float* meanLocal, __local_mem__ float* rstdLocal, | 950 | +__aicore__ inline void VFWelfordParallelFinalizeNonAlign(__ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, |
| 953 | - __local_mem__ float* tmpMeanLocal, | 951 | + __ubuf__ float* tmpMeanLocal, __ubuf__ float* tmpVarLocal, |
| 954 | - __local_mem__ float* tmpVarLocal, | 952 | + __ubuf__ float* dichotomyAddLocal, uint32_t reduceCount, |
| 955 | - __local_mem__ float* dichotomyAddLocal, uint32_t reduceCount, | ||
| 956 | uint32_t dichotomyAddPower, uint32_t dichotomyAddK, | 953 | uint32_t dichotomyAddPower, uint32_t dichotomyAddK, |
| 957 | uint32_t dichotomyAddLastNum, uint32_t offset, | 954 | uint32_t dichotomyAddLastNum, uint32_t offset, |
| 958 | uint32_t tailSize, float reduceScale, float cnt, float eps) | 955 | uint32_t tailSize, float reduceScale, float cnt, float eps) |
| @@ -977,9 +974,9 @@ __aicore__ inline void VFWelfordParallelFinalizeNonAlign(__local_mem__ float* me | |||
| 977 | offset, tailSize, reduceScale, cnt, eps); | 974 | offset, tailSize, reduceScale, cnt, eps); |
| 978 | } | 975 | } |
| 979 | 976 | ||
| 980 | -__aicore__ inline void VFWelfordParallelFinalize(__local_mem__ float* meanLocal, __local_mem__ float* rstdLocal, | 977 | +__aicore__ inline void VFWelfordParallelFinalize(__ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, |
| 981 | - __local_mem__ float* tmpMeanLocal, __local_mem__ float* tmpVarLocal, | 978 | + __ubuf__ float* tmpMeanLocal, __ubuf__ float* tmpVarLocal, |
| 982 | - __local_mem__ float* dichotomyAddLocal, uint32_t reduceCount, | 979 | + __ubuf__ float* dichotomyAddLocal, uint32_t reduceCount, |
| 983 | uint32_t dichotomyAddPower, uint32_t dichotomyAddK, | 980 | uint32_t dichotomyAddPower, uint32_t dichotomyAddK, |
| 984 | uint32_t dichotomyAddLastNum, uint32_t offset, uint32_t tailSize, | 981 | uint32_t dichotomyAddLastNum, uint32_t offset, uint32_t tailSize, |
| 985 | float reduceScale, float scale, float cnt, float eps, | 982 | float reduceScale, float scale, float cnt, float eps, |
| @@ -998,12 +995,11 @@ __aicore__ inline void VFWelfordParallelFinalize(__local_mem__ float* meanLocal, | |||
| 998 | } | 995 | } |
| 999 | 996 | ||
| 1000 | template <typename T> | 997 | template <typename T> |
| 1001 | -__aicore__ inline void CalMeanAndRstdByDichotomyAdd(__local_mem__ T* xLocal, __local_mem__ float* meanLocal, | 998 | +__aicore__ inline void CalMeanAndRstdByDichotomyAdd(__ubuf__ T* xLocal, __ubuf__ float* meanLocal, |
| 1002 | - __local_mem__ float* rstdLocal, | 999 | + __ubuf__ float* rstdLocal, __ubuf__ float* dichotomyAddLocal, |
| 1003 | - __local_mem__ float* dichotomyAddLocal, uint16_t numPerCoreProcess, | 1000 | + uint16_t numPerCoreProcess, uint32_t dichotomyAddPower, |
| 1004 | - uint32_t dichotomyAddPower, uint32_t dichotomyAddK, | 1001 | + uint32_t dichotomyAddK, uint32_t dichotomyAddLastNum, |
| 1005 | - uint32_t dichotomyAddLastNum, uint64_t reduceCount, float scale, | 1002 | + uint64_t reduceCount, float scale, float eps) |
| 1006 | - float eps) | ||
| 1007 | { | 1003 | { |
| 1008 | uint32_t dichotomyAddReminder = reduceCount - dichotomyAddPower; | 1004 | uint32_t dichotomyAddReminder = reduceCount - dichotomyAddPower; |
| 1009 | uint16_t dichotomyAddReminderLoopCount = CeilDiv(dichotomyAddReminder, VL_FP32); | 1005 | uint16_t dichotomyAddReminderLoopCount = CeilDiv(dichotomyAddReminder, VL_FP32); |
| @@ -1035,9 +1031,9 @@ __aicore__ inline void CalMeanAndRstdByDichotomyAdd(__local_mem__ T* xLocal, __l | |||
| 1035 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); | 1031 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); |
| 1036 | Muls(dichotomyAddR, dichotomyAddR, scale, pregLoop); | 1032 | Muls(dichotomyAddR, dichotomyAddR, scale, pregLoop); |
| 1037 | Add(sumMean, dichotomyAddL, dichotomyAddR, pregMain); | 1033 | Add(sumMean, dichotomyAddL, dichotomyAddR, pregMain); |
| 1038 | - ReduceSum(mean, sumMean, pregMain); | 1034 | + Reduce<ReduceType::SUM>(mean, sumMean, pregMain); |
| 1039 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + j, mean, | 1035 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + j, mean, |
| 1040 | - pregMerge); | 1036 | + pregMerge); |
| 1041 | } | 1037 | } |
| 1042 | 1038 | ||
| 1043 | // 整块剩余部分vcadd回刷UB | 1039 | // 整块剩余部分vcadd回刷UB |
| @@ -1046,13 +1042,13 @@ __aicore__ inline void CalMeanAndRstdByDichotomyAdd(__local_mem__ T* xLocal, __l | |||
| 1046 | LoadInputData<T>(dichotomyAddL, xLocal, pregMain, | 1042 | LoadInputData<T>(dichotomyAddL, xLocal, pregMain, |
| 1047 | i * elemNumAlign + (j + dichotomyAddReminderLoopCount) * VL_FP32); | 1043 | i * elemNumAlign + (j + dichotomyAddReminderLoopCount) * VL_FP32); |
| 1048 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); | 1044 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); |
| 1049 | - ReduceSum(mean, dichotomyAddL, pregMain); | 1045 | + Reduce<ReduceType::SUM>(mean, dichotomyAddL, pregMain); |
| 1050 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 1046 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 1051 | dichotomyAddLocal + dichotomyAddReminderLoopCount + j, mean, pregMerge); | 1047 | dichotomyAddLocal + dichotomyAddReminderLoopCount + j, mean, pregMerge); |
| 1052 | } | 1048 | } |
| 1053 | 1049 | ||
| 1054 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 1050 | NormCommon::DichotomyAdd(mean, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 1055 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + i, mean, pregMerge); | 1051 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + i, mean, pregMerge); |
| 1056 | // 计算rstd | 1052 | // 计算rstd |
| 1057 | Duplicate(one, float(1.0), pregMain); | 1053 | Duplicate(one, float(1.0), pregMain); |
| 1058 | Duplicate(mean, mean, pregMain); | 1054 | Duplicate(mean, mean, pregMain); |
| @@ -1068,9 +1064,9 @@ __aicore__ inline void CalMeanAndRstdByDichotomyAdd(__local_mem__ T* xLocal, __l | |||
| 1068 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); | 1064 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); |
| 1069 | Muls(dichotomyAddR, dichotomyAddR, scale, pregLoop); | 1065 | Muls(dichotomyAddR, dichotomyAddR, scale, pregLoop); |
| 1070 | Add(sumVar, dichotomyAddL, dichotomyAddR, pregMain); | 1066 | Add(sumVar, dichotomyAddL, dichotomyAddR, pregMain); |
| 1071 | - ReduceSum(var, sumVar, pregMain); | 1067 | + Reduce<ReduceType::SUM>(var, sumVar, pregMain); |
| 1072 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + j, var, | 1068 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(dichotomyAddLocal + j, var, |
| 1073 | - pregMerge); | 1069 | + pregMerge); |
| 1074 | } | 1070 | } |
| 1075 | 1071 | ||
| 1076 | // 整块剩余部分vcadd回刷UB | 1072 | // 整块剩余部分vcadd回刷UB |
| @@ -1081,22 +1077,21 @@ __aicore__ inline void CalMeanAndRstdByDichotomyAdd(__local_mem__ T* xLocal, __l | |||
| 1081 | Sub(dichotomyAddL, dichotomyAddL, mean, pregMain); | 1077 | Sub(dichotomyAddL, dichotomyAddL, mean, pregMain); |
| 1082 | Mul(dichotomyAddL, dichotomyAddL, dichotomyAddL, pregMain); | 1078 | Mul(dichotomyAddL, dichotomyAddL, dichotomyAddL, pregMain); |
| 1083 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); | 1079 | Muls(dichotomyAddL, dichotomyAddL, scale, pregMain); |
| 1084 | - ReduceSum(var, dichotomyAddL, pregMain); | 1080 | + Reduce<ReduceType::SUM>(var, dichotomyAddL, pregMain); |
| 1085 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( | 1081 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 1086 | dichotomyAddLocal + dichotomyAddReminderLoopCount + j, var, pregMerge); | 1082 | dichotomyAddLocal + dichotomyAddReminderLoopCount + j, var, pregMerge); |
| 1087 | } | 1083 | } |
| 1088 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); | 1084 | NormCommon::DichotomyAdd(var, dichotomyAddLocal, dichotomyAddK, innerLoopCountOrigin, dichotomyAddLastNum); |
| 1089 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); | 1085 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); |
| 1090 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + i, rstd, pregMerge); | 1086 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + i, rstd, pregMerge); |
| 1091 | } | 1087 | } |
| 1092 | } | 1088 | } |
| 1093 | } | 1089 | } |
| 1094 | 1090 | ||
| 1095 | // R轴小于64 | 1091 | // R轴小于64 |
| 1096 | template <typename T> | 1092 | template <typename T> |
| 1097 | -__aicore__ inline void CalMeanAndRstdSpecial(__local_mem__ T* xLocal, __local_mem__ float* meanLocal, | 1093 | +__aicore__ inline void CalMeanAndRstdSpecial(__ubuf__ T* xLocal, __ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, |
| 1098 | - __local_mem__ float* rstdLocal, uint16_t numPerCoreProcess, | 1094 | + uint16_t numPerCoreProcess, uint64_t reduceCount, float scale, float eps) |
| 1099 | - uint64_t reduceCount, float scale, float eps) | ||
| 1100 | { | 1095 | { |
| 1101 | uint32_t elemNumAlign = RoundUp<T>(reduceCount); | 1096 | uint32_t elemNumAlign = RoundUp<T>(reduceCount); |
| 1102 | __VEC_SCOPE__ | 1097 | __VEC_SCOPE__ |
| @@ -1116,25 +1111,25 @@ __aicore__ inline void CalMeanAndRstdSpecial(__local_mem__ T* xLocal, __local_me | |||
| 1116 | pregLoop = UpdateMask<float>(sreg0); | 1111 | pregLoop = UpdateMask<float>(sreg0); |
| 1117 | LoadInputData<T>(x, xLocal, pregLoop, i * elemNumAlign); | 1112 | LoadInputData<T>(x, xLocal, pregLoop, i * elemNumAlign); |
| 1118 | Muls(xScale, x, scale, pregLoop); | 1113 | Muls(xScale, x, scale, pregLoop); |
| 1119 | - ReduceSum(mean, xScale, pregLoop); | 1114 | + Reduce<ReduceType::SUM>(mean, xScale, pregLoop); |
| 1120 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + i, mean, pregMerge); | 1115 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(meanLocal + i, mean, pregMerge); |
| 1121 | 1116 | ||
| 1122 | Duplicate(mean, mean, pregMain); | 1117 | Duplicate(mean, mean, pregMain); |
| 1123 | Sub(x, x, mean, pregLoop); | 1118 | Sub(x, x, mean, pregLoop); |
| 1124 | Mul(x, x, x, pregLoop); | 1119 | Mul(x, x, x, pregLoop); |
| 1125 | Muls(xScale, x, scale, pregLoop); | 1120 | Muls(xScale, x, scale, pregLoop); |
| 1126 | - ReduceSum(var, xScale, pregLoop); | 1121 | + Reduce<ReduceType::SUM>(var, xScale, pregLoop); |
| 1127 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); | 1122 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(var, rstd, pregMerge, eps); |
| 1128 | - DataCopy<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + i, rstd, pregMerge); | 1123 | + StoreAlign<float, AscendC::MicroAPI::StoreDist::DIST_FIRST_ELEMENT_B32>(rstdLocal + i, rstd, pregMerge); |
| 1129 | } | 1124 | } |
| 1130 | } | 1125 | } |
| 1131 | } | 1126 | } |
| 1132 | 1127 | ||
| 1133 | template <typename T> | 1128 | template <typename T> |
| 1134 | -__aicore__ inline void CalMeanAndRstd(__local_mem__ T* xLocal, __local_mem__ float* meanLocal, | 1129 | +__aicore__ inline void CalMeanAndRstd(__ubuf__ T* xLocal, __ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, |
| 1135 | - __local_mem__ float* rstdLocal, __local_mem__ float* dichotomyAddLocal, | 1130 | + __ubuf__ float* dichotomyAddLocal, uint16_t numPerCoreProcess, |
| 1136 | - uint16_t numPerCoreProcess, uint32_t dichotomyAddPower, uint32_t dichotomyAddK, | 1131 | + uint32_t dichotomyAddPower, uint32_t dichotomyAddK, uint32_t dichotomyAddLastNum, |
| 1137 | - uint32_t dichotomyAddLastNum, uint64_t reduceCount, float scale, float eps) | 1132 | + uint64_t reduceCount, float scale, float eps) |
| 1138 | { | 1133 | { |
| 1139 | if (dichotomyAddPower >= VL_FP32) { | 1134 | if (dichotomyAddPower >= VL_FP32) { |
| 1140 | CalMeanAndRstdByDichotomyAdd(xLocal, meanLocal, rstdLocal, dichotomyAddLocal, numPerCoreProcess, | 1135 | CalMeanAndRstdByDichotomyAdd(xLocal, meanLocal, rstdLocal, dichotomyAddLocal, numPerCoreProcess, |
| @@ -1155,10 +1150,9 @@ __aicore__ inline void VFInnerNormalize(RegTensor<float>& x, RegTensor<float>& m | |||
| 1155 | } | 1150 | } |
| 1156 | 1151 | ||
| 1157 | template <typename T1, typename T2> | 1152 | template <typename T1, typename T2> |
| 1158 | -__aicore__ inline void VFNormalizeUnAlign(__local_mem__ T1* xLocal, __local_mem__ T2* gammaLocal, | 1153 | +__aicore__ inline void VFNormalizeUnAlign(__ubuf__ T1* xLocal, __ubuf__ T2* gammaLocal, __ubuf__ T2* betaLocal, |
| 1159 | - __local_mem__ T2* betaLocal, __local_mem__ float* meanLocal, | 1154 | + __ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, __ubuf__ T1* yLocal, |
| 1160 | - __local_mem__ float* rstdLocal, __local_mem__ T1* yLocal, uint32_t rowsCount, | 1155 | + uint32_t rowsCount, int32_t reduceCount) |
| 1161 | - int32_t reduceCount) | ||
| 1162 | { | 1156 | { |
| 1163 | uint16_t VL = GetVLSize<T1>(); | 1157 | uint16_t VL = GetVLSize<T1>(); |
| 1164 | uint16_t loopCount = reduceCount / VL; | 1158 | uint16_t loopCount = reduceCount / VL; |
| @@ -1179,11 +1173,11 @@ __aicore__ inline void VFNormalizeUnAlign(__local_mem__ T1* xLocal, __local_mem_ | |||
| 1179 | MaskReg pregLoop; | 1173 | MaskReg pregLoop; |
| 1180 | MaskReg pregMain = CreateMask<T1, AscendC::MicroAPI::MaskPattern::ALL>(); | 1174 | MaskReg pregMain = CreateMask<T1, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 1181 | 1175 | ||
| 1182 | - UnalignReg uSrc; | 1176 | + UnalignRegForLoad uSrc; |
| 1183 | - UnalignReg uDst; | 1177 | + UnalignRegForStore uDst; |
| 1184 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstd, rstdLocal); | 1178 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstd, rstdLocal); |
| 1185 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(mean, meanLocal); | 1179 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(mean, meanLocal); |
| 1186 | - DataCopyUnAlignPre<T1>(uSrc, xLocal); | 1180 | + LoadUnAlignPre<T1>(uSrc, xLocal); |
| 1187 | for (uint16_t i = 0; i < static_cast<uint16_t>(rowsCount); i++) { | 1181 | for (uint16_t i = 0; i < static_cast<uint16_t>(rowsCount); i++) { |
| 1188 | LoadGammaAndBetaData<T2>(gamma, beta, gammaLocal, betaLocal, pregMain, i); | 1182 | LoadGammaAndBetaData<T2>(gamma, beta, gammaLocal, betaLocal, pregMain, i); |
| 1189 | if constexpr (IsSameType<T1, half>::value || IsSameType<T1, bfloat16_t>::value) { | 1183 | if constexpr (IsSameType<T1, half>::value || IsSameType<T1, bfloat16_t>::value) { |
| @@ -1192,7 +1186,7 @@ __aicore__ inline void VFNormalizeUnAlign(__local_mem__ T1* xLocal, __local_mem_ | |||
| 1192 | RegTensor<T1> yOddTmp; | 1186 | RegTensor<T1> yOddTmp; |
| 1193 | RegTensor<T1> yTmp; | 1187 | RegTensor<T1> yTmp; |
| 1194 | for (uint16_t j = 0; j < loopCount; j++) { | 1188 | for (uint16_t j = 0; j < loopCount; j++) { |
| 1195 | - DataCopyUnAlign(xTmp, uSrc, xLocal, VL); | 1189 | + LoadUnAlign(xTmp, uSrc, xLocal, VL); |
| 1196 | Cast<float, T1, castTraitB162B32Even>(xEven, xTmp, pregMain); | 1190 | Cast<float, T1, castTraitB162B32Even>(xEven, xTmp, pregMain); |
| 1197 | Cast<float, T1, castTraitB162B32Odd>(xOdd, xTmp, pregMain); | 1191 | Cast<float, T1, castTraitB162B32Odd>(xOdd, xTmp, pregMain); |
| 1198 | VFInnerNormalize(xEven, mean, rstd, gamma, beta, yEven, pregMain); | 1192 | VFInnerNormalize(xEven, mean, rstd, gamma, beta, yEven, pregMain); |
| @@ -1201,12 +1195,12 @@ __aicore__ inline void VFNormalizeUnAlign(__local_mem__ T1* xLocal, __local_mem_ | |||
| 1201 | Cast<T1, float, castTraitB322B16Odd>(yOddTmp, yOdd, pregMain); | 1195 | Cast<T1, float, castTraitB322B16Odd>(yOddTmp, yOdd, pregMain); |
| 1202 | Or((RegTensor<int16_t>&)yTmp, (RegTensor<int16_t>&)yEvenTmp, (RegTensor<int16_t>&)yOddTmp, | 1196 | Or((RegTensor<int16_t>&)yTmp, (RegTensor<int16_t>&)yEvenTmp, (RegTensor<int16_t>&)yOddTmp, |
| 1203 | pregMain); | 1197 | pregMain); |
| 1204 | - DataCopyUnAlign(yLocal, yTmp, uDst, VL); | 1198 | + StoreUnAlign(yLocal, yTmp, uDst, VL); |
| 1205 | } | 1199 | } |
| 1206 | uint32_t sreg0 = tailNum; | 1200 | uint32_t sreg0 = tailNum; |
| 1207 | for (uint16_t k = 0; k < tailLoop; k++) { | 1201 | for (uint16_t k = 0; k < tailLoop; k++) { |
| 1208 | pregLoop = UpdateMask<half>(sreg0); | 1202 | pregLoop = UpdateMask<half>(sreg0); |
| 1209 | - DataCopyUnAlign(xTmp, uSrc, xLocal, tailNum); | 1203 | + LoadUnAlign(xTmp, uSrc, xLocal, tailNum); |
| 1210 | Cast<float, T1, castTraitB162B32Even>(xEven, xTmp, pregLoop); | 1204 | Cast<float, T1, castTraitB162B32Even>(xEven, xTmp, pregLoop); |
| 1211 | Cast<float, T1, castTraitB162B32Odd>(xOdd, xTmp, pregLoop); | 1205 | Cast<float, T1, castTraitB162B32Odd>(xOdd, xTmp, pregLoop); |
| 1212 | VFInnerNormalize(xEven, mean, rstd, gamma, beta, yEven, pregLoop); | 1206 | VFInnerNormalize(xEven, mean, rstd, gamma, beta, yEven, pregLoop); |
| @@ -1215,33 +1209,32 @@ __aicore__ inline void VFNormalizeUnAlign(__local_mem__ T1* xLocal, __local_mem_ | |||
| 1215 | Cast<T1, float, castTraitB322B16Odd>(yOddTmp, yOdd, pregLoop); | 1209 | Cast<T1, float, castTraitB322B16Odd>(yOddTmp, yOdd, pregLoop); |
| 1216 | Or((RegTensor<int16_t>&)yTmp, (RegTensor<int16_t>&)yEvenTmp, (RegTensor<int16_t>&)yOddTmp, | 1210 | Or((RegTensor<int16_t>&)yTmp, (RegTensor<int16_t>&)yEvenTmp, (RegTensor<int16_t>&)yOddTmp, |
| 1217 | pregLoop); | 1211 | pregLoop); |
| 1218 | - DataCopyUnAlign(yLocal, yTmp, uDst, tailNum); | 1212 | + StoreUnAlign(yLocal, yTmp, uDst, tailNum); |
| 1219 | } | 1213 | } |
| 1220 | - DataCopyUnAlignPost(yLocal, uDst, 0); | 1214 | + StoreUnAlignPost(yLocal, uDst, 0); |
| 1221 | } else { | 1215 | } else { |
| 1222 | for (uint16_t j = 0; j < loopCount; j++) { | 1216 | for (uint16_t j = 0; j < loopCount; j++) { |
| 1223 | - DataCopyUnAlign(x, uSrc, xLocal, VL_FP32); | 1217 | + LoadUnAlign(x, uSrc, xLocal, VL_FP32); |
| 1224 | VFInnerNormalize(x, mean, rstd, gamma, beta, y, pregMain); | 1218 | VFInnerNormalize(x, mean, rstd, gamma, beta, y, pregMain); |
| 1225 | - DataCopyUnAlign(yLocal, y, uDst, VL_FP32); | 1219 | + StoreUnAlign(yLocal, y, uDst, VL_FP32); |
| 1226 | } | 1220 | } |
| 1227 | uint32_t sreg0 = tailNum; | 1221 | uint32_t sreg0 = tailNum; |
| 1228 | for (uint16_t k = 0; k < tailLoop; k++) { | 1222 | for (uint16_t k = 0; k < tailLoop; k++) { |
| 1229 | pregLoop = UpdateMask<float>(sreg0); | 1223 | pregLoop = UpdateMask<float>(sreg0); |
| 1230 | - DataCopyUnAlign(x, uSrc, xLocal, tailNum); | 1224 | + LoadUnAlign(x, uSrc, xLocal, tailNum); |
| 1231 | VFInnerNormalize(x, mean, rstd, gamma, beta, y, pregLoop); | 1225 | VFInnerNormalize(x, mean, rstd, gamma, beta, y, pregLoop); |
| 1232 | - DataCopyUnAlign(yLocal, y, uDst, tailNum); | 1226 | + StoreUnAlign(yLocal, y, uDst, tailNum); |
| 1233 | } | 1227 | } |
| 1234 | - DataCopyUnAlignPost(yLocal, uDst, 0); | 1228 | + StoreUnAlignPost(yLocal, uDst, 0); |
| 1235 | } | 1229 | } |
| 1236 | } | 1230 | } |
| 1237 | } | 1231 | } |
| 1238 | } | 1232 | } |
| 1239 | 1233 | ||
| 1240 | template <typename T1, typename T2> | 1234 | template <typename T1, typename T2> |
| 1241 | -__aicore__ inline void VFNormalizeAlign(__local_mem__ T1* xLocal, __local_mem__ T2* gammaLocal, | 1235 | +__aicore__ inline void VFNormalizeAlign(__ubuf__ T1* xLocal, __ubuf__ T2* gammaLocal, __ubuf__ T2* betaLocal, |
| 1242 | - __local_mem__ T2* betaLocal, __local_mem__ float* meanLocal, | 1236 | + __ubuf__ float* meanLocal, __ubuf__ float* rstdLocal, __ubuf__ T1* yLocal, |
| 1243 | - __local_mem__ float* rstdLocal, __local_mem__ T1* yLocal, uint16_t rowsCount, | 1237 | + uint16_t rowsCount, int32_t reduceCount) |
| 1244 | - int32_t reduceCount) | ||
| 1245 | { | 1238 | { |
| 1246 | uint16_t loopCount = CeilDiv(reduceCount, VL_FP32); | 1239 | uint16_t loopCount = CeilDiv(reduceCount, VL_FP32); |
| 1247 | uint32_t reduceCountAlign = RoundUp<T1>(reduceCount); | 1240 | uint32_t reduceCountAlign = RoundUp<T1>(reduceCount); |
| @@ -1255,8 +1248,8 @@ __aicore__ inline void VFNormalizeAlign(__local_mem__ T1* xLocal, __local_mem__ | |||
| 1255 | RegTensor<float> y; | 1248 | RegTensor<float> y; |
| 1256 | MaskReg pregLoop; | 1249 | MaskReg pregLoop; |
| 1257 | MaskReg pregMain = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 1250 | MaskReg pregMain = CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 1258 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstd, rstdLocal); | 1251 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstd, rstdLocal); |
| 1259 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(mean, meanLocal); | 1252 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(mean, meanLocal); |
| 1260 | for (uint16_t i = 0; i < rowsCount; i++) { | 1253 | for (uint16_t i = 0; i < rowsCount; i++) { |
| 1261 | uint32_t sreg0 = reduceCount; | 1254 | uint32_t sreg0 = reduceCount; |
| 1262 | LoadGammaAndBetaData<T2>(gamma, beta, gammaLocal, betaLocal, pregMain, i); | 1255 | LoadGammaAndBetaData<T2>(gamma, beta, gammaLocal, betaLocal, pregMain, i); |
| @@ -1338,10 +1331,10 @@ __aicore__ inline void ProcessMeanAndRstd(LocalTensor<float>& meanTensor, LocalT | |||
| 1338 | if constexpr (IsSameType<T1, float>::value) { | 1331 | if constexpr (IsSameType<T1, float>::value) { |
| 1339 | CopyMeanAndRstd2Gm<float>(meanGm[gmOffset], rstdGm[gmOffset], meanTensor, rstdTensor, 1, curNumPerCore); | 1332 | CopyMeanAndRstd2Gm<float>(meanGm[gmOffset], rstdGm[gmOffset], meanTensor, rstdTensor, 1, curNumPerCore); |
| 1340 | } else { | 1333 | } else { |
| 1341 | - __local_mem__ T1* meanOutLocal = (__local_mem__ T1*)meanOutTensor.GetPhyAddr(); | 1334 | + __ubuf__ T1* meanOutLocal = (__ubuf__ T1*)meanOutTensor.GetPhyAddr(); |
| 1342 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 1335 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 1343 | - __local_mem__ T1* rstdOutLocal = (__local_mem__ T1*)rstdOutTensor.GetPhyAddr(); | 1336 | + __ubuf__ T1* rstdOutLocal = (__ubuf__ T1*)rstdOutTensor.GetPhyAddr(); |
| 1344 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 1337 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 1345 | uint16_t loopCount = CeilDiv(curNumPerCore, VL_FP32); | 1338 | uint16_t loopCount = CeilDiv(curNumPerCore, VL_FP32); |
| 1346 | __VEC_SCOPE__ | 1339 | __VEC_SCOPE__ |
| 1347 | { | 1340 | { |
| @@ -1353,14 +1346,14 @@ __aicore__ inline void ProcessMeanAndRstd(LocalTensor<float>& meanTensor, LocalT | |||
| 1353 | RegTensor<T1> rstdOut; | 1346 | RegTensor<T1> rstdOut; |
| 1354 | for (uint16_t i = 0; i < loopCount; i++) { | 1347 | for (uint16_t i = 0; i < loopCount; i++) { |
| 1355 | pregLoop = UpdateMask<float>(sreg0); | 1348 | pregLoop = UpdateMask<float>(sreg0); |
| 1356 | - DataCopy(mean, meanLocal + i * VL_FP32); | 1349 | + LoadAlign(mean, meanLocal + i * VL_FP32); |
| 1357 | - DataCopy(rstd, rstdLocal + i * VL_FP32); | 1350 | + LoadAlign(rstd, rstdLocal + i * VL_FP32); |
| 1358 | Cast<T1, float, castTraitB322B16Even>(meanOut, mean, pregLoop); | 1351 | Cast<T1, float, castTraitB322B16Even>(meanOut, mean, pregLoop); |
| 1359 | Cast<T1, float, castTraitB322B16Even>(rstdOut, rstd, pregLoop); | 1352 | Cast<T1, float, castTraitB322B16Even>(rstdOut, rstd, pregLoop); |
| 1360 | - DataCopy<T1, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(meanOutLocal + i * VL_FP32, meanOut, | 1353 | + StoreAlign<T1, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(meanOutLocal + i * VL_FP32, meanOut, |
| 1361 | - pregLoop); | 1354 | + pregLoop); |
| 1362 | - DataCopy<T1, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(rstdOutLocal + i * VL_FP32, rstdOut, | 1355 | + StoreAlign<T1, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(rstdOutLocal + i * VL_FP32, rstdOut, |
| 1363 | - pregLoop); | 1356 | + pregLoop); |
| 1364 | } | 1357 | } |
| 1365 | } | 1358 | } |
| 1366 | event_t eventIdVToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); | 1359 | event_t eventIdVToMte3 = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3)); |
| @@ -85,7 +85,7 @@ private: | |||
| 85 | auto eventIDVToMte2Pong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>()); | 85 | auto eventIDVToMte2Pong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>()); |
| 86 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 86 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 87 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 87 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 88 | - __local_mem__ float* dichotomyAddLocal = (__local_mem__ float*)dichotomyAddTensor.GetPhyAddr(); | 88 | + __ubuf__ float* dichotomyAddLocal = (__ubuf__ float*)dichotomyAddTensor.GetPhyAddr(); |
| 89 | for (int64_t i = 0; i < numPerCoreExtent; i++) { | 89 | for (int64_t i = 0; i < numPerCoreExtent; i++) { |
| 90 | if (i == numPerCoreExtent - 1) { | 90 | if (i == numPerCoreExtent - 1) { |
| 91 | numPerCoreProcess = numPerCoreTail; | 91 | numPerCoreProcess = numPerCoreTail; |
| @@ -99,9 +99,9 @@ private: | |||
| 99 | CopyX2UB<T1>(xGm[xGmOffset], xTensor[xUbOffset], numPerCoreProcess, elemNum); | 99 | CopyX2UB<T1>(xGm[xGmOffset], xTensor[xUbOffset], numPerCoreProcess, elemNum); |
| 100 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 100 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 101 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 101 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 102 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xTensor[xUbOffset].GetPhyAddr(); | 102 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xTensor[xUbOffset].GetPhyAddr(); |
| 103 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[onceNumPerCore * i].GetPhyAddr(); | 103 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[onceNumPerCore * i].GetPhyAddr(); |
| 104 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[onceNumPerCore * i].GetPhyAddr(); | 104 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[onceNumPerCore * i].GetPhyAddr(); |
| 105 | if (i > 1) { | 105 | if (i > 1) { |
| 106 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); | 106 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); |
| 107 | } | 107 | } |
| @@ -131,20 +131,17 @@ private: | |||
| 131 | __aicore__ inline void NormalizeAndSwish(uint32_t xUbOffset, uint32_t numPerCoreoffset, int64_t numPerCoreProcess, | 131 | __aicore__ inline void NormalizeAndSwish(uint32_t xUbOffset, uint32_t numPerCoreoffset, int64_t numPerCoreProcess, |
| 132 | uint32_t numPerCoreLoop) | 132 | uint32_t numPerCoreLoop) |
| 133 | { | 133 | { |
| 134 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xTensor[xUbOffset].GetPhyAddr(); | 134 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xTensor[xUbOffset].GetPhyAddr(); |
| 135 | - __local_mem__ T1* yOutLocal = (__local_mem__ T1*)yTensor[xUbOffset].GetPhyAddr(); | 135 | + __ubuf__ T1* yOutLocal = (__ubuf__ T1*)yTensor[xUbOffset].GetPhyAddr(); |
| 136 | for (int64_t i = 0; i < numPerCoreProcess; i++) { | 136 | for (int64_t i = 0; i < numPerCoreProcess; i++) { |
| 137 | uint64_t gammaOffset = ((blockIdx * tiling->numPerCore + numPerCoreoffset + i) % numGroups) * shapeD; | 137 | uint64_t gammaOffset = ((blockIdx * tiling->numPerCore + numPerCoreoffset + i) % numGroups) * shapeD; |
| 138 | uint64_t betaOffset = gammaOffset; | 138 | uint64_t betaOffset = gammaOffset; |
| 139 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xTensor[xUbOffset + i * elemNumAlign].GetPhyAddr(); | 139 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xTensor[xUbOffset + i * elemNumAlign].GetPhyAddr(); |
| 140 | - __local_mem__ T1* yOutLocal = (__local_mem__ T1*)yTensor[xUbOffset + i * elemNumAlign].GetPhyAddr(); | 140 | + __ubuf__ T1* yOutLocal = (__ubuf__ T1*)yTensor[xUbOffset + i * elemNumAlign].GetPhyAddr(); |
| 141 | - __local_mem__ T2* gammaLocal = hasGamma ? (__local_mem__ T2*)gammaTensor[gammaOffset].GetPhyAddr() : | 141 | + __ubuf__ T2* gammaLocal = hasGamma ? (__ubuf__ T2*)gammaTensor[gammaOffset].GetPhyAddr() : nullptr; |
| 142 | - nullptr; | 142 | + __ubuf__ T2* betaLocal = hasBeta ? (__ubuf__ T2*)betaTensor[betaOffset].GetPhyAddr() : nullptr; |
| 143 | - __local_mem__ T2* betaLocal = hasBeta ? (__local_mem__ T2*)betaTensor[betaOffset].GetPhyAddr() : nullptr; | 143 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[numPerCoreLoop * onceNumPerCore + i].GetPhyAddr(); |
| 144 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[numPerCoreLoop * onceNumPerCore + i] | 144 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[numPerCoreLoop * onceNumPerCore + i].GetPhyAddr(); |
| 145 | - .GetPhyAddr(); | ||
| 146 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[numPerCoreLoop * onceNumPerCore + i] | ||
| 147 | - .GetPhyAddr(); | ||
| 148 | VFNormalizeUnAlign<T1, T2>(xLocal, gammaLocal, betaLocal, meanLocal, rstdLocal, yOutLocal, shapeD, hwNum); | 145 | VFNormalizeUnAlign<T1, T2>(xLocal, gammaLocal, betaLocal, meanLocal, rstdLocal, yOutLocal, shapeD, hwNum); |
| 149 | } | 146 | } |
| 150 | } | 147 | } |
| @@ -84,7 +84,7 @@ private: | |||
| 84 | auto eventIDVToMte2Pong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>()); | 84 | auto eventIDVToMte2Pong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>()); |
| 85 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 85 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 86 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 86 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 87 | - __local_mem__ float* dichotomyAddLocal = (__local_mem__ float*)dichotomyAddTensor.GetPhyAddr(); | 87 | + __ubuf__ float* dichotomyAddLocal = (__ubuf__ float*)dichotomyAddTensor.GetPhyAddr(); |
| 88 | for (int64_t i = 0; i < numPerCoreExtent; i++) { | 88 | for (int64_t i = 0; i < numPerCoreExtent; i++) { |
| 89 | if (i == numPerCoreExtent - 1) { | 89 | if (i == numPerCoreExtent - 1) { |
| 90 | numPerCoreProcess = numPerCoreTail; | 90 | numPerCoreProcess = numPerCoreTail; |
| @@ -98,9 +98,9 @@ private: | |||
| 98 | CopyX2UB<T1>(xGm[xGmOffset], xTensor[xUbOffset], numPerCoreProcess, elemNum); | 98 | CopyX2UB<T1>(xGm[xGmOffset], xTensor[xUbOffset], numPerCoreProcess, elemNum); |
| 99 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 99 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 100 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 100 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 101 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xTensor[xUbOffset].GetPhyAddr(); | 101 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xTensor[xUbOffset].GetPhyAddr(); |
| 102 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[onceNumPerCore * i].GetPhyAddr(); | 102 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[onceNumPerCore * i].GetPhyAddr(); |
| 103 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[onceNumPerCore * i].GetPhyAddr(); | 103 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[onceNumPerCore * i].GetPhyAddr(); |
| 104 | if (i > 1) { | 104 | if (i > 1) { |
| 105 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); | 105 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); |
| 106 | } | 106 | } |
| @@ -143,14 +143,12 @@ private: | |||
| 143 | for (int64_t i = 0; i < numPerCoreProcess; i++) { | 143 | for (int64_t i = 0; i < numPerCoreProcess; i++) { |
| 144 | uint64_t gammaOffset = ((blockIdx * tiling->numPerCore + numPerCoreoffset + i) % numGroups) * shapeD; | 144 | uint64_t gammaOffset = ((blockIdx * tiling->numPerCore + numPerCoreoffset + i) % numGroups) * shapeD; |
| 145 | uint64_t betaOffset = gammaOffset; | 145 | uint64_t betaOffset = gammaOffset; |
| 146 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xTensor[xUbOffset + i * elemNumAlign].GetPhyAddr(); | 146 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xTensor[xUbOffset + i * elemNumAlign].GetPhyAddr(); |
| 147 | - __local_mem__ T1* yOutLocal = (__local_mem__ T1*)yTensor[outputUbOffset + i * elemNumAlign].GetPhyAddr(); | 147 | + __ubuf__ T1* yOutLocal = (__ubuf__ T1*)yTensor[outputUbOffset + i * elemNumAlign].GetPhyAddr(); |
| 148 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[numPerCoreLoop * onceNumPerCore + i] | 148 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[numPerCoreLoop * onceNumPerCore + i].GetPhyAddr(); |
| 149 | - .GetPhyAddr(); | 149 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[numPerCoreLoop * onceNumPerCore + i].GetPhyAddr(); |
| 150 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[numPerCoreLoop * onceNumPerCore + i] | 150 | + __ubuf__ T2* gammaLocal = (__ubuf__ T2*)gammaTensor.GetPhyAddr(); |
| 151 | - .GetPhyAddr(); | 151 | + __ubuf__ T2* betaLocal = (__ubuf__ T2*)betaTensor.GetPhyAddr(); |
| 152 | - __local_mem__ T2* gammaLocal = (__local_mem__ T2*)gammaTensor.GetPhyAddr(); | ||
| 153 | - __local_mem__ T2* betaLocal = (__local_mem__ T2*)betaTensor.GetPhyAddr(); | ||
| 154 | if (i > 0) { | 152 | if (i > 0) { |
| 155 | WaitFlag<HardEvent::V_MTE2>(eventIDVToMte2); | 153 | WaitFlag<HardEvent::V_MTE2>(eventIDVToMte2); |
| 156 | } | 154 | } |
| @@ -86,11 +86,11 @@ private: | |||
| 86 | 86 | ||
| 87 | __aicore__ inline void CalMeanAndRstdByWelford(uint64_t curNumPerCore, uint64_t curInnerNumPerCore) | 87 | __aicore__ inline void CalMeanAndRstdByWelford(uint64_t curNumPerCore, uint64_t curInnerNumPerCore) |
| 88 | { | 88 | { |
| 89 | - __local_mem__ float* tmpMeanLocal = (__local_mem__ float*)tMeanTensor.GetPhyAddr(); | 89 | + __ubuf__ float* tmpMeanLocal = (__ubuf__ float*)tMeanTensor.GetPhyAddr(); |
| 90 | - __local_mem__ float* tmpVarLocal = (__local_mem__ float*)tVarTensor.GetPhyAddr(); | 90 | + __ubuf__ float* tmpVarLocal = (__ubuf__ float*)tVarTensor.GetPhyAddr(); |
| 91 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 91 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 92 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 92 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 93 | - __local_mem__ float* dichotomyAddLocal = (__local_mem__ float*)dichotomyAddTensor.GetPhyAddr(); | 93 | + __ubuf__ float* dichotomyAddLocal = (__ubuf__ float*)dichotomyAddTensor.GetPhyAddr(); |
| 94 | uint64_t xGmOffset = blockIdx * tiling->numPerCore * elemNum; | 94 | uint64_t xGmOffset = blockIdx * tiling->numPerCore * elemNum; |
| 95 | uint32_t welfordLen = parallelN; | 95 | uint32_t welfordLen = parallelN; |
| 96 | count = 0; | 96 | count = 0; |
| @@ -111,7 +111,7 @@ private: | |||
| 111 | welfordLen); | 111 | welfordLen); |
| 112 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 112 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 113 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 113 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 114 | - __local_mem__ T1* x1Local = (__local_mem__ T1*)xPhase1Tensor[xPhase1Offset].GetPhyAddr(); | 114 | + __ubuf__ T1* x1Local = (__ubuf__ T1*)xPhase1Tensor[xPhase1Offset].GetPhyAddr(); |
| 115 | count = count + 1; | 115 | count = count + 1; |
| 116 | float scale = (float)1.0 / static_cast<float>(count); | 116 | float scale = (float)1.0 / static_cast<float>(count); |
| 117 | VFWelfordParallelUpdate<T1>(x1Local, tmpMeanLocal, tmpVarLocal, i, welfordLen, scale); | 117 | VFWelfordParallelUpdate<T1>(x1Local, tmpMeanLocal, tmpVarLocal, i, welfordLen, scale); |
| @@ -194,13 +194,12 @@ private: | |||
| 194 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 194 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 195 | uint64_t gammaOffset = gammaBaseOffset + i * (processSize / hwNumAlign); | 195 | uint64_t gammaOffset = gammaBaseOffset + i * (processSize / hwNumAlign); |
| 196 | uint64_t betaOffset = gammaOffset; | 196 | uint64_t betaOffset = gammaOffset; |
| 197 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); | 197 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); |
| 198 | - __local_mem__ T2* gammaLocal = hasGamma ? (__local_mem__ T2*)gammaTensor[gammaOffset].GetPhyAddr() : | 198 | + __ubuf__ T2* gammaLocal = hasGamma ? (__ubuf__ T2*)gammaTensor[gammaOffset].GetPhyAddr() : nullptr; |
| 199 | - nullptr; | 199 | + __ubuf__ T2* betaLocal = hasBeta ? (__ubuf__ T2*)betaTensor[betaOffset].GetPhyAddr() : nullptr; |
| 200 | - __local_mem__ T2* betaLocal = hasBeta ? (__local_mem__ T2*)betaTensor[betaOffset].GetPhyAddr() : nullptr; | 200 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); |
| 201 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); | 201 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); |
| 202 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); | 202 | + __ubuf__ T1* yOutLocal = (__ubuf__ T1*)yTensor[inputUbOffset].GetPhyAddr(); |
| 203 | - __local_mem__ T1* yOutLocal = (__local_mem__ T1*)yTensor[inputUbOffset].GetPhyAddr(); | ||
| 204 | if (i > 1) { | 203 | if (i > 1) { |
| 205 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); | 204 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); |
| 206 | } | 205 | } |
| @@ -264,14 +263,12 @@ private: | |||
| 264 | CopyX2UB(xGm[inputOffset], xPhase2Tensor[inputUbOffset], 1, copyLen); | 263 | CopyX2UB(xGm[inputOffset], xPhase2Tensor[inputUbOffset], 1, copyLen); |
| 265 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 264 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 266 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 265 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 267 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); | 266 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); |
| 268 | - __local_mem__ T2* gammaLocal = hasGamma ? (__local_mem__ T2*)gammaTensor[gammaOffset].GetPhyAddr() : | 267 | + __ubuf__ T2* gammaLocal = hasGamma ? (__ubuf__ T2*)gammaTensor[gammaOffset].GetPhyAddr() : nullptr; |
| 269 | - nullptr; | 268 | + __ubuf__ T2* betaLocal = hasBeta ? (__ubuf__ T2*)betaTensor[betaOffset].GetPhyAddr() : nullptr; |
| 270 | - __local_mem__ T2* betaLocal = hasBeta ? (__local_mem__ T2*)betaTensor[betaOffset].GetPhyAddr() : | 269 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); |
| 271 | - nullptr; | 270 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); |
| 272 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); | 271 | + __ubuf__ T1* yOutLocal = (__ubuf__ T1*)yTensor[inputUbOffset].GetPhyAddr(); |
| 273 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); | ||
| 274 | - __local_mem__ T1* yOutLocal = (__local_mem__ T1*)yTensor[inputUbOffset].GetPhyAddr(); | ||
| 275 | if (extent > 1) { | 272 | if (extent > 1) { |
| 276 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); | 273 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); |
| 277 | } | 274 | } |
| @@ -86,11 +86,11 @@ private: | |||
| 86 | 86 | ||
| 87 | __aicore__ inline void CalMeanAndRstdByWelford(uint64_t curNumPerCore, uint64_t curInnerNumPerCore) | 87 | __aicore__ inline void CalMeanAndRstdByWelford(uint64_t curNumPerCore, uint64_t curInnerNumPerCore) |
| 88 | { | 88 | { |
| 89 | - __local_mem__ float* tmpMeanLocal = (__local_mem__ float*)tMeanTensor.GetPhyAddr(); | 89 | + __ubuf__ float* tmpMeanLocal = (__ubuf__ float*)tMeanTensor.GetPhyAddr(); |
| 90 | - __local_mem__ float* tmpVarLocal = (__local_mem__ float*)tVarTensor.GetPhyAddr(); | 90 | + __ubuf__ float* tmpVarLocal = (__ubuf__ float*)tVarTensor.GetPhyAddr(); |
| 91 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 91 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 92 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 92 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 93 | - __local_mem__ float* dichotomyAddLocal = (__local_mem__ float*)dichotomyAddTensor.GetPhyAddr(); | 93 | + __ubuf__ float* dichotomyAddLocal = (__ubuf__ float*)dichotomyAddTensor.GetPhyAddr(); |
| 94 | uint64_t xGmOffset = blockIdx * tiling->numPerCore * elemNum; | 94 | uint64_t xGmOffset = blockIdx * tiling->numPerCore * elemNum; |
| 95 | uint32_t welfordLen = parallelN; | 95 | uint32_t welfordLen = parallelN; |
| 96 | count = 0; | 96 | count = 0; |
| @@ -111,7 +111,7 @@ private: | |||
| 111 | welfordLen); | 111 | welfordLen); |
| 112 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 112 | SetFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 113 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 113 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 114 | - __local_mem__ T1* x1Local = (__local_mem__ T1*)xPhase1Tensor[xPhase1Offset].GetPhyAddr(); | 114 | + __ubuf__ T1* x1Local = (__ubuf__ T1*)xPhase1Tensor[xPhase1Offset].GetPhyAddr(); |
| 115 | count = count + 1; | 115 | count = count + 1; |
| 116 | float scale = (float)1.0 / static_cast<float>(count); | 116 | float scale = (float)1.0 / static_cast<float>(count); |
| 117 | VFWelfordParallelUpdate<T1>(x1Local, tmpMeanLocal, tmpVarLocal, i, welfordLen, scale); | 117 | VFWelfordParallelUpdate<T1>(x1Local, tmpMeanLocal, tmpVarLocal, i, welfordLen, scale); |
| @@ -163,10 +163,10 @@ private: | |||
| 163 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 163 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 164 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 164 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 165 | 165 | ||
| 166 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); | 166 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); |
| 167 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); | 167 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); |
| 168 | - __local_mem__ T2* gammaLocal = (__local_mem__ T2*)gammaTensor.GetPhyAddr(); | 168 | + __ubuf__ T2* gammaLocal = (__ubuf__ T2*)gammaTensor.GetPhyAddr(); |
| 169 | - __local_mem__ T2* betaLocal = (__local_mem__ T2*)betaTensor.GetPhyAddr(); | 169 | + __ubuf__ T2* betaLocal = (__ubuf__ T2*)betaTensor.GetPhyAddr(); |
| 170 | for (int64_t i = 0; i < loopNum; i++) { | 170 | for (int64_t i = 0; i < loopNum; i++) { |
| 171 | uint64_t inputGmOffset = inputBaseOffset + hwNum * rowsCount * i + elemNum * curNumPerCore; | 171 | uint64_t inputGmOffset = inputBaseOffset + hwNum * rowsCount * i + elemNum * curNumPerCore; |
| 172 | bool isPing = (i % BUFFER_NUM) == 0; | 172 | bool isPing = (i % BUFFER_NUM) == 0; |
| @@ -190,8 +190,8 @@ private: | |||
| 190 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); | 190 | WaitFlag<HardEvent::MTE2_V>(isPing ? eventIDMte2ToVPing : eventIDMte2ToVPong); |
| 191 | SetFlag<HardEvent::MTE2_V>(eventIDMte2ToV); | 191 | SetFlag<HardEvent::MTE2_V>(eventIDMte2ToV); |
| 192 | WaitFlag<HardEvent::MTE2_V>(eventIDMte2ToV); | 192 | WaitFlag<HardEvent::MTE2_V>(eventIDMte2ToV); |
| 193 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); | 193 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); |
| 194 | - __local_mem__ T1* yOutLocal = (__local_mem__ T1*)yTensor[inputUbOffset].GetPhyAddr(); | 194 | + __ubuf__ T1* yOutLocal = (__ubuf__ T1*)yTensor[inputUbOffset].GetPhyAddr(); |
| 195 | VFNormalizeAlign<T1, T2>(xLocal, gammaLocal, betaLocal, meanLocal, rstdLocal, yOutLocal, rowsCount, | 195 | VFNormalizeAlign<T1, T2>(xLocal, gammaLocal, betaLocal, meanLocal, rstdLocal, yOutLocal, rowsCount, |
| 196 | reduceCount); | 196 | reduceCount); |
| 197 | SetFlag<HardEvent::V_MTE3>(isPing ? eventIDVToMte3Ping : eventIDVToMte3Pong); | 197 | SetFlag<HardEvent::V_MTE3>(isPing ? eventIDVToMte3Ping : eventIDVToMte3Pong); |
| @@ -241,10 +241,10 @@ private: | |||
| 241 | auto eventIDVToMte2Pong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>()); | 241 | auto eventIDVToMte2Pong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::V_MTE2>()); |
| 242 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 242 | auto eventIDMte3ToVPing = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 243 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); | 243 | auto eventIDMte3ToVPong = static_cast<event_t>(GetTPipePtr()->AllocEventID<HardEvent::MTE3_V>()); |
| 244 | - __local_mem__ float* meanLocal = (__local_mem__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); | 244 | + __ubuf__ float* meanLocal = (__ubuf__ float*)meanTensor[curInnerNumPerCore].GetPhyAddr(); |
| 245 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); | 245 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor[curInnerNumPerCore].GetPhyAddr(); |
| 246 | - __local_mem__ T2* gammaLocal = (__local_mem__ T2*)gammaTensor.GetPhyAddr(); | 246 | + __ubuf__ T2* gammaLocal = (__ubuf__ T2*)gammaTensor.GetPhyAddr(); |
| 247 | - __local_mem__ T2* betaLocal = (__local_mem__ T2*)betaTensor.GetPhyAddr(); | 247 | + __ubuf__ T2* betaLocal = (__ubuf__ T2*)betaTensor.GetPhyAddr(); |
| 248 | for (int64_t i = 0; i < loopNum; i++) { // for D | 248 | for (int64_t i = 0; i < loopNum; i++) { // for D |
| 249 | int64_t copyLen = totalSize; | 249 | int64_t copyLen = totalSize; |
| 250 | uint64_t gammaOffset = gammaBaseOffset + i; | 250 | uint64_t gammaOffset = gammaBaseOffset + i; |
| @@ -256,8 +256,8 @@ private: | |||
| 256 | 256 | ||
| 257 | SetFlag<HardEvent::MTE2_V>(eventIDMte2ToV); | 257 | SetFlag<HardEvent::MTE2_V>(eventIDMte2ToV); |
| 258 | WaitFlag<HardEvent::MTE2_V>(eventIDMte2ToV); | 258 | WaitFlag<HardEvent::MTE2_V>(eventIDMte2ToV); |
| 259 | - __local_mem__ T2* gammaLocal = (__local_mem__ T2*)gammaTensor.GetPhyAddr(); | 259 | + __ubuf__ T2* gammaLocal = (__ubuf__ T2*)gammaTensor.GetPhyAddr(); |
| 260 | - __local_mem__ T2* betaLocal = (__local_mem__ T2*)betaTensor.GetPhyAddr(); | 260 | + __ubuf__ T2* betaLocal = (__ubuf__ T2*)betaTensor.GetPhyAddr(); |
| 261 | for (int64_t j = 0; j < innerLoopNum; j++) { // for HW | 261 | for (int64_t j = 0; j < innerLoopNum; j++) { // for HW |
| 262 | int64_t inputGmOffset = inputBaseOffset + totalSize * j + hwNum * i + elemNum * curNumPerCore; | 262 | int64_t inputGmOffset = inputBaseOffset + totalSize * j + hwNum * i + elemNum * curNumPerCore; |
| 263 | auto extent = i * innerLoopNum + j; | 263 | auto extent = i * innerLoopNum + j; |
| @@ -275,8 +275,8 @@ private: | |||
| 275 | if (extent > 1) { | 275 | if (extent > 1) { |
| 276 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); | 276 | WaitFlag<HardEvent::MTE3_V>(isPing ? eventIDMte3ToVPing : eventIDMte3ToVPong); |
| 277 | } | 277 | } |
| 278 | - __local_mem__ T1* xLocal = (__local_mem__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); | 278 | + __ubuf__ T1* xLocal = (__ubuf__ T1*)xPhase2Tensor[inputUbOffset].GetPhyAddr(); |
| 279 | - __local_mem__ T1* yOutLocal = (__local_mem__ T1*)yTensor[inputUbOffset].GetPhyAddr(); | 279 | + __ubuf__ T1* yOutLocal = (__ubuf__ T1*)yTensor[inputUbOffset].GetPhyAddr(); |
| 280 | int32_t reduceCount = copyLen; | 280 | int32_t reduceCount = copyLen; |
| 281 | VFNormalizeAlign<T1, T2>(xLocal, gammaLocal, betaLocal, meanLocal, rstdLocal, yOutLocal, 1, copyLen); | 281 | VFNormalizeAlign<T1, T2>(xLocal, gammaLocal, betaLocal, meanLocal, rstdLocal, yOutLocal, 1, copyLen); |
| 282 | SetFlag<HardEvent::V_MTE3>(isPing ? eventIDVToMte3Ping : eventIDVToMte3Pong); | 282 | SetFlag<HardEvent::V_MTE3>(isPing ? eventIDVToMte3Ping : eventIDVToMte3Pong); |
| @@ -9,5 +9,5 @@ | |||
| 9 | # ----------------------------------------------------------------------------------------------------------- | 9 | # ----------------------------------------------------------------------------------------------------------- |
| 10 | 10 | ||
| 11 | if ((UT_TEST_ALL OR OP_KERNEL_UT) AND NOT UT_DONE) | 11 | if ((UT_TEST_ALL OR OP_KERNEL_UT) AND NOT UT_DONE) |
| 12 | - AddOpTestCase(group_norm_v2 "ascend950pr_9599" "-DDTYPE_X=float") | 12 | + AddOpTestCase(group_norm_v2 "ascend950pr_9599" "-DDTYPE_X=float" norm_common) |
| 13 | -endif() | 13 | +endif() |
| @@ -29,6 +29,9 @@ using AscendC::MicroAPI::MemType; | |||
| 29 | using AscendC::MicroAPI::RegTensor; | 29 | using AscendC::MicroAPI::RegTensor; |
| 30 | using AscendC::MicroAPI::StoreDist; | 30 | using AscendC::MicroAPI::StoreDist; |
| 31 | using AscendC::MicroAPI::UpdateMask; | 31 | using AscendC::MicroAPI::UpdateMask; |
| 32 | +using AscendC::Reg::LoadAlign; | ||
| 33 | +using AscendC::Reg::Reduce; | ||
| 34 | +using AscendC::Reg::StoreAlign; | ||
| 32 | 35 | ||
| 33 | constexpr uint64_t ALIGN_32_FACTOR = 32; | 36 | constexpr uint64_t ALIGN_32_FACTOR = 32; |
| 34 | constexpr uint32_t NUM_ONE = 1; | 37 | constexpr uint32_t NUM_ONE = 1; |
| @@ -215,12 +218,12 @@ private: | |||
| 215 | LocalTensor<float>& meanFp32Local, LocalTensor<float>& rstdLocal, | 218 | LocalTensor<float>& meanFp32Local, LocalTensor<float>& rstdLocal, |
| 216 | uint32_t curRows, uint32_t numColAlign, uint32_t reduceNum) | 219 | uint32_t curRows, uint32_t numColAlign, uint32_t reduceNum) |
| 217 | { | 220 | { |
| 218 | - __local_mem__ T_X* xInUb = (__local_mem__ T_X*)xLocal.GetPhyAddr(); | 221 | + __ubuf__ T_X* xInUb = (__ubuf__ T_X*)xLocal.GetPhyAddr(); |
| 219 | - __local_mem__ T_X* yInUb = (__local_mem__ T_X*)yLocal.GetPhyAddr(); | 222 | + __ubuf__ T_X* yInUb = (__ubuf__ T_X*)yLocal.GetPhyAddr(); |
| 220 | - __local_mem__ float* meanFp32Ub = (__local_mem__ float*)meanFp32Local.GetPhyAddr(); | 223 | + __ubuf__ float* meanFp32Ub = (__ubuf__ float*)meanFp32Local.GetPhyAddr(); |
| 221 | - __local_mem__ float* rstdInUb = (__local_mem__ float*)rstdLocal.GetPhyAddr(); | 224 | + __ubuf__ float* rstdInUb = (__ubuf__ float*)rstdLocal.GetPhyAddr(); |
| 222 | - __local_mem__ T_BETA* gammaInUb = (__local_mem__ T_BETA*)gammaLocal.GetPhyAddr(); | 225 | + __ubuf__ T_BETA* gammaInUb = (__ubuf__ T_BETA*)gammaLocal.GetPhyAddr(); |
| 223 | - __local_mem__ T_BETA* betaInUb = (__local_mem__ T_BETA*)betaLocal.GetPhyAddr(); | 226 | + __ubuf__ T_BETA* betaInUb = (__ubuf__ T_BETA*)betaLocal.GetPhyAddr(); |
| 224 | 227 | ||
| 225 | uint16_t loopRows = static_cast<uint16_t>(curRows); | 228 | uint16_t loopRows = static_cast<uint16_t>(curRows); |
| 226 | uint16_t loopCols = static_cast<uint16_t>((reduceNum + VL_FP32 - 1) / VL_FP32); | 229 | uint16_t loopCols = static_cast<uint16_t>((reduceNum + VL_FP32 - 1) / VL_FP32); |
| @@ -260,13 +263,13 @@ private: | |||
| 260 | uint32_t numColAlign, uint32_t reduceNum, float avgFactor) | 263 | uint32_t numColAlign, uint32_t reduceNum, float avgFactor) |
| 261 | { | 264 | { |
| 262 | LocalTensor<float> binaryAddBuffTmp = binaryAddBuf_.Get<float>(); | 265 | LocalTensor<float> binaryAddBuffTmp = binaryAddBuf_.Get<float>(); |
| 263 | - __local_mem__ T_X* xInUb = (__local_mem__ T_X*)xLocal.GetPhyAddr(); | 266 | + __ubuf__ T_X* xInUb = (__ubuf__ T_X*)xLocal.GetPhyAddr(); |
| 264 | - __local_mem__ T_MEAN* meanUb = (__local_mem__ T_MEAN*)meanLocal.GetPhyAddr(); | 267 | + __ubuf__ T_MEAN* meanUb = (__ubuf__ T_MEAN*)meanLocal.GetPhyAddr(); |
| 265 | - __local_mem__ T_MEAN* varUb = (__local_mem__ T_MEAN*)varLocal.GetPhyAddr(); | 268 | + __ubuf__ T_MEAN* varUb = (__ubuf__ T_MEAN*)varLocal.GetPhyAddr(); |
| 266 | 269 | ||
| 267 | - __local_mem__ float* meanFp32Ub = (__local_mem__ float*)meanFp32Local.GetPhyAddr(); | 270 | + __ubuf__ float* meanFp32Ub = (__ubuf__ float*)meanFp32Local.GetPhyAddr(); |
| 268 | - __local_mem__ float* varFp32Ub = (__local_mem__ float*)varFp32Local.GetPhyAddr(); | 271 | + __ubuf__ float* varFp32Ub = (__ubuf__ float*)varFp32Local.GetPhyAddr(); |
| 269 | - __local_mem__ float* tmpUb = (__local_mem__ float*)binaryAddBuffTmp.GetPhyAddr(); // 二分累加 tmp buffer | 272 | + __ubuf__ float* tmpUb = (__ubuf__ float*)binaryAddBuffTmp.GetPhyAddr(); // 二分累加 tmp buffer |
| 270 | 273 | ||
| 271 | if (reduceNum <= VL_FP32) { | 274 | if (reduceNum <= VL_FP32) { |
| 272 | CalculateMeanLessThanVL(xInUb, meanUb, meanFp32Ub, curRows, numColAlign, reduceNum, avgFactor); | 275 | CalculateMeanLessThanVL(xInUb, meanUb, meanFp32Ub, curRows, numColAlign, reduceNum, avgFactor); |
| @@ -290,9 +293,9 @@ private: | |||
| 290 | } | 293 | } |
| 291 | 294 | ||
| 292 | // LessThanVL | 295 | // LessThanVL |
| 293 | - __aicore__ inline void CalculateMeanLessThanVL(__local_mem__ T_X* xInUb, __local_mem__ T_MEAN* meanUb, | 296 | + __aicore__ inline void CalculateMeanLessThanVL(__ubuf__ T_X* xInUb, __ubuf__ T_MEAN* meanUb, |
| 294 | - __local_mem__ float* meanFp32Ub, uint16_t curRows, | 297 | + __ubuf__ float* meanFp32Ub, uint16_t curRows, uint32_t numColAlign, |
| 295 | - uint32_t numColAlign, uint32_t reduceNum, float avgFactor) | 298 | + uint32_t reduceNum, float avgFactor) |
| 296 | { | 299 | { |
| 297 | __VEC_SCOPE__ | 300 | __VEC_SCOPE__ |
| 298 | { | 301 | { |
| @@ -305,18 +308,18 @@ private: | |||
| 305 | 308 | ||
| 306 | for (uint16_t i = 0; i < curRows; i++) { | 309 | for (uint16_t i = 0; i < curRows; i++) { |
| 307 | LoadTensorForDtypeTIn<T_X>(xInUb, x, pregLoop, i * numColAlign); | 310 | LoadTensorForDtypeTIn<T_X>(xInUb, x, pregLoop, i * numColAlign); |
| 308 | - ReduceSum(mean, x, pregLoop); | 311 | + Reduce<ReduceType::SUM>(mean, x, pregLoop); |
| 309 | Muls(mean, mean, avgFactor, pregOne); | 312 | Muls(mean, mean, avgFactor, pregOne); |
| 310 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, mean, pregOne, i); | 313 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, mean, pregOne, i); |
| 311 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, mean, pregOne); | 314 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, mean, pregOne); |
| 312 | } | 315 | } |
| 313 | } | 316 | } |
| 314 | } | 317 | } |
| 315 | 318 | ||
| 316 | - __aicore__ inline void CalculateSquareReduceSumLessThanVL(__local_mem__ T_X* xInUb, __local_mem__ float* meanFp32Ub, | 319 | + __aicore__ inline void CalculateSquareReduceSumLessThanVL(__ubuf__ T_X* xInUb, __ubuf__ float* meanFp32Ub, |
| 317 | - __local_mem__ T_MEAN* varUb, | 320 | + __ubuf__ T_MEAN* varUb, __ubuf__ float* varFp32Ub, |
| 318 | - __local_mem__ float* varFp32Ub, uint16_t curRows, | 321 | + uint16_t curRows, uint32_t numColAlign, |
| 319 | - uint32_t numColAlign, uint32_t reduceNum, float avgFactor) | 322 | + uint32_t reduceNum, float avgFactor) |
| 320 | { | 323 | { |
| 321 | __VEC_SCOPE__ | 324 | __VEC_SCOPE__ |
| 322 | { | 325 | { |
| @@ -333,17 +336,17 @@ private: | |||
| 333 | LoadScalarForDtypeTIn<float>(meanFp32Ub, mean, pregLoop, i); | 336 | LoadScalarForDtypeTIn<float>(meanFp32Ub, mean, pregLoop, i); |
| 334 | Sub(x, x, mean, pregLoop); | 337 | Sub(x, x, mean, pregLoop); |
| 335 | Mul(x, x, x, pregLoop); | 338 | Mul(x, x, x, pregLoop); |
| 336 | - ReduceSum(vMean, x, pregLoop); | 339 | + Reduce<ReduceType::SUM>(vMean, x, pregLoop); |
| 337 | Muls(vMean, vMean, avgFactor, pregOne); | 340 | Muls(vMean, vMean, avgFactor, pregOne); |
| 338 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); | 341 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); |
| 339 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); | 342 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); |
| 340 | } | 343 | } |
| 341 | } | 344 | } |
| 342 | } | 345 | } |
| 343 | 346 | ||
| 344 | // LessThanTwoVL | 347 | // LessThanTwoVL |
| 345 | - __aicore__ inline void CalculateMeanLessThanTwoVL(__local_mem__ T_X* xInUb, __local_mem__ T_MEAN* meanUb, | 348 | + __aicore__ inline void CalculateMeanLessThanTwoVL(__ubuf__ T_X* xInUb, __ubuf__ T_MEAN* meanUb, |
| 346 | - __local_mem__ float* meanFp32Ub, uint16_t curRows, | 349 | + __ubuf__ float* meanFp32Ub, uint16_t curRows, |
| 347 | uint32_t numColAlign, uint32_t reduceNum, float avgFactor) | 350 | uint32_t numColAlign, uint32_t reduceNum, float avgFactor) |
| 348 | { | 351 | { |
| 349 | uint32_t tailLen = reduceNum - VL_FP32; | 352 | uint32_t tailLen = reduceNum - VL_FP32; |
| @@ -362,17 +365,18 @@ private: | |||
| 362 | LoadTensorForDtypeTIn<T_X>(xInUb + VL_FP32, xFold, pregTail, i * numColAlign); | 365 | LoadTensorForDtypeTIn<T_X>(xInUb + VL_FP32, xFold, pregTail, i * numColAlign); |
| 363 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregTail); | 366 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), pregTail); |
| 364 | Add(x, x, xFold, pregFull); | 367 | Add(x, x, xFold, pregFull); |
| 365 | - ReduceSum(mean, x, pregFull); | 368 | + Reduce<ReduceType::SUM>(mean, x, pregFull); |
| 366 | Muls(mean, mean, avgFactor, pregOne); | 369 | Muls(mean, mean, avgFactor, pregOne); |
| 367 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, mean, pregOne, i); | 370 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, mean, pregOne, i); |
| 368 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, mean, pregOne); | 371 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, mean, pregOne); |
| 369 | } | 372 | } |
| 370 | } | 373 | } |
| 371 | } | 374 | } |
| 372 | 375 | ||
| 373 | - __aicore__ inline void CalculateSquareReduceSumLessThanTwoVL( | 376 | + __aicore__ inline void CalculateSquareReduceSumLessThanTwoVL(__ubuf__ T_X* xInUb, __ubuf__ float* meanFp32Ub, |
| 374 | - __local_mem__ T_X* xInUb, __local_mem__ float* meanFp32Ub, __local_mem__ T_MEAN* varUb, | 377 | + __ubuf__ T_MEAN* varUb, __ubuf__ float* varFp32Ub, |
| 375 | - __local_mem__ float* varFp32Ub, uint16_t curRows, uint32_t numColAlign, uint32_t reduceNum, float avgFactor) | 378 | + uint16_t curRows, uint32_t numColAlign, |
| 379 | + uint32_t reduceNum, float avgFactor) | ||
| 376 | { | 380 | { |
| 377 | uint32_t tailLen = reduceNum - VL_FP32; | 381 | uint32_t tailLen = reduceNum - VL_FP32; |
| 378 | __VEC_SCOPE__ | 382 | __VEC_SCOPE__ |
| @@ -396,20 +400,19 @@ private: | |||
| 396 | Mul(x, x, x, pregFull); | 400 | Mul(x, x, x, pregFull); |
| 397 | Mul(xFold, xFold, xFold, pregTail); | 401 | Mul(xFold, xFold, xFold, pregTail); |
| 398 | Add(x, x, xFold, pregFull); | 402 | Add(x, x, xFold, pregFull); |
| 399 | - ReduceSum(vMean, x, pregFull); | 403 | + Reduce<ReduceType::SUM>(vMean, x, pregFull); |
| 400 | Muls(vMean, vMean, avgFactor, pregOne); | 404 | Muls(vMean, vMean, avgFactor, pregOne); |
| 401 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); | 405 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); |
| 402 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); | 406 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); |
| 403 | } | 407 | } |
| 404 | } | 408 | } |
| 405 | } | 409 | } |
| 406 | 410 | ||
| 407 | // | 411 | // |
| 408 | template <int32_t LAST_LOOP_NUMS> | 412 | template <int32_t LAST_LOOP_NUMS> |
| 409 | - __aicore__ inline void CalculateMeanSumCommon(__local_mem__ T_X* xInUb, __local_mem__ T_MEAN* meanUb, | 413 | + __aicore__ inline void CalculateMeanSumCommon(__ubuf__ T_X* xInUb, __ubuf__ T_MEAN* meanUb, |
| 410 | - __local_mem__ float* meanFp32Ub, __local_mem__ float* tmpUb, | 414 | + __ubuf__ float* meanFp32Ub, __ubuf__ float* tmpUb, uint16_t curRows, |
| 411 | - uint16_t curRows, uint32_t numColAlign, uint32_t reduceNum, | 415 | + uint32_t numColAlign, uint32_t reduceNum, float avgFactor) |
| 412 | - float avgFactor) | ||
| 413 | { | 416 | { |
| 414 | uint32_t binaryAddQuotient = binaryAddQuotient_; | 417 | uint32_t binaryAddQuotient = binaryAddQuotient_; |
| 415 | uint16_t binaryAddQuotientLoop = (binaryAddQuotient + VL_FP32 - 1) / VL_FP32; | 418 | uint16_t binaryAddQuotientLoop = (binaryAddQuotient + VL_FP32 - 1) / VL_FP32; |
| @@ -438,8 +441,8 @@ private: | |||
| 438 | LoadTensorForDtypeTIn<T_X>(xInUb, x, pregFull, offset); | 441 | LoadTensorForDtypeTIn<T_X>(xInUb, x, pregFull, offset); |
| 439 | LoadTensorForDtypeTIn<T_X>(xInUb + binaryAddQuotient, xFold, pregFull, offset); | 442 | LoadTensorForDtypeTIn<T_X>(xInUb + binaryAddQuotient, xFold, pregFull, offset); |
| 440 | Add(x, x, xFold, pregFull); | 443 | Add(x, x, xFold, pregFull); |
| 441 | - ReduceSum(mean, x, pregFull); | 444 | + Reduce<ReduceType::SUM>(mean, x, pregFull); |
| 442 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 445 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 443 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + r), mean, pregOne); | 446 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + r), mean, pregOne); |
| 444 | } | 447 | } |
| 445 | uint32_t sregRemainder = binaryAddRemainder - binaryAddRemainderFloorLoop * VL_FP32; | 448 | uint32_t sregRemainder = binaryAddRemainder - binaryAddRemainderFloorLoop * VL_FP32; |
| @@ -453,8 +456,8 @@ private: | |||
| 453 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), | 456 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), |
| 454 | pregLoop); | 457 | pregLoop); |
| 455 | Add(x, x, xFold, pregFull); | 458 | Add(x, x, xFold, pregFull); |
| 456 | - ReduceSum(mean, x, pregFull); | 459 | + Reduce<ReduceType::SUM>(mean, x, pregFull); |
| 457 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 460 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 458 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, | 461 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, |
| 459 | pregOne); | 462 | pregOne); |
| 460 | } | 463 | } |
| @@ -462,8 +465,8 @@ private: | |||
| 462 | r++) { | 465 | r++) { |
| 463 | uint32_t offset = r * VL_FP32 + baseOffset; | 466 | uint32_t offset = r * VL_FP32 + baseOffset; |
| 464 | LoadTensorForDtypeTIn<T_X>(xInUb + binaryAddRemainderCeilLoop * VL_FP32, x, pregFull, offset); | 467 | LoadTensorForDtypeTIn<T_X>(xInUb + binaryAddRemainderCeilLoop * VL_FP32, x, pregFull, offset); |
| 465 | - ReduceSum(mean, x, pregFull); | 468 | + Reduce<ReduceType::SUM>(mean, x, pregFull); |
| 466 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 469 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 467 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, | 470 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, |
| 468 | pregOne); | 471 | pregOne); |
| 469 | } | 472 | } |
| @@ -472,34 +475,34 @@ private: | |||
| 472 | if constexpr (LAST_LOOP_NUMS == 1) { | 475 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 473 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); | 476 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); |
| 474 | for (uint16_t i = 0; i < curRows; ++i) { | 477 | for (uint16_t i = 0; i < curRows; ++i) { |
| 475 | - DataCopy(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); | 478 | + LoadAlign(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); |
| 476 | - ReduceSum(vMean, x, pregLast); | 479 | + Reduce<ReduceType::SUM>(vMean, x, pregLast); |
| 477 | Muls(vMean, vMean, avgFactor, pregOne); | 480 | Muls(vMean, vMean, avgFactor, pregOne); |
| 478 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, vMean, pregOne, i); | 481 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, vMean, pregOne, i); |
| 479 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, vMean, pregOne); | 482 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, vMean, pregOne); |
| 480 | } | 483 | } |
| 481 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 484 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 482 | lastBinaryAddNum -= VL_FP32; | 485 | lastBinaryAddNum -= VL_FP32; |
| 483 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); | 486 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); |
| 484 | for (uint16_t i = 0; i < curRows; ++i) { | 487 | for (uint16_t i = 0; i < curRows; ++i) { |
| 485 | - DataCopy(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); | 488 | + LoadAlign(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); |
| 486 | - DataCopy(xFold, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + VL_FP32)); | 489 | + LoadAlign(xFold, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + VL_FP32)); |
| 487 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), | 490 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), |
| 488 | pregLast); | 491 | pregLast); |
| 489 | Add(sumReg, x, xFold, pregFull); | 492 | Add(sumReg, x, xFold, pregFull); |
| 490 | - ReduceSum(vMean, sumReg, pregFull); | 493 | + Reduce<ReduceType::SUM>(vMean, sumReg, pregFull); |
| 491 | Muls(vMean, vMean, avgFactor, pregOne); | 494 | Muls(vMean, vMean, avgFactor, pregOne); |
| 492 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, vMean, pregOne, i); | 495 | StoreOneElementForDtypeTOut<T_MEAN>(meanUb, vMean, pregOne, i); |
| 493 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, vMean, pregOne); | 496 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanFp32Ub + i, vMean, pregOne); |
| 494 | } | 497 | } |
| 495 | } | 498 | } |
| 496 | } | 499 | } |
| 497 | } | 500 | } |
| 498 | 501 | ||
| 499 | template <int32_t LAST_LOOP_NUMS> | 502 | template <int32_t LAST_LOOP_NUMS> |
| 500 | - __aicore__ inline void CalculateSquareReduceSumCommon(__local_mem__ T_X* xInUb, __local_mem__ float* meanFp32Ub, | 503 | + __aicore__ inline void CalculateSquareReduceSumCommon(__ubuf__ T_X* xInUb, __ubuf__ float* meanFp32Ub, |
| 501 | - __local_mem__ T_MEAN* varUb, __local_mem__ float* tmpUb, | 504 | + __ubuf__ T_MEAN* varUb, __ubuf__ float* tmpUb, |
| 502 | - __local_mem__ float* varFp32Ub, uint16_t curRows, | 505 | + __ubuf__ float* varFp32Ub, uint16_t curRows, |
| 503 | uint32_t numColAlign, uint32_t reduceNum, float avgFactor) | 506 | uint32_t numColAlign, uint32_t reduceNum, float avgFactor) |
| 504 | { | 507 | { |
| 505 | uint32_t binaryAddQuotient = binaryAddQuotient_; | 508 | uint32_t binaryAddQuotient = binaryAddQuotient_; |
| @@ -535,8 +538,8 @@ private: | |||
| 535 | Mul(x, x, x, pregFull); | 538 | Mul(x, x, x, pregFull); |
| 536 | Mul(xFold, xFold, xFold, pregFull); | 539 | Mul(xFold, xFold, xFold, pregFull); |
| 537 | Add(sumReg, x, xFold, pregFull); | 540 | Add(sumReg, x, xFold, pregFull); |
| 538 | - ReduceSum(vMean, sumReg, pregFull); | 541 | + Reduce<ReduceType::SUM>(vMean, sumReg, pregFull); |
| 539 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 542 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 540 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + r), vMean, pregOne); | 543 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + r), vMean, pregOne); |
| 541 | } | 544 | } |
| 542 | uint32_t sregRemainder = binaryAddRemainder - binaryAddRemainderFloorLoop * VL_FP32; | 545 | uint32_t sregRemainder = binaryAddRemainder - binaryAddRemainderFloorLoop * VL_FP32; |
| @@ -554,8 +557,8 @@ private: | |||
| 554 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), | 557 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), |
| 555 | pregLoop); | 558 | pregLoop); |
| 556 | Add(sumReg, x, xFold, pregFull); | 559 | Add(sumReg, x, xFold, pregFull); |
| 557 | - ReduceSum(vMean, sumReg, pregFull); | 560 | + Reduce<ReduceType::SUM>(vMean, sumReg, pregFull); |
| 558 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 561 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 559 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), vMean, | 562 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), vMean, |
| 560 | pregOne); | 563 | pregOne); |
| 561 | } | 564 | } |
| @@ -565,8 +568,8 @@ private: | |||
| 565 | LoadTensorForDtypeTIn<T_X>(xInUb + binaryAddRemainderCeilLoop * VL_FP32, x, pregFull, offset); | 568 | LoadTensorForDtypeTIn<T_X>(xInUb + binaryAddRemainderCeilLoop * VL_FP32, x, pregFull, offset); |
| 566 | Sub(x, x, mean, pregFull); | 569 | Sub(x, x, mean, pregFull); |
| 567 | Mul(x, x, x, pregFull); | 570 | Mul(x, x, x, pregFull); |
| 568 | - ReduceSum(vMean, x, pregFull); | 571 | + Reduce<ReduceType::SUM>(vMean, x, pregFull); |
| 569 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 572 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 570 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), | 573 | tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), |
| 571 | vMean, pregOne); | 574 | vMean, pregOne); |
| 572 | } | 575 | } |
| @@ -575,25 +578,25 @@ private: | |||
| 575 | if constexpr (LAST_LOOP_NUMS == 1) { | 578 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 576 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); | 579 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); |
| 577 | for (uint16_t i = 0; i < curRows; ++i) { | 580 | for (uint16_t i = 0; i < curRows; ++i) { |
| 578 | - DataCopy(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); | 581 | + LoadAlign(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); |
| 579 | - ReduceSum(vMean, x, pregLast); | 582 | + Reduce<ReduceType::SUM>(vMean, x, pregLast); |
| 580 | Muls(vMean, vMean, avgFactor, pregOne); | 583 | Muls(vMean, vMean, avgFactor, pregOne); |
| 581 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); | 584 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); |
| 582 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); | 585 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); |
| 583 | } | 586 | } |
| 584 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 587 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 585 | lastBinaryAddNum -= VL_FP32; | 588 | lastBinaryAddNum -= VL_FP32; |
| 586 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); | 589 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); |
| 587 | for (uint16_t i = 0; i < curRows; ++i) { | 590 | for (uint16_t i = 0; i < curRows; ++i) { |
| 588 | - DataCopy(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); | 591 | + LoadAlign(x, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign)); |
| 589 | - DataCopy(xFold, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + VL_FP32)); | 592 | + LoadAlign(xFold, tmpUb + static_cast<uint32_t>(i * lastBinaryAddNumAlign + VL_FP32)); |
| 590 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), | 593 | ShiftLefts((RegTensor<uint32_t>&)xFold, (RegTensor<uint32_t>&)xFold, static_cast<int16_t>(0), |
| 591 | pregLast); | 594 | pregLast); |
| 592 | Add(sumReg, x, xFold, pregFull); | 595 | Add(sumReg, x, xFold, pregFull); |
| 593 | - ReduceSum(vMean, sumReg, pregFull); | 596 | + Reduce<ReduceType::SUM>(vMean, sumReg, pregFull); |
| 594 | Muls(vMean, vMean, avgFactor, pregOne); | 597 | Muls(vMean, vMean, avgFactor, pregOne); |
| 595 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); | 598 | StoreOneElementForDtypeTOut<T_MEAN>(varUb, vMean, pregOne, i); |
| 596 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); | 599 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(varFp32Ub + i, vMean, pregOne); |
| 597 | } | 600 | } |
| 598 | } | 601 | } |
| 599 | } | 602 | } |
| @@ -608,8 +611,8 @@ private: | |||
| 608 | static constexpr float SCALAR3 = 0.5; | 611 | static constexpr float SCALAR3 = 0.5; |
| 609 | static constexpr float SCALAR0 = -99.99; | 612 | static constexpr float SCALAR0 = -99.99; |
| 610 | 613 | ||
| 611 | - __local_mem__ float* rstdInUb = (__local_mem__ float*)rstdLocal.GetPhyAddr(); | 614 | + __ubuf__ float* rstdInUb = (__ubuf__ float*)rstdLocal.GetPhyAddr(); |
| 612 | - __local_mem__ float* varUb = (__local_mem__ float*)varLocal.GetPhyAddr(); | 615 | + __ubuf__ float* varUb = (__ubuf__ float*)varLocal.GetPhyAddr(); |
| 613 | uint16_t loopRows = static_cast<uint16_t>((curRows + VL_FP32 - 1) / VL_FP32); | 616 | uint16_t loopRows = static_cast<uint16_t>((curRows + VL_FP32 - 1) / VL_FP32); |
| 614 | __VEC_SCOPE__ | 617 | __VEC_SCOPE__ |
| 615 | { | 618 | { |
| @@ -642,7 +645,7 @@ private: | |||
| 642 | Duplicate(t1, SCALAR2, pregLoop); | 645 | Duplicate(t1, SCALAR2, pregLoop); |
| 643 | Duplicate(s, float(1.0), pregLoop); | 646 | Duplicate(s, float(1.0), pregLoop); |
| 644 | // rstd | 647 | // rstd |
| 645 | - DataCopy(var, varUb + i * VL_FP32); | 648 | + LoadAlign(var, varUb + i * VL_FP32); |
| 646 | Adds(var, var, epsilon, pregLoop); | 649 | Adds(var, var, epsilon, pregLoop); |
| 647 | Maxs(var, var, SCALAR0, pregLoop); | 650 | Maxs(var, var, SCALAR0, pregLoop); |
| 648 | Div(r, one, var, pregLoop); | 651 | Div(r, one, var, pregLoop); |
| @@ -658,11 +661,11 @@ private: | |||
| 658 | Mula(s, var, r, pregLoop); // s + x * t | 661 | Mula(s, var, r, pregLoop); // s + x * t |
| 659 | Mul(s, s, rstd, pregLoop); // e * y | 662 | Mul(s, s, rstd, pregLoop); // e * y |
| 660 | Mula(rstd, s, scalar1, pregLoop); // y + y * e * 0.5 | 663 | Mula(rstd, s, scalar1, pregLoop); // y + y * e * 0.5 |
| 661 | - CompareScalar(cmpRegZero, var, POS_INF, pregLoop); | 664 | + Compares(cmpRegZero, var, POS_INF, pregLoop); |
| 662 | Select(rstd, scalarZero, rstd, cmpRegZero); | 665 | Select(rstd, scalarZero, rstd, cmpRegZero); |
| 663 | - CompareScalar(cmpRegInf, var, float(0.0), pregLoop); | 666 | + Compares(cmpRegInf, var, float(0.0), pregLoop); |
| 664 | Select(rstd, scalarInf, rstd, cmpRegInf); | 667 | Select(rstd, scalarInf, rstd, cmpRegInf); |
| 665 | - DataCopy(rstdInUb + i * VL_FP32, rstd, pregLoop); | 668 | + StoreAlign(rstdInUb + i * VL_FP32, rstd, pregLoop); |
| 666 | } | 669 | } |
| 667 | } | 670 | } |
| 668 | } | 671 | } |
| @@ -30,6 +30,8 @@ using AscendC::MicroAPI::MemType; | |||
| 30 | using AscendC::MicroAPI::RegTensor; | 30 | using AscendC::MicroAPI::RegTensor; |
| 31 | using AscendC::MicroAPI::StoreDist; | 31 | using AscendC::MicroAPI::StoreDist; |
| 32 | using AscendC::MicroAPI::UpdateMask; | 32 | using AscendC::MicroAPI::UpdateMask; |
| 33 | +using AscendC::Reg::LoadAlign; | ||
| 34 | +using AscendC::Reg::StoreAlign; | ||
| 33 | 35 | ||
| 34 | template <typename T> | 36 | template <typename T> |
| 35 | __aicore__ inline void CopyIn(const LocalTensor<T>& dstTensor, const GlobalTensor<T>& srcTensor, const int64_t rowSize) | 37 | __aicore__ inline void CopyIn(const LocalTensor<T>& dstTensor, const GlobalTensor<T>& srcTensor, const int64_t rowSize) |
| @@ -197,7 +199,7 @@ public: | |||
| 197 | private: | 199 | private: |
| 198 | __aicore__ inline void CaculateCountBuf(LocalTensor<int32_t>& tCountTensor) | 200 | __aicore__ inline void CaculateCountBuf(LocalTensor<int32_t>& tCountTensor) |
| 199 | { | 201 | { |
| 200 | - __local_mem__ int32_t* tmpCountLocal = (__local_mem__ int32_t*)tCountTensor.GetPhyAddr(); | 202 | + __ubuf__ int32_t* tmpCountLocal = (__ubuf__ int32_t*)tCountTensor.GetPhyAddr(); |
| 201 | 203 | ||
| 202 | int64_t tailcoreProcessNum = welfordUpdateTimes_; | 204 | int64_t tailcoreProcessNum = welfordUpdateTimes_; |
| 203 | int64_t formercoreProcessNum = welfordUpdateTimes_; | 205 | int64_t formercoreProcessNum = welfordUpdateTimes_; |
| @@ -226,13 +228,13 @@ private: | |||
| 226 | Duplicate(tmpCount, tailCoreAddCount, pregMain); | 228 | Duplicate(tmpCount, tailCoreAddCount, pregMain); |
| 227 | for (uint16_t i = 0; i < fisrstLoopCount; i++) { | 229 | for (uint16_t i = 0; i < fisrstLoopCount; i++) { |
| 228 | pregLoop = AscendC::MicroAPI::UpdateMask<int32_t>(sreg3); | 230 | pregLoop = AscendC::MicroAPI::UpdateMask<int32_t>(sreg3); |
| 229 | - DataCopy(((__local_mem__ int32_t*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); | 231 | + StoreAlign(((__ubuf__ int32_t*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); |
| 230 | } | 232 | } |
| 231 | uint32_t sreg4 = secondNum; | 233 | uint32_t sreg4 = secondNum; |
| 232 | Duplicate(tmpCount, formerCoreAddCount, pregMain); | 234 | Duplicate(tmpCount, formerCoreAddCount, pregMain); |
| 233 | for (uint16_t i = 0; i < secondLoopCount; i++) { | 235 | for (uint16_t i = 0; i < secondLoopCount; i++) { |
| 234 | pregLoop = AscendC::MicroAPI::UpdateMask<int32_t>(sreg4); | 236 | pregLoop = AscendC::MicroAPI::UpdateMask<int32_t>(sreg4); |
| 235 | - DataCopy(((__local_mem__ int32_t*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); | 237 | + StoreAlign(((__ubuf__ int32_t*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); |
| 236 | } | 238 | } |
| 237 | } | 239 | } |
| 238 | } | 240 | } |
| @@ -246,10 +248,10 @@ private: | |||
| 246 | 248 | ||
| 247 | __aicore__ inline void CastBatchMeanVariance(uint64_t currentANum) | 249 | __aicore__ inline void CastBatchMeanVariance(uint64_t currentANum) |
| 248 | { | 250 | { |
| 249 | - __local_mem__ float* batchMeanInAddr = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 251 | + __ubuf__ float* batchMeanInAddr = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 250 | - __local_mem__ float* batchVarianceInAddr = (__local_mem__ float*)varianceTensor.GetPhyAddr(); | 252 | + __ubuf__ float* batchVarianceInAddr = (__ubuf__ float*)varianceTensor.GetPhyAddr(); |
| 251 | - __local_mem__ T_MEAN* batchMeanOutAddr = (__local_mem__ T_MEAN*)meanTensor.GetPhyAddr(); | 253 | + __ubuf__ T_MEAN* batchMeanOutAddr = (__ubuf__ T_MEAN*)meanTensor.GetPhyAddr(); |
| 252 | - __local_mem__ T_MEAN* batchVarianceOutAddr = (__local_mem__ T_MEAN*)varianceTensor.GetPhyAddr(); | 254 | + __ubuf__ T_MEAN* batchVarianceOutAddr = (__ubuf__ T_MEAN*)varianceTensor.GetPhyAddr(); |
| 253 | 255 | ||
| 254 | uint32_t castCount = static_cast<uint32_t>(currentANum); | 256 | uint32_t castCount = static_cast<uint32_t>(currentANum); |
| 255 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); | 257 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); |
| @@ -262,15 +264,15 @@ private: | |||
| 262 | MicroAPI::MaskReg pregLoop; | 264 | MicroAPI::MaskReg pregLoop; |
| 263 | for (uint16_t i = 0; i < castLoops; i++) { | 265 | for (uint16_t i = 0; i < castLoops; i++) { |
| 264 | pregLoop = MicroAPI::UpdateMask<float>(castCount); | 266 | pregLoop = MicroAPI::UpdateMask<float>(castCount); |
| 265 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); | 267 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); |
| 266 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_variance, | 268 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_variance, |
| 267 | - batchVarianceInAddr + VL_F32 * i); | 269 | + batchVarianceInAddr + VL_F32 * i); |
| 268 | Cast<T_MEAN, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 270 | Cast<T_MEAN, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 269 | Cast<T_MEAN, float, castTraitB322B16>(output_variance, input_variance, pregLoop); | 271 | Cast<T_MEAN, float, castTraitB322B16>(output_variance, input_variance, pregLoop); |
| 270 | - DataCopy<T_MEAN, StoreDist::DIST_PACK_B32>(((__local_mem__ T_MEAN*)batchMeanOutAddr + i * VL_MEAN), | 272 | + StoreAlign<T_MEAN, StoreDist::DIST_PACK_B32>(((__ubuf__ T_MEAN*)batchMeanOutAddr + i * VL_MEAN), |
| 271 | - output_mean, pregLoop); | 273 | + output_mean, pregLoop); |
| 272 | - DataCopy<T_MEAN, StoreDist::DIST_PACK_B32>(((__local_mem__ T_MEAN*)batchVarianceOutAddr + i * VL_MEAN), | 274 | + StoreAlign<T_MEAN, StoreDist::DIST_PACK_B32>(((__ubuf__ T_MEAN*)batchVarianceOutAddr + i * VL_MEAN), |
| 273 | - output_variance, pregLoop); | 275 | + output_variance, pregLoop); |
| 274 | } | 276 | } |
| 275 | } | 277 | } |
| 276 | } | 278 | } |
| @@ -336,16 +338,16 @@ private: | |||
| 336 | uint16_t loopTimes = (elemCnt + VL_B32 - 1) / VL_B32; | 338 | uint16_t loopTimes = (elemCnt + VL_B32 - 1) / VL_B32; |
| 337 | __VEC_SCOPE__ | 339 | __VEC_SCOPE__ |
| 338 | { | 340 | { |
| 339 | - __local_mem__ float* meamPtr = (__local_mem__ float*)mean.GetPhyAddr(); | 341 | + __ubuf__ float* meamPtr = (__ubuf__ float*)mean.GetPhyAddr(); |
| 340 | - __local_mem__ float* variancePtr = (__local_mem__ float*)variance.GetPhyAddr(); | 342 | + __ubuf__ float* variancePtr = (__ubuf__ float*)variance.GetPhyAddr(); |
| 341 | uint32_t count = static_cast<uint32_t>(elemCnt); | 343 | uint32_t count = static_cast<uint32_t>(elemCnt); |
| 342 | AscendC::MicroAPI::RegTensor<float> xReg; | 344 | AscendC::MicroAPI::RegTensor<float> xReg; |
| 343 | AscendC::MicroAPI::MaskReg pMask; | 345 | AscendC::MicroAPI::MaskReg pMask; |
| 344 | Duplicate(xReg, 0.0f); | 346 | Duplicate(xReg, 0.0f); |
| 345 | for (uint16_t i = 0; i < loopTimes; ++i) { | 347 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 346 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 348 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 347 | - DataCopy((__local_mem__ float*)meamPtr + i * VL_B32, xReg, pMask); | 349 | + StoreAlign((__ubuf__ float*)meamPtr + i * VL_B32, xReg, pMask); |
| 348 | - DataCopy((__local_mem__ float*)variancePtr + i * VL_B32, xReg, pMask); | 350 | + StoreAlign((__ubuf__ float*)variancePtr + i * VL_B32, xReg, pMask); |
| 349 | } | 351 | } |
| 350 | } | 352 | } |
| 351 | } | 353 | } |
| @@ -371,8 +373,8 @@ private: | |||
| 371 | 373 | ||
| 372 | __aicore__ inline void CalculateBatchRstd(const int64_t offset) | 374 | __aicore__ inline void CalculateBatchRstd(const int64_t offset) |
| 373 | { | 375 | { |
| 374 | - __local_mem__ float* batchRstdTensorAddr = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 376 | + __ubuf__ float* batchRstdTensorAddr = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 375 | - __local_mem__ float* batchVarTensorAddr = (__local_mem__ float*)varianceTensor.GetPhyAddr(); | 377 | + __ubuf__ float* batchVarTensorAddr = (__ubuf__ float*)varianceTensor.GetPhyAddr(); |
| 376 | __VEC_SCOPE__ | 378 | __VEC_SCOPE__ |
| 377 | { | 379 | { |
| 378 | RegTensor<float> var; | 380 | RegTensor<float> var; |
| @@ -394,7 +396,7 @@ private: | |||
| 394 | MaskReg pregOne = CreateMask<float, MaskPattern::VL1>(); | 396 | MaskReg pregOne = CreateMask<float, MaskPattern::VL1>(); |
| 395 | Duplicate(one, 1.0, pregOne); | 397 | Duplicate(one, 1.0, pregOne); |
| 396 | 398 | ||
| 397 | - DataCopy<float, LoadDist::DIST_BRC_B32>(var, ((__local_mem__ float*)batchVarTensorAddr + offset)); | 399 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(var, ((__ubuf__ float*)batchVarTensorAddr + offset)); |
| 398 | Duplicate(scalar1, float(0.5), pregOne); | 400 | Duplicate(scalar1, float(0.5), pregOne); |
| 399 | Duplicate(scalarInf, POS_INF, pregOne); | 401 | Duplicate(scalarInf, POS_INF, pregOne); |
| 400 | Duplicate(scalarZero, float(0.0), pregOne); | 402 | Duplicate(scalarZero, float(0.0), pregOne); |
| @@ -414,12 +416,12 @@ private: | |||
| 414 | Mula(s, var, r, pregOne); // s + x * t | 416 | Mula(s, var, r, pregOne); // s + x * t |
| 415 | Mul(s, s, rstd, pregOne); // e * y | 417 | Mul(s, s, rstd, pregOne); // e * y |
| 416 | Mula(rstd, s, scalar1, pregOne); // y + y * e * 0.5 | 418 | Mula(rstd, s, scalar1, pregOne); // y + y * e * 0.5 |
| 417 | - CompareScalar(cmpRegZero, var, POS_INF, pregOne); | 419 | + Compares(cmpRegZero, var, POS_INF, pregOne); |
| 418 | Select(rstd, scalarZero, rstd, cmpRegZero); | 420 | Select(rstd, scalarZero, rstd, cmpRegZero); |
| 419 | - CompareScalar(cmpRegInf, var, float(0.0), pregOne); | 421 | + Compares(cmpRegInf, var, float(0.0), pregOne); |
| 420 | Select(rstd, scalarInf, rstd, cmpRegInf); | 422 | Select(rstd, scalarInf, rstd, cmpRegInf); |
| 421 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(((__local_mem__ float*)batchRstdTensorAddr + offset), | 423 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(((__ubuf__ float*)batchRstdTensorAddr + offset), rstd, |
| 422 | - rstd, pregOne); | 424 | + pregOne); |
| 423 | } | 425 | } |
| 424 | } | 426 | } |
| 425 | 427 | ||
| @@ -437,47 +439,47 @@ private: | |||
| 437 | } | 439 | } |
| 438 | 440 | ||
| 439 | template <typename T_SRC> | 441 | template <typename T_SRC> |
| 440 | - __aicore__ inline void LoadTensorForDtypeT(RegTensor<float>& dst, __local_mem__ T_SRC* input, MaskReg& preg, | 442 | + __aicore__ inline void LoadTensorForDtypeT(RegTensor<float>& dst, __ubuf__ T_SRC* input, MaskReg& preg, |
| 441 | uint32_t offset) | 443 | uint32_t offset) |
| 442 | { | 444 | { |
| 443 | if constexpr (IsSameType<T_SRC, half>::value) { | 445 | if constexpr (IsSameType<T_SRC, half>::value) { |
| 444 | RegTensor<half> xFp16; | 446 | RegTensor<half> xFp16; |
| 445 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ half*)(input) + (offset))); | 447 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ half*)(input) + (offset))); |
| 446 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); | 448 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); |
| 447 | } else if constexpr (IsSameType<T_SRC, bfloat16_t>::value) { | 449 | } else if constexpr (IsSameType<T_SRC, bfloat16_t>::value) { |
| 448 | RegTensor<bfloat16_t> xBf16; | 450 | RegTensor<bfloat16_t> xBf16; |
| 449 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__local_mem__ bfloat16_t*)(input) + (offset))); | 451 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__ubuf__ bfloat16_t*)(input) + (offset))); |
| 450 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); | 452 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); |
| 451 | } else { | 453 | } else { |
| 452 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset))); | 454 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset))); |
| 453 | } | 455 | } |
| 454 | } | 456 | } |
| 455 | 457 | ||
| 456 | template <typename T_SRC_GAMMA> | 458 | template <typename T_SRC_GAMMA> |
| 457 | - __aicore__ inline void LoadOneNumberTensorForDtypeT(RegTensor<float>& dst, __local_mem__ T_SRC_GAMMA* input, | 459 | + __aicore__ inline void LoadOneNumberTensorForDtypeT(RegTensor<float>& dst, __ubuf__ T_SRC_GAMMA* input, |
| 458 | MaskReg& preg, uint32_t offset) | 460 | MaskReg& preg, uint32_t offset) |
| 459 | { | 461 | { |
| 460 | if constexpr (IsSameType<T_SRC_GAMMA, half>::value) { | 462 | if constexpr (IsSameType<T_SRC_GAMMA, half>::value) { |
| 461 | RegTensor<half> xFp16; | 463 | RegTensor<half> xFp16; |
| 462 | - DataCopy<half, LoadDist::DIST_BRC_B16>(xFp16, ((__local_mem__ half*)(input) + (offset))); | 464 | + LoadAlign<half, LoadDist::DIST_BRC_B16>(xFp16, ((__ubuf__ half*)(input) + (offset))); |
| 463 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); | 465 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); |
| 464 | } else if constexpr (IsSameType<T_SRC_GAMMA, bfloat16_t>::value) { | 466 | } else if constexpr (IsSameType<T_SRC_GAMMA, bfloat16_t>::value) { |
| 465 | RegTensor<bfloat16_t> xBf16; | 467 | RegTensor<bfloat16_t> xBf16; |
| 466 | - DataCopy<bfloat16_t, LoadDist::DIST_BRC_B16>(xBf16, ((__local_mem__ bfloat16_t*)(input) + (offset))); | 468 | + LoadAlign<bfloat16_t, LoadDist::DIST_BRC_B16>(xBf16, ((__ubuf__ bfloat16_t*)(input) + (offset))); |
| 467 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); | 469 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); |
| 468 | } else { | 470 | } else { |
| 469 | - DataCopy<float, LoadDist::DIST_BRC_B32>(dst, ((__local_mem__ float*)(input) + (offset))); | 471 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(dst, ((__ubuf__ float*)(input) + (offset))); |
| 470 | } | 472 | } |
| 471 | } | 473 | } |
| 472 | 474 | ||
| 473 | __aicore__ inline void CalY(const int64_t elemCnt, const int64_t gammaUbOffset) | 475 | __aicore__ inline void CalY(const int64_t elemCnt, const int64_t gammaUbOffset) |
| 474 | { | 476 | { |
| 475 | - __local_mem__ float* batchMeanTensorAddr = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 477 | + __ubuf__ float* batchMeanTensorAddr = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 476 | - __local_mem__ float* batchRstdTensorAddr = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 478 | + __ubuf__ float* batchRstdTensorAddr = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 477 | - __local_mem__ T* xTensorAddr = (__local_mem__ T*)xTensor.GetPhyAddr(); | 479 | + __ubuf__ T* xTensorAddr = (__ubuf__ T*)xTensor.GetPhyAddr(); |
| 478 | - __local_mem__ T* yTensorAddr = (__local_mem__ T*)yTensor.GetPhyAddr(); | 480 | + __ubuf__ T* yTensorAddr = (__ubuf__ T*)yTensor.GetPhyAddr(); |
| 479 | - __local_mem__ T_BETA* gammaTensorAddr = (__local_mem__ T_BETA*)gammaTensor.GetPhyAddr(); | 481 | + __ubuf__ T_BETA* gammaTensorAddr = (__ubuf__ T_BETA*)gammaTensor.GetPhyAddr(); |
| 480 | - __local_mem__ T_BETA* betaTensorAddr = (__local_mem__ T_BETA*)betaTensor.GetPhyAddr(); | 482 | + __ubuf__ T_BETA* betaTensorAddr = (__ubuf__ T_BETA*)betaTensor.GetPhyAddr(); |
| 481 | 483 | ||
| 482 | uint16_t numLoop = CEIL_DIV(elemCnt, VL_F32); | 484 | uint16_t numLoop = CEIL_DIV(elemCnt, VL_F32); |
| 483 | __VEC_SCOPE__ | 485 | __VEC_SCOPE__ |
| @@ -488,8 +490,8 @@ private: | |||
| 488 | RegTensor<float> gamma; | 490 | RegTensor<float> gamma; |
| 489 | RegTensor<float> beta; | 491 | RegTensor<float> beta; |
| 490 | RegTensor<float> y; | 492 | RegTensor<float> y; |
| 491 | - DataCopy<float, LoadDist::DIST_BRC_B32>(mean, ((__local_mem__ float*)batchMeanTensorAddr + gammaUbOffset)); | 493 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(mean, ((__ubuf__ float*)batchMeanTensorAddr + gammaUbOffset)); |
| 492 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rstd, ((__local_mem__ float*)batchRstdTensorAddr + gammaUbOffset)); | 494 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rstd, ((__ubuf__ float*)batchRstdTensorAddr + gammaUbOffset)); |
| 493 | 495 | ||
| 494 | MaskReg mask0; | 496 | MaskReg mask0; |
| 495 | uint32_t sreg0 = elemCnt; | 497 | uint32_t sreg0 = elemCnt; |
| @@ -506,13 +508,13 @@ private: | |||
| 506 | if constexpr (IsSameType<T, half>::value) { | 508 | if constexpr (IsSameType<T, half>::value) { |
| 507 | RegTensor<half> yFp16; | 509 | RegTensor<half> yFp16; |
| 508 | Cast<half, float, castTraitB322B16>(yFp16, y, mask0); | 510 | Cast<half, float, castTraitB322B16>(yFp16, y, mask0); |
| 509 | - DataCopy<half, StoreDist::DIST_PACK_B32>(yTensorAddr + i * VL_F32, yFp16, mask0); | 511 | + StoreAlign<half, StoreDist::DIST_PACK_B32>(yTensorAddr + i * VL_F32, yFp16, mask0); |
| 510 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 512 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 511 | RegTensor<bfloat16_t> xBf16; | 513 | RegTensor<bfloat16_t> xBf16; |
| 512 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, y, mask0); | 514 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, y, mask0); |
| 513 | - DataCopy<bfloat16_t, StoreDist::DIST_PACK_B32>(yTensorAddr + i * VL_F32, xBf16, mask0); | 515 | + StoreAlign<bfloat16_t, StoreDist::DIST_PACK_B32>(yTensorAddr + i * VL_F32, xBf16, mask0); |
| 514 | } else { | 516 | } else { |
| 515 | - DataCopy(yTensorAddr + i * VL_F32, y, mask0); | 517 | + StoreAlign(yTensorAddr + i * VL_F32, y, mask0); |
| 516 | } | 518 | } |
| 517 | } | 519 | } |
| 518 | } | 520 | } |
| @@ -145,14 +145,14 @@ public: | |||
| 145 | LocalTensor<T_X> yOutUb = yQueue_.AllocTensor<T_X>(); | 145 | LocalTensor<T_X> yOutUb = yQueue_.AllocTensor<T_X>(); |
| 146 | LocalTensor<float> meanOutUb = meanQueue_.AllocTensor<float>(); | 146 | LocalTensor<float> meanOutUb = meanQueue_.AllocTensor<float>(); |
| 147 | LocalTensor<float> varianceOutUb = varianceQueue_.AllocTensor<float>(); | 147 | LocalTensor<float> varianceOutUb = varianceQueue_.AllocTensor<float>(); |
| 148 | - __local_mem__ T_X* xInUbAddr = (__local_mem__ T_X*)xInUb.GetPhyAddr(); | 148 | + __ubuf__ T_X* xInUbAddr = (__ubuf__ T_X*)xInUb.GetPhyAddr(); |
| 149 | - __local_mem__ float* xFp32InUbAddr = (__local_mem__ float*)xInUbAddr; | 149 | + __ubuf__ float* xFp32InUbAddr = (__ubuf__ float*)xInUbAddr; |
| 150 | - __local_mem__ T_X* yOutUbAddr = (__local_mem__ T_X*)yOutUb.GetPhyAddr(); | 150 | + __ubuf__ T_X* yOutUbAddr = (__ubuf__ T_X*)yOutUb.GetPhyAddr(); |
| 151 | - __local_mem__ float* meanOutUbAddr = (__local_mem__ float*)meanOutUb.GetPhyAddr(); | 151 | + __ubuf__ float* meanOutUbAddr = (__ubuf__ float*)meanOutUb.GetPhyAddr(); |
| 152 | - __local_mem__ float* varianceOutUbAddr = (__local_mem__ float*)varianceOutUb.GetPhyAddr(); | 152 | + __ubuf__ float* varianceOutUbAddr = (__ubuf__ float*)varianceOutUb.GetPhyAddr(); |
| 153 | if constexpr (IsSameType<T_X, half>::value || IsSameType<T_X, bfloat16_t>::value) { | 153 | if constexpr (IsSameType<T_X, half>::value || IsSameType<T_X, bfloat16_t>::value) { |
| 154 | LocalTensor<float> castInUb = castBuf_.Get<float>(); | 154 | LocalTensor<float> castInUb = castBuf_.Get<float>(); |
| 155 | - xFp32InUbAddr = (__local_mem__ float*)castInUb.GetPhyAddr(); | 155 | + xFp32InUbAddr = (__ubuf__ float*)castInUb.GetPhyAddr(); |
| 156 | CastToFp32(xInUbAddr, xFp32InUbAddr, curTileA0Len); | 156 | CastToFp32(xInUbAddr, xFp32InUbAddr, curTileA0Len); |
| 157 | CalculateMean(xFp32InUbAddr, yOutUbAddr, meanOutUbAddr, curTileA0Len); | 157 | CalculateMean(xFp32InUbAddr, yOutUbAddr, meanOutUbAddr, curTileA0Len); |
| 158 | CalculateVar(xFp32InUbAddr, yOutUbAddr, meanOutUbAddr, varianceOutUbAddr, curTileA0Len); | 158 | CalculateVar(xFp32InUbAddr, yOutUbAddr, meanOutUbAddr, varianceOutUbAddr, curTileA0Len); |
| @@ -167,14 +167,14 @@ public: | |||
| 167 | varianceQueue_.DeQue<float>(); | 167 | varianceQueue_.DeQue<float>(); |
| 168 | 168 | ||
| 169 | LocalTensor<float> rstdUb = rstdBuf_.Get<float>(); | 169 | LocalTensor<float> rstdUb = rstdBuf_.Get<float>(); |
| 170 | - __local_mem__ float* rstdUbAddr = (__local_mem__ float*)rstdUb.GetPhyAddr(); | 170 | + __ubuf__ float* rstdUbAddr = (__ubuf__ float*)rstdUb.GetPhyAddr(); |
| 171 | ComputeRstd(varianceOutUbAddr, rstdUbAddr, curTileA0Len, epsilon_); | 171 | ComputeRstd(varianceOutUbAddr, rstdUbAddr, curTileA0Len, epsilon_); |
| 172 | 172 | ||
| 173 | CopyInGammaBeta(GammaBetaOffset, curTileA0Len); | 173 | CopyInGammaBeta(GammaBetaOffset, curTileA0Len); |
| 174 | LocalTensor<T_B> betaInUb = betaQueue_.template DeQue<T_B>(); | 174 | LocalTensor<T_B> betaInUb = betaQueue_.template DeQue<T_B>(); |
| 175 | LocalTensor<T_B> gammaInUb = gammaQueue_.template DeQue<T_B>(); | 175 | LocalTensor<T_B> gammaInUb = gammaQueue_.template DeQue<T_B>(); |
| 176 | - __local_mem__ T_B* betaInUbAddr = (__local_mem__ T_B*)betaInUb.GetPhyAddr(); | 176 | + __ubuf__ T_B* betaInUbAddr = (__ubuf__ T_B*)betaInUb.GetPhyAddr(); |
| 177 | - __local_mem__ T_B* gammaInUbAddr = (__local_mem__ T_B*)gammaInUb.GetPhyAddr(); | 177 | + __ubuf__ T_B* gammaInUbAddr = (__ubuf__ T_B*)gammaInUb.GetPhyAddr(); |
| 178 | 178 | ||
| 179 | CalculateY(xFp32InUbAddr, yOutUbAddr, betaInUbAddr, gammaInUbAddr, meanOutUbAddr, rstdUbAddr, curTileA0Len); | 179 | CalculateY(xFp32InUbAddr, yOutUbAddr, betaInUbAddr, gammaInUbAddr, meanOutUbAddr, rstdUbAddr, curTileA0Len); |
| 180 | if constexpr (!IsSameType<T_M, float>::value) { | 180 | if constexpr (!IsSameType<T_M, float>::value) { |
| @@ -196,7 +196,7 @@ public: | |||
| 196 | return 0; | 196 | return 0; |
| 197 | } | 197 | } |
| 198 | 198 | ||
| 199 | - __aicore__ inline void CastToFp32(__local_mem__ T_X* xInUb, __local_mem__ float* castInUb, uint32_t curTileA0Len) | 199 | + __aicore__ inline void CastToFp32(__ubuf__ T_X* xInUb, __ubuf__ float* castInUb, uint32_t curTileA0Len) |
| 200 | { | 200 | { |
| 201 | uint32_t tileA0Len = tileA0Len_; | 201 | uint32_t tileA0Len = tileA0Len_; |
| 202 | uint16_t curTileRLen = totalRLen_; | 202 | uint16_t curTileRLen = totalRLen_; |
| @@ -214,7 +214,7 @@ public: | |||
| 214 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; | 214 | uint32_t xOffset = i * tileA0Len + k * VL_FP32; |
| 215 | // load x | 215 | // load x |
| 216 | LoadOneTensorForDtypeT(xInUb, x, pregMask, xOffset); | 216 | LoadOneTensorForDtypeT(xInUb, x, pregMask, xOffset); |
| 217 | - DataCopy(((__local_mem__ float*)castInUb) + xOffset, x, pregMask); | 217 | + StoreAlign(((__ubuf__ float*)castInUb) + xOffset, x, pregMask); |
| 218 | } | 218 | } |
| 219 | } | 219 | } |
| 220 | } | 220 | } |
| @@ -258,81 +258,81 @@ public: | |||
| 258 | } | 258 | } |
| 259 | 259 | ||
| 260 | template <typename T_SRC> | 260 | template <typename T_SRC> |
| 261 | - __aicore__ inline void LoadTwoTensorForDtypeT(__local_mem__ T_SRC* src1, __local_mem__ T_SRC* src2, | 261 | + __aicore__ inline void LoadTwoTensorForDtypeT(__ubuf__ T_SRC* src1, __ubuf__ T_SRC* src2, RegTensor<float>& dst1, |
| 262 | - RegTensor<float>& dst1, RegTensor<float>& dst2, MaskReg& dst1Preg, | 262 | + RegTensor<float>& dst2, MaskReg& dst1Preg, MaskReg& dst2Preg, |
| 263 | - MaskReg& dst2Preg, uint32_t src1Offset, uint32_t src2Offset) | 263 | + uint32_t src1Offset, uint32_t src2Offset) |
| 264 | { | 264 | { |
| 265 | if constexpr (IsSameType<T_SRC, half>::value) { | 265 | if constexpr (IsSameType<T_SRC, half>::value) { |
| 266 | RegTensor<half> xFp16Q; | 266 | RegTensor<half> xFp16Q; |
| 267 | RegTensor<half> xFp16R; | 267 | RegTensor<half> xFp16R; |
| 268 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ half*)(src1) + (src1Offset))); | 268 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__ubuf__ half*)(src1) + (src1Offset))); |
| 269 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ half*)(src2) + (src2Offset))); | 269 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__ubuf__ half*)(src2) + (src2Offset))); |
| 270 | Cast<float, half, castTraitB162B32>(dst1, xFp16Q, dst1Preg); | 270 | Cast<float, half, castTraitB162B32>(dst1, xFp16Q, dst1Preg); |
| 271 | Cast<float, half, castTraitB162B32>(dst2, xFp16R, dst2Preg); | 271 | Cast<float, half, castTraitB162B32>(dst2, xFp16R, dst2Preg); |
| 272 | } else if constexpr (IsSameType<T_SRC, bfloat16_t>::value) { | 272 | } else if constexpr (IsSameType<T_SRC, bfloat16_t>::value) { |
| 273 | RegTensor<bfloat16_t> xFp16Q; | 273 | RegTensor<bfloat16_t> xFp16Q; |
| 274 | RegTensor<bfloat16_t> xFp16R; | 274 | RegTensor<bfloat16_t> xFp16R; |
| 275 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__local_mem__ bfloat16_t*)(src1) + (src1Offset))); | 275 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16Q, ((__ubuf__ bfloat16_t*)(src1) + (src1Offset))); |
| 276 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__local_mem__ bfloat16_t*)(src2) + (src2Offset))); | 276 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xFp16R, ((__ubuf__ bfloat16_t*)(src2) + (src2Offset))); |
| 277 | Cast<float, bfloat16_t, castTraitB162B32>(dst1, xFp16Q, dst1Preg); | 277 | Cast<float, bfloat16_t, castTraitB162B32>(dst1, xFp16Q, dst1Preg); |
| 278 | Cast<float, bfloat16_t, castTraitB162B32>(dst2, xFp16R, dst2Preg); | 278 | Cast<float, bfloat16_t, castTraitB162B32>(dst2, xFp16R, dst2Preg); |
| 279 | } else { | 279 | } else { |
| 280 | - DataCopy(dst1, ((__local_mem__ float*)(src1) + (src1Offset))); | 280 | + LoadAlign(dst1, ((__ubuf__ float*)(src1) + (src1Offset))); |
| 281 | - DataCopy(dst2, ((__local_mem__ float*)(src2) + (src2Offset))); | 281 | + LoadAlign(dst2, ((__ubuf__ float*)(src2) + (src2Offset))); |
| 282 | } | 282 | } |
| 283 | } | 283 | } |
| 284 | 284 | ||
| 285 | - __aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T_X* input, RegTensor<float>& dst, MaskReg& preg, | 285 | + __aicore__ inline void LoadOneTensorForDtypeT(__ubuf__ T_X* input, RegTensor<float>& dst, MaskReg& preg, |
| 286 | uint32_t offset) | 286 | uint32_t offset) |
| 287 | { | 287 | { |
| 288 | if constexpr (IsSameType<T_X, half>::value) { | 288 | if constexpr (IsSameType<T_X, half>::value) { |
| 289 | RegTensor<half> xFp16; | 289 | RegTensor<half> xFp16; |
| 290 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ half*)(input) + (offset))); | 290 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ half*)(input) + (offset))); |
| 291 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); | 291 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); |
| 292 | } else if constexpr (IsSameType<T_X, bfloat16_t>::value) { | 292 | } else if constexpr (IsSameType<T_X, bfloat16_t>::value) { |
| 293 | RegTensor<bfloat16_t> xBf16; | 293 | RegTensor<bfloat16_t> xBf16; |
| 294 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__local_mem__ bfloat16_t*)(input) + (offset))); | 294 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__ubuf__ bfloat16_t*)(input) + (offset))); |
| 295 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); | 295 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); |
| 296 | } else { | 296 | } else { |
| 297 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset))); | 297 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset))); |
| 298 | } | 298 | } |
| 299 | } | 299 | } |
| 300 | 300 | ||
| 301 | - __aicore__ inline void TwoRowAddForMeanWithTail(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, | 301 | + __aicore__ inline void TwoRowAddForMeanWithTail(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 302 | uint32_t offset1, uint32_t offset2, uint32_t offset3, | 302 | uint32_t offset1, uint32_t offset2, uint32_t offset3, |
| 303 | uint32_t offset4, RegTensor<float>& rem, RegTensor<float>& nextRow, | 303 | uint32_t offset4, RegTensor<float>& rem, RegTensor<float>& nextRow, |
| 304 | RegTensor<float>& remNextRow, float n) | 304 | RegTensor<float>& remNextRow, float n) |
| 305 | { | 305 | { |
| 306 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 306 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 307 | - DataCopy(rem, ((__local_mem__ float*)(input) + (offset2))); | 307 | + LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2))); |
| 308 | Muls(dst, dst, n, preg); | 308 | Muls(dst, dst, n, preg); |
| 309 | Muls(rem, rem, n, preg); | 309 | Muls(rem, rem, n, preg); |
| 310 | Add(dst, dst, rem, preg); | 310 | Add(dst, dst, rem, preg); |
| 311 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset3))); | 311 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset3))); |
| 312 | - DataCopy(remNextRow, ((__local_mem__ float*)(input) + (offset4))); | 312 | + LoadAlign(remNextRow, ((__ubuf__ float*)(input) + (offset4))); |
| 313 | Muls(nextRow, nextRow, n, preg); | 313 | Muls(nextRow, nextRow, n, preg); |
| 314 | Muls(remNextRow, remNextRow, n, preg); | 314 | Muls(remNextRow, remNextRow, n, preg); |
| 315 | Add(nextRow, nextRow, remNextRow, preg); | 315 | Add(nextRow, nextRow, remNextRow, preg); |
| 316 | Add(dst, dst, nextRow, preg); | 316 | Add(dst, dst, nextRow, preg); |
| 317 | } | 317 | } |
| 318 | 318 | ||
| 319 | - __aicore__ inline void TwoRowAddForMean(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, | 319 | + __aicore__ inline void TwoRowAddForMean(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 320 | uint32_t offset1, uint32_t offset2, RegTensor<float>& nextRow, float n) | 320 | uint32_t offset1, uint32_t offset2, RegTensor<float>& nextRow, float n) |
| 321 | { | 321 | { |
| 322 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 322 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 323 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset2))); | 323 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset2))); |
| 324 | Muls(dst, dst, n, preg); | 324 | Muls(dst, dst, n, preg); |
| 325 | Muls(nextRow, nextRow, n, preg); | 325 | Muls(nextRow, nextRow, n, preg); |
| 326 | Add(dst, dst, nextRow, preg); | 326 | Add(dst, dst, nextRow, preg); |
| 327 | } | 327 | } |
| 328 | 328 | ||
| 329 | - __aicore__ inline void TwoRowAddForVarWithTail(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, | 329 | + __aicore__ inline void TwoRowAddForVarWithTail(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 330 | uint32_t offset1, uint32_t offset2, uint32_t offset3, | 330 | uint32_t offset1, uint32_t offset2, uint32_t offset3, |
| 331 | uint32_t offset4, RegTensor<float>& mean, RegTensor<float>& rem, | 331 | uint32_t offset4, RegTensor<float>& mean, RegTensor<float>& rem, |
| 332 | RegTensor<float>& nextRow, RegTensor<float>& remNextRow, float n) | 332 | RegTensor<float>& nextRow, RegTensor<float>& remNextRow, float n) |
| 333 | { | 333 | { |
| 334 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 334 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 335 | - DataCopy(rem, ((__local_mem__ float*)(input) + (offset2))); | 335 | + LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2))); |
| 336 | Sub(dst, dst, mean, preg); | 336 | Sub(dst, dst, mean, preg); |
| 337 | Sub(rem, rem, mean, preg); | 337 | Sub(rem, rem, mean, preg); |
| 338 | Mul(dst, dst, dst, preg); | 338 | Mul(dst, dst, dst, preg); |
| @@ -340,8 +340,8 @@ public: | |||
| 340 | Muls(dst, dst, n, preg); | 340 | Muls(dst, dst, n, preg); |
| 341 | Muls(rem, rem, n, preg); | 341 | Muls(rem, rem, n, preg); |
| 342 | Add(dst, dst, rem, preg); | 342 | Add(dst, dst, rem, preg); |
| 343 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset3))); | 343 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset3))); |
| 344 | - DataCopy(remNextRow, ((__local_mem__ float*)(input) + (offset4))); | 344 | + LoadAlign(remNextRow, ((__ubuf__ float*)(input) + (offset4))); |
| 345 | Sub(nextRow, nextRow, mean, preg); | 345 | Sub(nextRow, nextRow, mean, preg); |
| 346 | Sub(remNextRow, remNextRow, mean, preg); | 346 | Sub(remNextRow, remNextRow, mean, preg); |
| 347 | Mul(nextRow, nextRow, nextRow, preg); | 347 | Mul(nextRow, nextRow, nextRow, preg); |
| @@ -352,12 +352,12 @@ public: | |||
| 352 | Add(dst, dst, nextRow, preg); | 352 | Add(dst, dst, nextRow, preg); |
| 353 | } | 353 | } |
| 354 | 354 | ||
| 355 | - __aicore__ inline void TwoRowAddForVar(RegTensor<float>& dst, __local_mem__ float* input, MaskReg& preg, | 355 | + __aicore__ inline void TwoRowAddForVar(RegTensor<float>& dst, __ubuf__ float* input, MaskReg& preg, |
| 356 | uint32_t offset1, uint32_t offset2, RegTensor<float>& mean, | 356 | uint32_t offset1, uint32_t offset2, RegTensor<float>& mean, |
| 357 | RegTensor<float>& nextRow, float n) | 357 | RegTensor<float>& nextRow, float n) |
| 358 | { | 358 | { |
| 359 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 359 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 360 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset2))); | 360 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset2))); |
| 361 | Sub(dst, dst, mean, preg); | 361 | Sub(dst, dst, mean, preg); |
| 362 | Sub(nextRow, nextRow, mean, preg); | 362 | Sub(nextRow, nextRow, mean, preg); |
| 363 | Mul(dst, dst, dst, preg); | 363 | Mul(dst, dst, dst, preg); |
| @@ -367,7 +367,7 @@ public: | |||
| 367 | Add(dst, dst, nextRow, preg); | 367 | Add(dst, dst, nextRow, preg); |
| 368 | } | 368 | } |
| 369 | 369 | ||
| 370 | - __aicore__ inline void BinaryAddVF(__local_mem__ float* binaryAddTmpAddr, uint32_t rLoopStride, uint32_t offset, | 370 | + __aicore__ inline void BinaryAddVF(__ubuf__ float* binaryAddTmpAddr, uint32_t rLoopStride, uint32_t offset, |
| 371 | uint16_t binaryAddKLoop, uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop, | 371 | uint16_t binaryAddKLoop, uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop, |
| 372 | MaskReg& pregLoop, RegTensor<float>& x1, RegTensor<float>& x2, | 372 | MaskReg& pregLoop, RegTensor<float>& x1, RegTensor<float>& x2, |
| 373 | RegTensor<float>& x3, RegTensor<float>& x4) | 373 | RegTensor<float>& x3, RegTensor<float>& x4) |
| @@ -376,31 +376,30 @@ public: | |||
| 376 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { | 376 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { |
| 377 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET; | 377 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET; |
| 378 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { | 378 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { |
| 379 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset)); | 379 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset)); |
| 380 | - DataCopy(x2, | 380 | + LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset)); |
| 381 | - ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset)); | ||
| 382 | Add(x1, x1, x2, pregLoop); | 381 | Add(x1, x1, x2, pregLoop); |
| 383 | - DataCopy(x3, ((__local_mem__ float*)binaryAddTmpAddr + | 382 | + LoadAlign(x3, ((__ubuf__ float*)binaryAddTmpAddr + |
| 384 | - (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset)); | 383 | + (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset)); |
| 385 | - DataCopy(x4, ((__local_mem__ float*)binaryAddTmpAddr + | 384 | + LoadAlign(x4, ((__ubuf__ float*)binaryAddTmpAddr + |
| 386 | - (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset)); | 385 | + (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset)); |
| 387 | Add(x3, x3, x4, pregLoop); | 386 | Add(x3, x3, x4, pregLoop); |
| 388 | Add(x1, x1, x3, pregLoop); | 387 | Add(x1, x1, x3, pregLoop); |
| 389 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop); | 388 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop); |
| 390 | } | 389 | } |
| 391 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 390 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 392 | } | 391 | } |
| 393 | for (uint16_t i = 0; i < binaryAddLastLoop; i++) { | 392 | for (uint16_t i = 0; i < binaryAddLastLoop; i++) { |
| 394 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + offset)); | 393 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + offset)); |
| 395 | - DataCopy(x2, ((__local_mem__ float*)binaryAddTmpAddr + rLoopStride + offset)); | 394 | + LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + rLoopStride + offset)); |
| 396 | Add(x1, x1, x2, pregLoop); | 395 | Add(x1, x1, x2, pregLoop); |
| 397 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + offset), x1, pregLoop); | 396 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + offset), x1, pregLoop); |
| 398 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 397 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 399 | } | 398 | } |
| 400 | } | 399 | } |
| 401 | 400 | ||
| 402 | - __aicore__ inline void CalculateMean(__local_mem__ float* xInUb, __local_mem__ T_X* yOutUb, | 401 | + __aicore__ inline void CalculateMean(__ubuf__ float* xInUb, __ubuf__ T_X* yOutUb, __ubuf__ float* meanOutUbAddr, |
| 403 | - __local_mem__ float* meanOutUbAddr, int64_t curTileA0Len) | 402 | + int64_t curTileA0Len) |
| 404 | { | 403 | { |
| 405 | if (totalRLen_ <= SCALE_COEF_TWO) { | 404 | if (totalRLen_ <= SCALE_COEF_TWO) { |
| 406 | CalculateMeanRLessThan2(xInUb, meanOutUbAddr, curTileA0Len); | 405 | CalculateMeanRLessThan2(xInUb, meanOutUbAddr, curTileA0Len); |
| @@ -413,7 +412,7 @@ public: | |||
| 413 | } | 412 | } |
| 414 | } | 413 | } |
| 415 | 414 | ||
| 416 | - __aicore__ inline void CalculateMeanRLessThan2(__local_mem__ float* xInUb, __local_mem__ float* meanOutUbAddr, | 415 | + __aicore__ inline void CalculateMeanRLessThan2(__ubuf__ float* xInUb, __ubuf__ float* meanOutUbAddr, |
| 417 | int64_t curTileA0Len) | 416 | int64_t curTileA0Len) |
| 418 | { | 417 | { |
| 419 | uint32_t rStride = tileA0Len_; | 418 | uint32_t rStride = tileA0Len_; |
| @@ -433,16 +432,16 @@ public: | |||
| 433 | pregLoop = UpdateMask<float>(sreg0); | 432 | pregLoop = UpdateMask<float>(sreg0); |
| 434 | Duplicate(sum, 0.0, pregLoop); | 433 | Duplicate(sum, 0.0, pregLoop); |
| 435 | for (uint16_t i = 0; i < rLoopCount; i++) { | 434 | for (uint16_t i = 0; i < rLoopCount; i++) { |
| 436 | - DataCopy(xld, ((__local_mem__ float*)xInUb + i * rStride + k * VL_FP32)); | 435 | + LoadAlign(xld, ((__ubuf__ float*)xInUb + i * rStride + k * VL_FP32)); |
| 437 | Muls(xmuls, xld, n, pregLoop); | 436 | Muls(xmuls, xld, n, pregLoop); |
| 438 | Add(sum, sum, xmuls, pregLoop); | 437 | Add(sum, sum, xmuls, pregLoop); |
| 439 | } | 438 | } |
| 440 | - DataCopy(((__local_mem__ float*)meanOutUbAddr + k * VL_FP32), sum, pregLoop); | 439 | + StoreAlign(((__ubuf__ float*)meanOutUbAddr + k * VL_FP32), sum, pregLoop); |
| 441 | } | 440 | } |
| 442 | } | 441 | } |
| 443 | } | 442 | } |
| 444 | 443 | ||
| 445 | - __aicore__ inline void CalculateMeanRLessThan4(__local_mem__ float* xInUb, __local_mem__ float* meanOutUbAddr, | 444 | + __aicore__ inline void CalculateMeanRLessThan4(__ubuf__ float* xInUb, __ubuf__ float* meanOutUbAddr, |
| 446 | int64_t curTileA0Len) | 445 | int64_t curTileA0Len) |
| 447 | { | 446 | { |
| 448 | uint32_t remainderOffset = SCALE_COEF_TWO * tileA0Len_; | 447 | uint32_t remainderOffset = SCALE_COEF_TWO * tileA0Len_; |
| @@ -473,18 +472,18 @@ public: | |||
| 473 | for (uint16_t k = 0; k < aLoopCount; k++) { | 472 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 474 | pregLoop = UpdateMask<float>(sreg0); | 473 | pregLoop = UpdateMask<float>(sreg0); |
| 475 | uint32_t aLoopOffset = k * VL_FP32; | 474 | uint32_t aLoopOffset = k * VL_FP32; |
| 476 | - DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop); | 475 | + StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop); |
| 477 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 476 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 478 | TwoRowAddForMeanWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 477 | TwoRowAddForMeanWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 479 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, | 478 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, |
| 480 | remNextRow, n); | 479 | remNextRow, n); |
| 481 | Muls(x1, x1, nCorrection, pregLoop); | 480 | Muls(x1, x1, nCorrection, pregLoop); |
| 482 | - DataCopy(((__local_mem__ float*)meanOutUbAddr + aLoopOffset), x1, pregLoop); | 481 | + StoreAlign(((__ubuf__ float*)meanOutUbAddr + aLoopOffset), x1, pregLoop); |
| 483 | } | 482 | } |
| 484 | } | 483 | } |
| 485 | } | 484 | } |
| 486 | 485 | ||
| 487 | - __aicore__ inline void CalculateMeanRLessThan8(__local_mem__ float* xInUb, __local_mem__ float* meanOutUbAddr, | 486 | + __aicore__ inline void CalculateMeanRLessThan8(__ubuf__ float* xInUb, __ubuf__ float* meanOutUbAddr, |
| 488 | int64_t curTileA0Len) | 487 | int64_t curTileA0Len) |
| 489 | { | 488 | { |
| 490 | uint32_t remainderOffset = SCALE_COEF_FOUR * tileA0Len_; | 489 | uint32_t remainderOffset = SCALE_COEF_FOUR * tileA0Len_; |
| @@ -520,7 +519,7 @@ public: | |||
| 520 | for (uint16_t k = 0; k < aLoopCount; k++) { | 519 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 521 | pregLoop = UpdateMask<float>(sreg0); | 520 | pregLoop = UpdateMask<float>(sreg0); |
| 522 | uint32_t aLoopOffset = k * VL_FP32; | 521 | uint32_t aLoopOffset = k * VL_FP32; |
| 523 | - DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop); | 522 | + StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop); |
| 524 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 523 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 525 | TwoRowAddForMeanWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 524 | TwoRowAddForMeanWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 526 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, | 525 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, rem, nextRow, |
| @@ -530,13 +529,13 @@ public: | |||
| 530 | remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow, n); | 529 | remainderTailOffset3 + aLoopOffset, rem, nextRow, remNextRow, n); |
| 531 | Add(x1, x1, x2, pregLoop); | 530 | Add(x1, x1, x2, pregLoop); |
| 532 | Muls(x1, x1, nCorrection, pregLoop); | 531 | Muls(x1, x1, nCorrection, pregLoop); |
| 533 | - DataCopy(((__local_mem__ float*)meanOutUbAddr + aLoopOffset), x1, pregLoop); | 532 | + StoreAlign(((__ubuf__ float*)meanOutUbAddr + aLoopOffset), x1, pregLoop); |
| 534 | } | 533 | } |
| 535 | } | 534 | } |
| 536 | } | 535 | } |
| 537 | 536 | ||
| 538 | - __aicore__ inline void CalculateMeanRMoreThan8(__local_mem__ float* xInUb, __local_mem__ T_X* yOutUb, | 537 | + __aicore__ inline void CalculateMeanRMoreThan8(__ubuf__ float* xInUb, __ubuf__ T_X* yOutUb, |
| 539 | - __local_mem__ float* meanOutUbAddr, int64_t curTileA0Len) | 538 | + __ubuf__ float* meanOutUbAddr, int64_t curTileA0Len) |
| 540 | { | 539 | { |
| 541 | uint16_t remainderLoopCount = (totalRLen_ - r1Quotient_ + SCALE_COEF_EIGHT - 1) / SCALE_COEF_EIGHT; | 540 | uint16_t remainderLoopCount = (totalRLen_ - r1Quotient_ + SCALE_COEF_EIGHT - 1) / SCALE_COEF_EIGHT; |
| 542 | uint16_t quotientLoopCount = (r1Quotient_ / SCALE_COEF_EIGHT) - remainderLoopCount; | 541 | uint16_t quotientLoopCount = (r1Quotient_ / SCALE_COEF_EIGHT) - remainderLoopCount; |
| @@ -597,7 +596,7 @@ public: | |||
| 597 | for (uint16_t k = 0; k < aLoopCount; k++) { | 596 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 598 | pregLoop = UpdateMask<float>(sreg0); | 597 | pregLoop = UpdateMask<float>(sreg0); |
| 599 | uint32_t aLoopOffset = k * VL_FP32; | 598 | uint32_t aLoopOffset = k * VL_FP32; |
| 600 | - DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop); | 599 | + StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), zero, pregLoop); |
| 601 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 600 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 602 | // 前半部分与后半部分中，都为8行的部分 | 601 | // 前半部分与后半部分中，都为8行的部分 |
| 603 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainderLoopCount - 1); i++) { | 602 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainderLoopCount - 1); i++) { |
| @@ -620,7 +619,7 @@ public: | |||
| 620 | remOffset + ROW_SEVEN_OFFSET * aLength, rem, nextRow, remNextRow, n); | 619 | remOffset + ROW_SEVEN_OFFSET * aLength, rem, nextRow, remNextRow, n); |
| 621 | Add(x3, x3, x4, pregLoop); | 620 | Add(x3, x3, x4, pregLoop); |
| 622 | Add(x1, x1, x3, pregLoop); | 621 | Add(x1, x1, x3, pregLoop); |
| 623 | - DataCopy(((__local_mem__ float*)yOutUb + i * aLength + aLoopOffset), x1, pregLoop); | 622 | + StoreAlign(((__ubuf__ float*)yOutUb + i * aLength + aLoopOffset), x1, pregLoop); |
| 624 | } | 623 | } |
| 625 | // 前半部分为8行，后半部分可能不足8行 | 624 | // 前半部分为8行，后半部分可能不足8行 |
| 626 | { | 625 | { |
| @@ -646,8 +645,8 @@ public: | |||
| 646 | remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow, n); | 645 | remainderTailOffset7 + aLoopOffset, rem, nextRow, remNextRow, n); |
| 647 | Add(x3, x3, x4, pregLoop); | 646 | Add(x3, x3, x4, pregLoop); |
| 648 | Add(x1, x1, x3, pregLoop); | 647 | Add(x1, x1, x3, pregLoop); |
| 649 | - DataCopy(((__local_mem__ float*)yOutUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, | 648 | + StoreAlign(((__ubuf__ float*)yOutUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, |
| 650 | - pregLoop); | 649 | + pregLoop); |
| 651 | } | 650 | } |
| 652 | // 剩余的前半部分，一次for循环，处理8行 | 651 | // 剩余的前半部分，一次for循环，处理8行 |
| 653 | for (uint16_t i = 0; i < quotientLoopCount; i++) { | 652 | for (uint16_t i = 0; i < quotientLoopCount; i++) { |
| @@ -662,22 +661,21 @@ public: | |||
| 662 | baseOffset + ROW_SEVEN_OFFSET * aLength, nextRow, n); | 661 | baseOffset + ROW_SEVEN_OFFSET * aLength, nextRow, n); |
| 663 | Add(x3, x3, x4, pregLoop); | 662 | Add(x3, x3, x4, pregLoop); |
| 664 | Add(x1, x1, x3, pregLoop); | 663 | Add(x1, x1, x3, pregLoop); |
| 665 | - DataCopy(((__local_mem__ float*)yOutUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1, | 664 | + StoreAlign(((__ubuf__ float*)yOutUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1, |
| 666 | - pregLoop); | 665 | + pregLoop); |
| 667 | } | 666 | } |
| 668 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 667 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 669 | - BinaryAddVF((__local_mem__ float*)yOutUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop, | 668 | + BinaryAddVF((__ubuf__ float*)yOutUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop, |
| 670 | binaryAddLastLoop, pregLoop, x1, x2, x3, x4); | 669 | binaryAddLastLoop, pregLoop, x1, x2, x3, x4); |
| 671 | - DataCopy(x1, ((__local_mem__ float*)yOutUb + aLoopOffset)); | 670 | + LoadAlign(x1, ((__ubuf__ float*)yOutUb + aLoopOffset)); |
| 672 | Muls(x1, x1, nCorrection, pregLoop); | 671 | Muls(x1, x1, nCorrection, pregLoop); |
| 673 | - DataCopy(((__local_mem__ float*)meanOutUbAddr + aLoopOffset), x1, pregLoop); | 672 | + StoreAlign(((__ubuf__ float*)meanOutUbAddr + aLoopOffset), x1, pregLoop); |
| 674 | } | 673 | } |
| 675 | } | 674 | } |
| 676 | } | 675 | } |
| 677 | 676 | ||
| 678 | - __aicore__ inline void CalculateVar(__local_mem__ float* xInUb, __local_mem__ T_X* yOutUb, | 677 | + __aicore__ inline void CalculateVar(__ubuf__ float* xInUb, __ubuf__ T_X* yOutUb, __ubuf__ float* meanOutUbAddr, |
| 679 | - __local_mem__ float* meanOutUbAddr, __local_mem__ float* varianceOutUbAddr, | 678 | + __ubuf__ float* varianceOutUbAddr, int64_t curTileA0Len) |
| 680 | - int64_t curTileA0Len) | ||
| 681 | { | 679 | { |
| 682 | if (totalRLen_ <= SCALE_COEF_TWO) { | 680 | if (totalRLen_ <= SCALE_COEF_TWO) { |
| 683 | CalculateVarRLessThan2(xInUb, meanOutUbAddr, varianceOutUbAddr, curTileA0Len); | 681 | CalculateVarRLessThan2(xInUb, meanOutUbAddr, varianceOutUbAddr, curTileA0Len); |
| @@ -690,8 +688,8 @@ public: | |||
| 690 | } | 688 | } |
| 691 | } | 689 | } |
| 692 | 690 | ||
| 693 | - __aicore__ inline void CalculateVarRLessThan2(__local_mem__ float* xInUb, __local_mem__ float* meanOutUbAddr, | 691 | + __aicore__ inline void CalculateVarRLessThan2(__ubuf__ float* xInUb, __ubuf__ float* meanOutUbAddr, |
| 694 | - __local_mem__ float* varianceOutUbAddr, int64_t curTileA0Len) | 692 | + __ubuf__ float* varianceOutUbAddr, int64_t curTileA0Len) |
| 695 | { | 693 | { |
| 696 | uint32_t rStride = tileA0Len_; | 694 | uint32_t rStride = tileA0Len_; |
| 697 | uint16_t rLoopCount = totalRLen_; | 695 | uint16_t rLoopCount = totalRLen_; |
| @@ -712,21 +710,21 @@ public: | |||
| 712 | for (uint16_t k = 0; k < aLoopCount; k++) { | 710 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 713 | pregLoop = UpdateMask<float>(sreg0); | 711 | pregLoop = UpdateMask<float>(sreg0); |
| 714 | Duplicate(sum, 0.0, pregLoop); | 712 | Duplicate(sum, 0.0, pregLoop); |
| 715 | - DataCopy(mean, ((__local_mem__ float*)meanOutUbAddr + k * VL_FP32)); | 713 | + LoadAlign(mean, ((__ubuf__ float*)meanOutUbAddr + k * VL_FP32)); |
| 716 | for (uint16_t i = 0; i < rLoopCount; i++) { | 714 | for (uint16_t i = 0; i < rLoopCount; i++) { |
| 717 | - DataCopy(xld, ((__local_mem__ float*)xInUb + i * rStride + k * VL_FP32)); | 715 | + LoadAlign(xld, ((__ubuf__ float*)xInUb + i * rStride + k * VL_FP32)); |
| 718 | Sub(xsub, xld, mean, pregLoop); | 716 | Sub(xsub, xld, mean, pregLoop); |
| 719 | Mul(xpow, xsub, xsub, pregLoop); | 717 | Mul(xpow, xsub, xsub, pregLoop); |
| 720 | Muls(xmuls, xpow, n, pregLoop); | 718 | Muls(xmuls, xpow, n, pregLoop); |
| 721 | Add(sum, sum, xmuls, pregLoop); | 719 | Add(sum, sum, xmuls, pregLoop); |
| 722 | } | 720 | } |
| 723 | - DataCopy(((__local_mem__ float*)varianceOutUbAddr + k * VL_FP32), sum, pregLoop); | 721 | + StoreAlign(((__ubuf__ float*)varianceOutUbAddr + k * VL_FP32), sum, pregLoop); |
| 724 | } | 722 | } |
| 725 | } | 723 | } |
| 726 | } | 724 | } |
| 727 | 725 | ||
| 728 | - __aicore__ inline void CalculateVarRLessThan4(__local_mem__ float* xInUb, __local_mem__ float* meanOutUbAddr, | 726 | + __aicore__ inline void CalculateVarRLessThan4(__ubuf__ float* xInUb, __ubuf__ float* meanOutUbAddr, |
| 729 | - __local_mem__ float* varianceOutUbAddr, int64_t curTileA0Len) | 727 | + __ubuf__ float* varianceOutUbAddr, int64_t curTileA0Len) |
| 730 | { | 728 | { |
| 731 | uint32_t remainderOffset = SCALE_COEF_TWO * tileA0Len_; | 729 | uint32_t remainderOffset = SCALE_COEF_TWO * tileA0Len_; |
| 732 | uint32_t aLength = tileA0Len_; | 730 | uint32_t aLength = tileA0Len_; |
| @@ -753,20 +751,20 @@ public: | |||
| 753 | for (uint16_t k = 0; k < aLoopCount; k++) { | 751 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 754 | pregLoop = UpdateMask<float>(sreg0); | 752 | pregLoop = UpdateMask<float>(sreg0); |
| 755 | uint32_t aLoopOffset = k * VL_FP32; | 753 | uint32_t aLoopOffset = k * VL_FP32; |
| 756 | - DataCopy(mean, ((__local_mem__ float*)meanOutUbAddr + aLoopOffset)); | 754 | + LoadAlign(mean, ((__ubuf__ float*)meanOutUbAddr + aLoopOffset)); |
| 757 | - DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), mean, pregLoop); | 755 | + StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), mean, pregLoop); |
| 758 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 756 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 759 | TwoRowAddForVarWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 757 | TwoRowAddForVarWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 760 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, mean, rem, nextRow, | 758 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, mean, rem, nextRow, |
| 761 | remNextRow, n); | 759 | remNextRow, n); |
| 762 | Muls(x1, x1, nCorrection, pregLoop); | 760 | Muls(x1, x1, nCorrection, pregLoop); |
| 763 | - DataCopy(((__local_mem__ float*)varianceOutUbAddr + aLoopOffset), x1, pregLoop); | 761 | + StoreAlign(((__ubuf__ float*)varianceOutUbAddr + aLoopOffset), x1, pregLoop); |
| 764 | } | 762 | } |
| 765 | } | 763 | } |
| 766 | } | 764 | } |
| 767 | 765 | ||
| 768 | - __aicore__ inline void CalculateVarRLessThan8(__local_mem__ float* xInUb, __local_mem__ float* meanOutUbAddr, | 766 | + __aicore__ inline void CalculateVarRLessThan8(__ubuf__ float* xInUb, __ubuf__ float* meanOutUbAddr, |
| 769 | - __local_mem__ float* varianceOutUbAddr, int64_t curTileA0Len) | 767 | + __ubuf__ float* varianceOutUbAddr, int64_t curTileA0Len) |
| 770 | { | 768 | { |
| 771 | uint32_t remainderOffset = SCALE_COEF_FOUR * tileA0Len_; | 769 | uint32_t remainderOffset = SCALE_COEF_FOUR * tileA0Len_; |
| 772 | uint32_t aLength = tileA0Len_; | 770 | uint32_t aLength = tileA0Len_; |
| @@ -798,8 +796,8 @@ public: | |||
| 798 | for (uint16_t k = 0; k < aLoopCount; k++) { | 796 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 799 | pregLoop = UpdateMask<float>(sreg0); | 797 | pregLoop = UpdateMask<float>(sreg0); |
| 800 | uint32_t aLoopOffset = k * VL_FP32; | 798 | uint32_t aLoopOffset = k * VL_FP32; |
| 801 | - DataCopy(mean, ((__local_mem__ float*)meanOutUbAddr + aLoopOffset)); | 799 | + LoadAlign(mean, ((__ubuf__ float*)meanOutUbAddr + aLoopOffset)); |
| 802 | - DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), mean, pregLoop); | 800 | + StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), mean, pregLoop); |
| 803 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 801 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 804 | TwoRowAddForVarWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, | 802 | TwoRowAddForVarWithTail(x1, xInUb, pregLoop, aLoopOffset, remainderTailOffset0 + aLoopOffset, |
| 805 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, mean, rem, nextRow, | 803 | aLength + aLoopOffset, remainderTailOffset1 + aLoopOffset, mean, rem, nextRow, |
| @@ -809,14 +807,14 @@ public: | |||
| 809 | remainderTailOffset3 + aLoopOffset, mean, rem, nextRow, remNextRow, n); | 807 | remainderTailOffset3 + aLoopOffset, mean, rem, nextRow, remNextRow, n); |
| 810 | Add(x1, x1, x2, pregLoop); | 808 | Add(x1, x1, x2, pregLoop); |
| 811 | Muls(x1, x1, nCorrection, pregLoop); | 809 | Muls(x1, x1, nCorrection, pregLoop); |
| 812 | - DataCopy(((__local_mem__ float*)varianceOutUbAddr + aLoopOffset), x1, pregLoop); | 810 | + StoreAlign(((__ubuf__ float*)varianceOutUbAddr + aLoopOffset), x1, pregLoop); |
| 813 | } | 811 | } |
| 814 | } | 812 | } |
| 815 | } | 813 | } |
| 816 | 814 | ||
| 817 | - __aicore__ inline void CalculateVarRMoreThan8(__local_mem__ float* xInUb, __local_mem__ T_X* yOutUb, | 815 | + __aicore__ inline void CalculateVarRMoreThan8(__ubuf__ float* xInUb, __ubuf__ T_X* yOutUb, |
| 818 | - __local_mem__ float* meanOutUbAddr, | 816 | + __ubuf__ float* meanOutUbAddr, __ubuf__ float* varianceOutUbAddr, |
| 819 | - __local_mem__ float* varianceOutUbAddr, int64_t curTileA0Len) | 817 | + int64_t curTileA0Len) |
| 820 | { | 818 | { |
| 821 | uint16_t remainderLoopCount = (totalRLen_ - r1Quotient_ + SCALE_COEF_EIGHT - 1) / SCALE_COEF_EIGHT; | 819 | uint16_t remainderLoopCount = (totalRLen_ - r1Quotient_ + SCALE_COEF_EIGHT - 1) / SCALE_COEF_EIGHT; |
| 822 | uint16_t quotientLoopCount = (r1Quotient_ / SCALE_COEF_EIGHT) - remainderLoopCount; | 820 | uint16_t quotientLoopCount = (r1Quotient_ / SCALE_COEF_EIGHT) - remainderLoopCount; |
| @@ -874,8 +872,8 @@ public: | |||
| 874 | for (uint16_t k = 0; k < aLoopCount; k++) { | 872 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 875 | pregLoop = UpdateMask<float>(sreg0); | 873 | pregLoop = UpdateMask<float>(sreg0); |
| 876 | uint32_t aLoopOffset = k * VL_FP32; | 874 | uint32_t aLoopOffset = k * VL_FP32; |
| 877 | - DataCopy(mean, ((__local_mem__ float*)meanOutUbAddr + aLoopOffset)); | 875 | + LoadAlign(mean, ((__ubuf__ float*)meanOutUbAddr + aLoopOffset)); |
| 878 | - DataCopy(((__local_mem__ float*)xInUb + validNumInXUb + aLoopOffset), mean, pregLoop); | 876 | + StoreAlign(((__ubuf__ float*)xInUb + validNumInXUb + aLoopOffset), mean, pregLoop); |
| 879 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 877 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 880 | // 前半部分与后半部分中，都为8行的部分 | 878 | // 前半部分与后半部分中，都为8行的部分 |
| 881 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainderLoopCount - 1); i++) { | 879 | for (uint16_t i = 0; i < static_cast<uint16_t>(remainderLoopCount - 1); i++) { |
| @@ -898,7 +896,7 @@ public: | |||
| 898 | remOffset + ROW_SEVEN_OFFSET * aLength, mean, rem, nextRow, remNextRow, n); | 896 | remOffset + ROW_SEVEN_OFFSET * aLength, mean, rem, nextRow, remNextRow, n); |
| 899 | Add(x3, x3, x4, pregLoop); | 897 | Add(x3, x3, x4, pregLoop); |
| 900 | Add(x1, x1, x3, pregLoop); | 898 | Add(x1, x1, x3, pregLoop); |
| 901 | - DataCopy(((__local_mem__ float*)yOutUb + i * aLength + aLoopOffset), x1, pregLoop); | 899 | + StoreAlign(((__ubuf__ float*)yOutUb + i * aLength + aLoopOffset), x1, pregLoop); |
| 902 | } | 900 | } |
| 903 | // 前半部分为8行，后半部分可能不足8行 | 901 | // 前半部分为8行，后半部分可能不足8行 |
| 904 | { | 902 | { |
| @@ -924,8 +922,8 @@ public: | |||
| 924 | remainderTailOffset7 + aLoopOffset, mean, rem, nextRow, remNextRow, n); | 922 | remainderTailOffset7 + aLoopOffset, mean, rem, nextRow, remNextRow, n); |
| 925 | Add(x3, x3, x4, pregLoop); | 923 | Add(x3, x3, x4, pregLoop); |
| 926 | Add(x1, x1, x3, pregLoop); | 924 | Add(x1, x1, x3, pregLoop); |
| 927 | - DataCopy(((__local_mem__ float*)yOutUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, | 925 | + StoreAlign(((__ubuf__ float*)yOutUb + (remainderLoopCount - 1) * aLength + aLoopOffset), x1, |
| 928 | - pregLoop); | 926 | + pregLoop); |
| 929 | } | 927 | } |
| 930 | // 剩余的前半部分，一次for循环，处理8行 | 928 | // 剩余的前半部分，一次for循环，处理8行 |
| 931 | for (uint16_t i = 0; i < quotientLoopCount; i++) { | 929 | for (uint16_t i = 0; i < quotientLoopCount; i++) { |
| @@ -940,20 +938,20 @@ public: | |||
| 940 | baseOffset + ROW_SEVEN_OFFSET * aLength, mean, nextRow, n); | 938 | baseOffset + ROW_SEVEN_OFFSET * aLength, mean, nextRow, n); |
| 941 | Add(x3, x3, x4, pregLoop); | 939 | Add(x3, x3, x4, pregLoop); |
| 942 | Add(x1, x1, x3, pregLoop); | 940 | Add(x1, x1, x3, pregLoop); |
| 943 | - DataCopy(((__local_mem__ float*)yOutUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1, | 941 | + StoreAlign(((__ubuf__ float*)yOutUb + (remainderLoopCount + i) * aLength + aLoopOffset), x1, |
| 944 | - pregLoop); | 942 | + pregLoop); |
| 945 | } | 943 | } |
| 946 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 944 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 947 | - BinaryAddVF((__local_mem__ float*)yOutUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop, | 945 | + BinaryAddVF((__ubuf__ float*)yOutUb, aLength, aLoopOffset, binaryAddKLoop, binaryAddInnerLoop, |
| 948 | binaryAddLastLoop, pregLoop, x1, x2, x3, x4); | 946 | binaryAddLastLoop, pregLoop, x1, x2, x3, x4); |
| 949 | - DataCopy(x1, ((__local_mem__ float*)yOutUb + aLoopOffset)); | 947 | + LoadAlign(x1, ((__ubuf__ float*)yOutUb + aLoopOffset)); |
| 950 | Muls(x1, x1, nCorrection, pregLoop); | 948 | Muls(x1, x1, nCorrection, pregLoop); |
| 951 | - DataCopy(((__local_mem__ float*)varianceOutUbAddr + aLoopOffset), x1, pregLoop); | 949 | + StoreAlign(((__ubuf__ float*)varianceOutUbAddr + aLoopOffset), x1, pregLoop); |
| 952 | } | 950 | } |
| 953 | } | 951 | } |
| 954 | } | 952 | } |
| 955 | 953 | ||
| 956 | - __aicore__ inline void ComputeRstd(__local_mem__ float* varianceOutUbAddr, __local_mem__ float* rstdUbAddr, | 954 | + __aicore__ inline void ComputeRstd(__ubuf__ float* varianceOutUbAddr, __ubuf__ float* rstdUbAddr, |
| 957 | uint64_t curTileA0Len, float epsilon) | 955 | uint64_t curTileA0Len, float epsilon) |
| 958 | { | 956 | { |
| 959 | uint16_t aLoop = static_cast<uint16_t>((curTileA0Len + VL_FP32 - 1) / VL_FP32); | 957 | uint16_t aLoop = static_cast<uint16_t>((curTileA0Len + VL_FP32 - 1) / VL_FP32); |
| @@ -990,7 +988,7 @@ public: | |||
| 990 | Duplicate(s, float(1.0), pregLoop); | 988 | Duplicate(s, float(1.0), pregLoop); |
| 991 | 989 | ||
| 992 | // rstd | 990 | // rstd |
| 993 | - DataCopy(var, varianceOutUbAddr + a * VL_FP32); | 991 | + LoadAlign(var, varianceOutUbAddr + a * VL_FP32); |
| 994 | Adds(var, var, epsilon, pregLoop); | 992 | Adds(var, var, epsilon, pregLoop); |
| 995 | Div(r, one, var, pregLoop); | 993 | Div(r, one, var, pregLoop); |
| 996 | Sqrt(y, r, pregLoop); | 994 | Sqrt(y, r, pregLoop); |
| @@ -1005,19 +1003,18 @@ public: | |||
| 1005 | Mula(s, var, r, pregLoop); // s + x * t | 1003 | Mula(s, var, r, pregLoop); // s + x * t |
| 1006 | Mul(s, s, rstd, pregLoop); // e * y | 1004 | Mul(s, s, rstd, pregLoop); // e * y |
| 1007 | Mula(rstd, s, scalar1, pregLoop); // y + y * e * 0.5 | 1005 | Mula(rstd, s, scalar1, pregLoop); // y + y * e * 0.5 |
| 1008 | - CompareScalar(cmpRegZero, var, RMS_POS_INF, pregLoop); | 1006 | + Compares(cmpRegZero, var, RMS_POS_INF, pregLoop); |
| 1009 | Select(rstd, scalarZero, rstd, cmpRegZero); | 1007 | Select(rstd, scalarZero, rstd, cmpRegZero); |
| 1010 | - CompareScalar(cmpRegInf, var, RMS_ZERO, pregLoop); | 1008 | + Compares(cmpRegInf, var, RMS_ZERO, pregLoop); |
| 1011 | Select(rstd, scalarInf, rstd, cmpRegInf); | 1009 | Select(rstd, scalarInf, rstd, cmpRegInf); |
| 1012 | - DataCopy(rstdUbAddr + a * VL_FP32, rstd, pregLoop); | 1010 | + StoreAlign(rstdUbAddr + a * VL_FP32, rstd, pregLoop); |
| 1013 | } | 1011 | } |
| 1014 | } | 1012 | } |
| 1015 | } | 1013 | } |
| 1016 | 1014 | ||
| 1017 | - __aicore__ inline void CalculateY(__local_mem__ float* xInUb, __local_mem__ T_X* yOutUb, | 1015 | + __aicore__ inline void CalculateY(__ubuf__ float* xInUb, __ubuf__ T_X* yOutUb, __ubuf__ T_B* betaInUb, |
| 1018 | - __local_mem__ T_B* betaInUb, __local_mem__ T_B* gammaInUb, | 1016 | + __ubuf__ T_B* gammaInUb, __ubuf__ float* meanOutUbAddr, |
| 1019 | - __local_mem__ float* meanOutUbAddr, __local_mem__ float* rstdUbAddr, | 1017 | + __ubuf__ float* rstdUbAddr, uint16_t curTileA0Len) |
| 1020 | - uint16_t curTileA0Len) | ||
| 1021 | { | 1018 | { |
| 1022 | uint16_t rLoopCount = totalRLen_; | 1019 | uint16_t rLoopCount = totalRLen_; |
| 1023 | uint16_t aLoopCount = CEIL_DIV(curTileA0Len, VL_FP32); | 1020 | uint16_t aLoopCount = CEIL_DIV(curTileA0Len, VL_FP32); |
| @@ -1038,10 +1035,10 @@ public: | |||
| 1038 | for (uint16_t k = 0; k < aLoopCount; k++) { | 1035 | for (uint16_t k = 0; k < aLoopCount; k++) { |
| 1039 | pregLoop = UpdateMask<float>(sreg2); | 1036 | pregLoop = UpdateMask<float>(sreg2); |
| 1040 | LoadTwoTensorForDtypeT(betaInUb, gammaInUb, beta, gamma, pregLoop, pregLoop, k * VL_FP32, k * VL_FP32); | 1037 | LoadTwoTensorForDtypeT(betaInUb, gammaInUb, beta, gamma, pregLoop, pregLoop, k * VL_FP32, k * VL_FP32); |
| 1041 | - DataCopy(mean, ((__local_mem__ float*)meanOutUbAddr + k * VL_FP32)); | 1038 | + LoadAlign(mean, ((__ubuf__ float*)meanOutUbAddr + k * VL_FP32)); |
| 1042 | - DataCopy(rsqrtVar, ((__local_mem__ float*)rstdUbAddr + k * VL_FP32)); | 1039 | + LoadAlign(rsqrtVar, ((__ubuf__ float*)rstdUbAddr + k * VL_FP32)); |
| 1043 | for (uint16_t r = 0; r < rLoopCount; r++) { | 1040 | for (uint16_t r = 0; r < rLoopCount; r++) { |
| 1044 | - DataCopy(x2, ((__local_mem__ float*)xInUb + r * rStride + k * VL_FP32)); | 1041 | + LoadAlign(x2, ((__ubuf__ float*)xInUb + r * rStride + k * VL_FP32)); |
| 1045 | Sub(x2, x2, mean, pregLoop); | 1042 | Sub(x2, x2, mean, pregLoop); |
| 1046 | Mul(y2, x2, rsqrtVar, pregLoop); | 1043 | Mul(y2, x2, rsqrtVar, pregLoop); |
| 1047 | Mul(y2, y2, gamma, pregLoop); | 1044 | Mul(y2, y2, gamma, pregLoop); |
| @@ -1049,15 +1046,15 @@ public: | |||
| 1049 | if constexpr (IsSameType<T_X, half>::value) { | 1046 | if constexpr (IsSameType<T_X, half>::value) { |
| 1050 | RegTensor<half> yFp16; | 1047 | RegTensor<half> yFp16; |
| 1051 | Cast<half, float, castTraitB322B16>(yFp16, y2, pregLoop); | 1048 | Cast<half, float, castTraitB322B16>(yFp16, y2, pregLoop); |
| 1052 | - DataCopy<half, StoreDist::DIST_PACK_B32>( | 1049 | + StoreAlign<half, StoreDist::DIST_PACK_B32>(((__ubuf__ half*)yOutUb + r * rStride + k * VL_FP32), |
| 1053 | - ((__local_mem__ half*)yOutUb + r * rStride + k * VL_FP32), yFp16, pregLoop); | 1050 | + yFp16, pregLoop); |
| 1054 | } else if constexpr (IsSameType<T_X, bfloat16_t>::value) { | 1051 | } else if constexpr (IsSameType<T_X, bfloat16_t>::value) { |
| 1055 | RegTensor<bfloat16_t> xBf16; | 1052 | RegTensor<bfloat16_t> xBf16; |
| 1056 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, y2, pregLoop); | 1053 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, y2, pregLoop); |
| 1057 | - DataCopy<bfloat16_t, StoreDist::DIST_PACK_B32>( | 1054 | + StoreAlign<bfloat16_t, StoreDist::DIST_PACK_B32>( |
| 1058 | - ((__local_mem__ bfloat16_t*)yOutUb + r * rStride + k * VL_FP32), xBf16, pregLoop); | 1055 | + ((__ubuf__ bfloat16_t*)yOutUb + r * rStride + k * VL_FP32), xBf16, pregLoop); |
| 1059 | } else { | 1056 | } else { |
| 1060 | - DataCopy(((__local_mem__ float*)yOutUb + r * rStride + k * VL_FP32), y2, pregLoop); | 1057 | + StoreAlign(((__ubuf__ float*)yOutUb + r * rStride + k * VL_FP32), y2, pregLoop); |
| 1061 | } | 1058 | } |
| 1062 | } | 1059 | } |
| 1063 | } | 1060 | } |
| @@ -1066,10 +1063,10 @@ public: | |||
| 1066 | __aicore__ inline void CastMeanRstd(LocalTensor<float> meanOutUb, LocalTensor<float> varianceOutUb, | 1063 | __aicore__ inline void CastMeanRstd(LocalTensor<float> meanOutUb, LocalTensor<float> varianceOutUb, |
| 1067 | uint64_t curTileA0Len) | 1064 | uint64_t curTileA0Len) |
| 1068 | { | 1065 | { |
| 1069 | - __local_mem__ float* meanInAddr = (__local_mem__ float*)meanOutUb.GetPhyAddr(); | 1066 | + __ubuf__ float* meanInAddr = (__ubuf__ float*)meanOutUb.GetPhyAddr(); |
| 1070 | - __local_mem__ float* varianceInAddr = (__local_mem__ float*)varianceOutUb.GetPhyAddr(); | 1067 | + __ubuf__ float* varianceInAddr = (__ubuf__ float*)varianceOutUb.GetPhyAddr(); |
| 1071 | - __local_mem__ T_M* meanOutAddr = (__local_mem__ T_M*)meanOutUb.GetPhyAddr(); | 1068 | + __ubuf__ T_M* meanOutAddr = (__ubuf__ T_M*)meanOutUb.GetPhyAddr(); |
| 1072 | - __local_mem__ T_M* varianceOutAddr = (__local_mem__ T_M*)varianceOutUb.GetPhyAddr(); | 1069 | + __ubuf__ T_M* varianceOutAddr = (__ubuf__ T_M*)varianceOutUb.GetPhyAddr(); |
| 1073 | 1070 | ||
| 1074 | uint32_t castCount = static_cast<uint32_t>(curTileA0Len); | 1071 | uint32_t castCount = static_cast<uint32_t>(curTileA0Len); |
| 1075 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_FP32 - 1) / VL_FP32); | 1072 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_FP32 - 1) / VL_FP32); |
| @@ -1082,14 +1079,14 @@ public: | |||
| 1082 | MicroAPI::MaskReg pregLoop; | 1079 | MicroAPI::MaskReg pregLoop; |
| 1083 | for (uint16_t i = 0; i < castLoops; i++) { | 1080 | for (uint16_t i = 0; i < castLoops; i++) { |
| 1084 | pregLoop = MicroAPI::UpdateMask<float>(castCount); | 1081 | pregLoop = MicroAPI::UpdateMask<float>(castCount); |
| 1085 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, meanInAddr + VL_FP32 * i); | 1082 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, meanInAddr + VL_FP32 * i); |
| 1086 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_rstd, varianceInAddr + VL_FP32 * i); | 1083 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_rstd, varianceInAddr + VL_FP32 * i); |
| 1087 | Cast<T_M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 1084 | Cast<T_M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 1088 | Cast<T_M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); | 1085 | Cast<T_M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); |
| 1089 | - DataCopy<T_M, StoreDist::DIST_PACK_B32>(((__local_mem__ T_M*)meanOutAddr + i * VL_MEAN), output_mean, | 1086 | + StoreAlign<T_M, StoreDist::DIST_PACK_B32>(((__ubuf__ T_M*)meanOutAddr + i * VL_MEAN), output_mean, |
| 1090 | - pregLoop); | 1087 | + pregLoop); |
| 1091 | - DataCopy<T_M, StoreDist::DIST_PACK_B32>(((__local_mem__ T_M*)varianceOutAddr + i * VL_MEAN), | 1088 | + StoreAlign<T_M, StoreDist::DIST_PACK_B32>(((__ubuf__ T_M*)varianceOutAddr + i * VL_MEAN), output_rstd, |
| 1092 | - output_rstd, pregLoop); | 1089 | + pregLoop); |
| 1093 | } | 1090 | } |
| 1094 | } | 1091 | } |
| 1095 | } | 1092 | } |
| @@ -18,6 +18,7 @@ | |||
| 18 | 18 | ||
| 19 | 19 | ||
| 20 | 20 | ||
| 21 | + | ||
| 21 | 22 | ||
| 22 | namespace InstanceNormOps { | 23 | namespace InstanceNormOps { |
| 23 | using namespace AscendC; | 24 | using namespace AscendC; |
| @@ -30,6 +31,8 @@ using AscendC::MicroAPI::MemType; | |||
| 30 | using AscendC::MicroAPI::RegTensor; | 31 | using AscendC::MicroAPI::RegTensor; |
| 31 | using AscendC::MicroAPI::StoreDist; | 32 | using AscendC::MicroAPI::StoreDist; |
| 32 | using AscendC::MicroAPI::UpdateMask; | 33 | using AscendC::MicroAPI::UpdateMask; |
| 34 | +using AscendC::Reg::LoadAlign; | ||
| 35 | +using AscendC::Reg::StoreAlign; | ||
| 33 | 36 | ||
| 34 | template <typename T, typename T_BETA, typename T_MEAN> | 37 | template <typename T, typename T_BETA, typename T_MEAN> |
| 35 | class InstanceNormARAWelford { | 38 | class InstanceNormARAWelford { |
| @@ -120,7 +123,7 @@ private: | |||
| 120 | __aicore__ inline void CalculateCountBuf() | 123 | __aicore__ inline void CalculateCountBuf() |
| 121 | { | 124 | { |
| 122 | LocalTensor<float> tCountTensor = tCountBuff.Get<float>(); | 125 | LocalTensor<float> tCountTensor = tCountBuff.Get<float>(); |
| 123 | - __local_mem__ float* tmpCountLocal = (__local_mem__ float*)tCountTensor.GetPhyAddr(); | 126 | + __ubuf__ float* tmpCountLocal = (__ubuf__ float*)tCountTensor.GetPhyAddr(); |
| 124 | int64_t parallelCount = this->r / this->rFactor; | 127 | int64_t parallelCount = this->r / this->rFactor; |
| 125 | int64_t parallelReminder = this->r % this->rFactor; | 128 | int64_t parallelReminder = this->r % this->rFactor; |
| 126 | float quotientAddCount = static_cast<float>(parallelCount); | 129 | float quotientAddCount = static_cast<float>(parallelCount); |
| @@ -141,13 +144,13 @@ private: | |||
| 141 | Duplicate(tmpCount, quotientAddCount, pregMain); | 144 | Duplicate(tmpCount, quotientAddCount, pregMain); |
| 142 | for (uint16_t i = 0; i < quotientLoopCount; i++) { | 145 | for (uint16_t i = 0; i < quotientLoopCount; i++) { |
| 143 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg1); | 146 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg1); |
| 144 | - DataCopy(((__local_mem__ float*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); | 147 | + StoreAlign(((__ubuf__ float*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); |
| 145 | } | 148 | } |
| 146 | uint32_t sreg2 = remainderNum; | 149 | uint32_t sreg2 = remainderNum; |
| 147 | Duplicate(tmpCount, remaninderAddCount, pregMain); | 150 | Duplicate(tmpCount, remaninderAddCount, pregMain); |
| 148 | for (uint16_t i = 0; i < remainderLoopCount; i++) { | 151 | for (uint16_t i = 0; i < remainderLoopCount; i++) { |
| 149 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg2); | 152 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg2); |
| 150 | - DataCopy(((__local_mem__ float*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); | 153 | + StoreAlign(((__ubuf__ float*)tmpCountLocal + i * VL_F32), tmpCount, pregLoop); |
| 151 | } | 154 | } |
| 152 | } | 155 | } |
| 153 | } | 156 | } |
| @@ -158,18 +161,18 @@ private: | |||
| 158 | LocalTensor<float> tMeanTensor = tMeanBuff.Get<float>(); | 161 | LocalTensor<float> tMeanTensor = tMeanBuff.Get<float>(); |
| 159 | LocalTensor<float> tVarTensor = tVarBuff.Get<float>(); | 162 | LocalTensor<float> tVarTensor = tVarBuff.Get<float>(); |
| 160 | LocalTensor<float> tCountTensor = tCountBuff.Get<float>(); | 163 | LocalTensor<float> tCountTensor = tCountBuff.Get<float>(); |
| 161 | - __local_mem__ float* rstdLocal = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 164 | + __ubuf__ float* rstdLocal = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 162 | - __local_mem__ float* tmpMeanLocal = (__local_mem__ float*)tMeanTensor.GetPhyAddr(); | 165 | + __ubuf__ float* tmpMeanLocal = (__ubuf__ float*)tMeanTensor.GetPhyAddr(); |
| 163 | - __local_mem__ float* tmpVarLocal = (__local_mem__ float*)tVarTensor.GetPhyAddr(); | 166 | + __ubuf__ float* tmpVarLocal = (__ubuf__ float*)tVarTensor.GetPhyAddr(); |
| 164 | - __local_mem__ float* tmpCountLocal = (__local_mem__ float*)tCountTensor.GetPhyAddr(); | 167 | + __ubuf__ float* tmpCountLocal = (__ubuf__ float*)tCountTensor.GetPhyAddr(); |
| 165 | 168 | ||
| 166 | ProcessWelfordUpdate(curA0Idx, curA1Idx, currentANum, tmpMeanLocal, tmpVarLocal); | 169 | ProcessWelfordUpdate(curA0Idx, curA1Idx, currentANum, tmpMeanLocal, tmpVarLocal); |
| 167 | CopyInGammaBeta(curA0Idx, currentANum); | 170 | CopyInGammaBeta(curA0Idx, currentANum); |
| 168 | 171 | ||
| 169 | LocalTensor<float> batchMeanOutUb = batchMeanQueue.AllocTensor<float>(); | 172 | LocalTensor<float> batchMeanOutUb = batchMeanQueue.AllocTensor<float>(); |
| 170 | LocalTensor<float> batchVarOutUb = batchVarQueue.AllocTensor<float>(); | 173 | LocalTensor<float> batchVarOutUb = batchVarQueue.AllocTensor<float>(); |
| 171 | - __local_mem__ float* batchMeanInUbAddr = (__local_mem__ float*)batchMeanOutUb.GetPhyAddr(); | 174 | + __ubuf__ float* batchMeanInUbAddr = (__ubuf__ float*)batchMeanOutUb.GetPhyAddr(); |
| 172 | - __local_mem__ float* batchVarInUbAddr = (__local_mem__ float*)batchVarOutUb.GetPhyAddr(); | 175 | + __ubuf__ float* batchVarInUbAddr = (__ubuf__ float*)batchVarOutUb.GetPhyAddr(); |
| 173 | ProcessWelfordFinalize(currentANum, tmpMeanLocal, tmpVarLocal, tmpCountLocal, batchMeanInUbAddr, | 176 | ProcessWelfordFinalize(currentANum, tmpMeanLocal, tmpVarLocal, tmpCountLocal, batchMeanInUbAddr, |
| 174 | batchVarInUbAddr); | 177 | batchVarInUbAddr); |
| 175 | // 此时batchMean 和 batchVar都是累加计算结果，大小是aFactor，输出的mean和var应该是这个，后面开始计算rstd | 178 | // 此时batchMean 和 batchVar都是累加计算结果，大小是aFactor，输出的mean和var应该是这个，后面开始计算rstd |
| @@ -181,7 +184,7 @@ private: | |||
| 181 | } | 184 | } |
| 182 | 185 | ||
| 183 | __aicore__ inline void ProcessWelfordUpdate(int64_t curA0Idx, int64_t curA1Idx, int64_t currentANum, | 186 | __aicore__ inline void ProcessWelfordUpdate(int64_t curA0Idx, int64_t curA1Idx, int64_t currentANum, |
| 184 | - __local_mem__ float* tmpMeanLocal, __local_mem__ float* tmpVarLocal) | 187 | + __ubuf__ float* tmpMeanLocal, __ubuf__ float* tmpVarLocal) |
| 185 | { | 188 | { |
| 186 | int64_t quotient = (this->r + this->rFactor - 1) / this->rFactor; | 189 | int64_t quotient = (this->r + this->rFactor - 1) / this->rFactor; |
| 187 | for (int64_t rLoopIdx = 0; rLoopIdx < quotient; rLoopIdx++) { | 190 | for (int64_t rLoopIdx = 0; rLoopIdx < quotient; rLoopIdx++) { |
| @@ -193,7 +196,7 @@ private: | |||
| 193 | CopyInX(copyXOffset, currentR, currentANum); | 196 | CopyInX(copyXOffset, currentR, currentANum); |
| 194 | 197 | ||
| 195 | LocalTensor<T> xInUb = xQueue.DeQue<T>(); | 198 | LocalTensor<T> xInUb = xQueue.DeQue<T>(); |
| 196 | - __local_mem__ T* xLocal = (__local_mem__ T*)xInUb.GetPhyAddr(); | 199 | + __ubuf__ T* xLocal = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 197 | // process welford after copy ubSize data into ub. | 200 | // process welford after copy ubSize data into ub. |
| 198 | float scale = (float)1.0 / static_cast<float>(rLoopIdx + 1); | 201 | float scale = (float)1.0 / static_cast<float>(rLoopIdx + 1); |
| 199 | uint64_t processNum = currentR * currentANumAlign; | 202 | uint64_t processNum = currentR * currentANumAlign; |
| @@ -227,24 +230,24 @@ private: | |||
| 227 | } | 230 | } |
| 228 | 231 | ||
| 229 | template <typename T_SRC> | 232 | template <typename T_SRC> |
| 230 | - __aicore__ inline void LoadOneTensorForDtypeT(__local_mem__ T_SRC* input, RegTensor<float>& dst, MaskReg& preg, | 233 | + __aicore__ inline void LoadOneTensorForDtypeT(__ubuf__ T_SRC* input, RegTensor<float>& dst, MaskReg& preg, |
| 231 | uint32_t offset) | 234 | uint32_t offset) |
| 232 | { | 235 | { |
| 233 | if constexpr (IsSameType<T_SRC, half>::value) { | 236 | if constexpr (IsSameType<T_SRC, half>::value) { |
| 234 | RegTensor<half> xFp16; | 237 | RegTensor<half> xFp16; |
| 235 | - DataCopy<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__local_mem__ half*)(input) + (offset))); | 238 | + LoadAlign<half, LoadDist::DIST_UNPACK_B16>(xFp16, ((__ubuf__ half*)(input) + (offset))); |
| 236 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); | 239 | Cast<float, half, castTraitB162B32>(dst, xFp16, preg); |
| 237 | } else if constexpr (IsSameType<T_SRC, bfloat16_t>::value) { | 240 | } else if constexpr (IsSameType<T_SRC, bfloat16_t>::value) { |
| 238 | RegTensor<bfloat16_t> xBf16; | 241 | RegTensor<bfloat16_t> xBf16; |
| 239 | - DataCopy<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__local_mem__ bfloat16_t*)(input) + (offset))); | 242 | + LoadAlign<bfloat16_t, LoadDist::DIST_UNPACK_B16>(xBf16, ((__ubuf__ bfloat16_t*)(input) + (offset))); |
| 240 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); | 243 | Cast<float, bfloat16_t, castTraitB162B32>(dst, xBf16, preg); |
| 241 | } else { | 244 | } else { |
| 242 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset))); | 245 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset))); |
| 243 | } | 246 | } |
| 244 | } | 247 | } |
| 245 | 248 | ||
| 246 | - __aicore__ inline void WelfordParallelUpdateWithInitVF(__local_mem__ T* x1Local, __local_mem__ float* tmpMeanLocal, | 249 | + __aicore__ inline void WelfordParallelUpdateWithInitVF(__ubuf__ T* x1Local, __ubuf__ float* tmpMeanLocal, |
| 247 | - __local_mem__ float* tmpVarLocal, uint64_t calLen, | 250 | + __ubuf__ float* tmpVarLocal, uint64_t calLen, |
| 248 | uint16_t loopCount, float scale) | 251 | uint16_t loopCount, float scale) |
| 249 | { | 252 | { |
| 250 | __VEC_SCOPE__ | 253 | __VEC_SCOPE__ |
| @@ -269,7 +272,7 @@ private: | |||
| 269 | Muls(delta2, delta1, scale, pregLoop); | 272 | Muls(delta2, delta1, scale, pregLoop); |
| 270 | // mean = mean + delta2 | 273 | // mean = mean + delta2 |
| 271 | Add(tmpMean, tmpMean, delta2, pregLoop); | 274 | Add(tmpMean, tmpMean, delta2, pregLoop); |
| 272 | - DataCopy(tmpMeanLocal + i * VL_F32, tmpMean, pregLoop); | 275 | + StoreAlign(tmpMeanLocal + i * VL_F32, tmpMean, pregLoop); |
| 273 | 276 | ||
| 274 | Duplicate(tmpVar, 0.0, pregLoop); | 277 | Duplicate(tmpVar, 0.0, pregLoop); |
| 275 | // delta3 = x1 - mean | 278 | // delta3 = x1 - mean |
| @@ -278,14 +281,14 @@ private: | |||
| 278 | Mul(delat4, delta1, delta3, pregLoop); | 281 | Mul(delat4, delta1, delta3, pregLoop); |
| 279 | // var = var + delta4 | 282 | // var = var + delta4 |
| 280 | Add(tmpVar, tmpVar, delat4, pregLoop); | 283 | Add(tmpVar, tmpVar, delat4, pregLoop); |
| 281 | - DataCopy(tmpVarLocal + i * VL_F32, tmpVar, pregLoop); | 284 | + StoreAlign(tmpVarLocal + i * VL_F32, tmpVar, pregLoop); |
| 282 | } | 285 | } |
| 283 | } | 286 | } |
| 284 | } | 287 | } |
| 285 | 288 | ||
| 286 | - __aicore__ inline void WelfordParallelUpdateVF(__local_mem__ T* x1Local, __local_mem__ float* tmpMeanLocal, | 289 | + __aicore__ inline void WelfordParallelUpdateVF(__ubuf__ T* x1Local, __ubuf__ float* tmpMeanLocal, |
| 287 | - __local_mem__ float* tmpVarLocal, uint64_t calLen, | 290 | + __ubuf__ float* tmpVarLocal, uint64_t calLen, uint16_t loopCount, |
| 288 | - uint16_t loopCount, float scale) | 291 | + float scale) |
| 289 | { | 292 | { |
| 290 | __VEC_SCOPE__ | 293 | __VEC_SCOPE__ |
| 291 | { | 294 | { |
| @@ -303,23 +306,23 @@ private: | |||
| 303 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0); | 306 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0); |
| 304 | LoadOneTensorForDtypeT(x1Local, x1, pregLoop, i * VL_F32); | 307 | LoadOneTensorForDtypeT(x1Local, x1, pregLoop, i * VL_F32); |
| 305 | 308 | ||
| 306 | - DataCopy(tmpMean, tmpMeanLocal + i * VL_F32); | 309 | + LoadAlign(tmpMean, tmpMeanLocal + i * VL_F32); |
| 307 | // delata1 = x1 - mean | 310 | // delata1 = x1 - mean |
| 308 | Sub(delta1, x1, tmpMean, pregLoop); | 311 | Sub(delta1, x1, tmpMean, pregLoop); |
| 309 | // delta2 = delta1 * scale | 312 | // delta2 = delta1 * scale |
| 310 | Muls(delta2, delta1, scale, pregLoop); | 313 | Muls(delta2, delta1, scale, pregLoop); |
| 311 | // mean = mean + delta2 | 314 | // mean = mean + delta2 |
| 312 | Add(tmpMean, tmpMean, delta2, pregLoop); | 315 | Add(tmpMean, tmpMean, delta2, pregLoop); |
| 313 | - DataCopy(tmpMeanLocal + i * VL_F32, tmpMean, pregLoop); | 316 | + StoreAlign(tmpMeanLocal + i * VL_F32, tmpMean, pregLoop); |
| 314 | 317 | ||
| 315 | - DataCopy(tmpVar, tmpVarLocal + i * VL_F32); | 318 | + LoadAlign(tmpVar, tmpVarLocal + i * VL_F32); |
| 316 | // delta3 = x1 - mean | 319 | // delta3 = x1 - mean |
| 317 | Sub(delta3, x1, tmpMean, pregLoop); | 320 | Sub(delta3, x1, tmpMean, pregLoop); |
| 318 | // delta4 = delta1 * delta3 | 321 | // delta4 = delta1 * delta3 |
| 319 | Mul(delat4, delta1, delta3, pregLoop); | 322 | Mul(delat4, delta1, delta3, pregLoop); |
| 320 | // var = var + delta4 | 323 | // var = var + delta4 |
| 321 | Add(tmpVar, tmpVar, delat4, pregLoop); | 324 | Add(tmpVar, tmpVar, delat4, pregLoop); |
| 322 | - DataCopy(tmpVarLocal + i * VL_F32, tmpVar, pregLoop); | 325 | + StoreAlign(tmpVarLocal + i * VL_F32, tmpVar, pregLoop); |
| 323 | } | 326 | } |
| 324 | } | 327 | } |
| 325 | } | 328 | } |
| @@ -345,13 +348,12 @@ private: | |||
| 345 | gammaQueue.EnQue(gammaInUb); | 348 | gammaQueue.EnQue(gammaInUb); |
| 346 | } | 349 | } |
| 347 | 350 | ||
| 348 | - __aicore__ inline void ProcessWelfordFinalize(int64_t currentANum, __local_mem__ float* tmpMeanLocal, | 351 | + __aicore__ inline void ProcessWelfordFinalize(int64_t currentANum, __ubuf__ float* tmpMeanLocal, |
| 349 | - __local_mem__ float* tmpVarLocal, __local_mem__ float* tmpCountLocal, | 352 | + __ubuf__ float* tmpVarLocal, __ubuf__ float* tmpCountLocal, |
| 350 | - __local_mem__ float* batchMeanInUbAddr, | 353 | + __ubuf__ float* batchMeanInUbAddr, __ubuf__ float* batchVarInUbAddr) |
| 351 | - __local_mem__ float* batchVarInUbAddr) | ||
| 352 | { | 354 | { |
| 353 | LocalTensor<T> yInUb = yQueue.AllocTensor<T>(); | 355 | LocalTensor<T> yInUb = yQueue.AllocTensor<T>(); |
| 354 | - __local_mem__ float* yInUbAddr = (__local_mem__ float*)yInUb.GetPhyAddr(); | 356 | + __ubuf__ float* yInUbAddr = (__ubuf__ float*)yInUb.GetPhyAddr(); |
| 355 | WelfordFinalizeMeanVF(currentANum, tmpMeanLocal, tmpVarLocal, tmpCountLocal, batchMeanInUbAddr, | 357 | WelfordFinalizeMeanVF(currentANum, tmpMeanLocal, tmpVarLocal, tmpCountLocal, batchMeanInUbAddr, |
| 356 | batchVarInUbAddr, yInUbAddr); | 358 | batchVarInUbAddr, yInUbAddr); |
| 357 | WelfordFinalizeVarVF(currentANum, tmpMeanLocal, tmpVarLocal, tmpCountLocal, batchMeanInUbAddr, batchVarInUbAddr, | 359 | WelfordFinalizeVarVF(currentANum, tmpMeanLocal, tmpVarLocal, tmpCountLocal, batchMeanInUbAddr, batchVarInUbAddr, |
| @@ -359,11 +361,10 @@ private: | |||
| 359 | yQueue.FreeTensor(yInUb); | 361 | yQueue.FreeTensor(yInUb); |
| 360 | } | 362 | } |
| 361 | 363 | ||
| 362 | - __aicore__ inline void WelfordFinalizeMeanVF(int64_t currentANum, __local_mem__ float* tmpMeanLocal, | 364 | + __aicore__ inline void WelfordFinalizeMeanVF(int64_t currentANum, __ubuf__ float* tmpMeanLocal, |
| 363 | - __local_mem__ float* tmpVarLocal, __local_mem__ float* tmpCountLocal, | 365 | + __ubuf__ float* tmpVarLocal, __ubuf__ float* tmpCountLocal, |
| 364 | - __local_mem__ float* batchMeanInUbAddr, | 366 | + __ubuf__ float* batchMeanInUbAddr, __ubuf__ float* batchVarInUbAddr, |
| 365 | - __local_mem__ float* batchVarInUbAddr, | 367 | + __ubuf__ float* binaryAddTmpAddr) |
| 366 | - __local_mem__ float* binaryAddTmpAddr) | ||
| 367 | { | 368 | { |
| 368 | uint16_t rLoopCount = this->rFactor; | 369 | uint16_t rLoopCount = this->rFactor; |
| 369 | uint16_t aLoopCount = this->currentALoopCount; | 370 | uint16_t aLoopCount = this->currentALoopCount; |
| @@ -435,7 +436,7 @@ private: | |||
| 435 | remCountOffset + ROW_THREE_OFFSET, rem, nextRow, remNextRow, rowCount, | 436 | remCountOffset + ROW_THREE_OFFSET, rem, nextRow, remNextRow, rowCount, |
| 436 | nextRowCount, remCount, nextRemCount, numScale); | 437 | nextRowCount, remCount, nextRemCount, numScale); |
| 437 | Add(x1, x1, x2, pregLoop); | 438 | Add(x1, x1, x2, pregLoop); |
| 438 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + i * rLoopStride + aLoopOffset), x1, pregLoop); | 439 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + i * rLoopStride + aLoopOffset), x1, pregLoop); |
| 439 | } | 440 | } |
| 440 | // 剩余的前半部分，一次for循环，处理8行 | 441 | // 剩余的前半部分，一次for循环，处理8行 |
| 441 | for (uint16_t i = 0; i < quotientLoopCount; i++) { | 442 | for (uint16_t i = 0; i < quotientLoopCount; i++) { |
| @@ -447,25 +448,24 @@ private: | |||
| 447 | baseOffset + threeRLoopSize, baseCountOffset + ROW_TWO_OFFSET, | 448 | baseOffset + threeRLoopSize, baseCountOffset + ROW_TWO_OFFSET, |
| 448 | baseCountOffset + ROW_THREE_OFFSET, rem, rowCount, nextRowCount, numScale); | 449 | baseCountOffset + ROW_THREE_OFFSET, rem, rowCount, nextRowCount, numScale); |
| 449 | Add(x1, x1, x2, pregLoop); | 450 | Add(x1, x1, x2, pregLoop); |
| 450 | - DataCopy( | 451 | + StoreAlign( |
| 451 | - ((__local_mem__ float*)binaryAddTmpAddr + (remainderLoopCount + i) * rLoopStride + aLoopOffset), | 452 | + ((__ubuf__ float*)binaryAddTmpAddr + (remainderLoopCount + i) * rLoopStride + aLoopOffset), x1, |
| 452 | - x1, pregLoop); | 453 | + pregLoop); |
| 453 | } | 454 | } |
| 454 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 455 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 455 | BinaryAddVF(binaryAddTmpAddr, rLoopStride, binaryAddKLoop, binaryAddInnerLoop, binaryAddLastLoop, | 456 | BinaryAddVF(binaryAddTmpAddr, rLoopStride, binaryAddKLoop, binaryAddInnerLoop, binaryAddLastLoop, |
| 456 | pregLoop, aLoopOffset, x1, x2, x3, x4); | 457 | pregLoop, aLoopOffset, x1, x2, x3, x4); |
| 457 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + aLoopOffset)); | 458 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + aLoopOffset)); |
| 458 | Muls(x1, x1, scaleCorrection, pregLoop); | 459 | Muls(x1, x1, scaleCorrection, pregLoop); |
| 459 | - DataCopy(((__local_mem__ float*)batchMeanInUbAddr + aLoopOffset), x1, pregLoop); | 460 | + StoreAlign(((__ubuf__ float*)batchMeanInUbAddr + aLoopOffset), x1, pregLoop); |
| 460 | } | 461 | } |
| 461 | } | 462 | } |
| 462 | } | 463 | } |
| 463 | 464 | ||
| 464 | - __aicore__ inline void WelfordFinalizeVarVF(int64_t currentANum, __local_mem__ float* tmpMeanLocal, | 465 | + __aicore__ inline void WelfordFinalizeVarVF(int64_t currentANum, __ubuf__ float* tmpMeanLocal, |
| 465 | - __local_mem__ float* tmpVarLocal, __local_mem__ float* tmpCountLocal, | 466 | + __ubuf__ float* tmpVarLocal, __ubuf__ float* tmpCountLocal, |
| 466 | - __local_mem__ float* batchMeanInUbAddr, | 467 | + __ubuf__ float* batchMeanInUbAddr, __ubuf__ float* batchVarInUbAddr, |
| 467 | - __local_mem__ float* batchVarInUbAddr, | 468 | + __ubuf__ float* binaryAddTmpAddr) |
| 468 | - __local_mem__ float* binaryAddTmpAddr) | ||
| 469 | { | 469 | { |
| 470 | uint16_t rLoopCount = this->rFactor; | 470 | uint16_t rLoopCount = this->rFactor; |
| 471 | uint16_t aLoopCount = this->currentALoopCount; | 471 | uint16_t aLoopCount = this->currentALoopCount; |
| @@ -519,7 +519,7 @@ private: | |||
| 519 | for (uint16_t aIndex = 0; aIndex < aLoopCount; aIndex++) { | 519 | for (uint16_t aIndex = 0; aIndex < aLoopCount; aIndex++) { |
| 520 | uint32_t aLoopOffset = aIndex * VL_F32; | 520 | uint32_t aLoopOffset = aIndex * VL_F32; |
| 521 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0); | 521 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(sreg0); |
| 522 | - DataCopy(saveMean, ((__local_mem__ float*)batchMeanInUbAddr + aLoopOffset)); | 522 | + LoadAlign(saveMean, ((__ubuf__ float*)batchMeanInUbAddr + aLoopOffset)); |
| 523 | for (uint16_t i = 0; i < remainderLoopCount; i++) { | 523 | for (uint16_t i = 0; i < remainderLoopCount; i++) { |
| 524 | uint32_t quotOffset = i * baseLineOffset + aLoopOffset; | 524 | uint32_t quotOffset = i * baseLineOffset + aLoopOffset; |
| 525 | uint32_t remOffset = i * baseLineOffset + remainderOffset + aLoopOffset; | 525 | uint32_t remOffset = i * baseLineOffset + remainderOffset + aLoopOffset; |
| @@ -538,7 +538,7 @@ private: | |||
| 538 | saveMean, rem, nextRow, remNextRow, rowCount, nextRowCount, remCount, | 538 | saveMean, rem, nextRow, remNextRow, rowCount, nextRowCount, remCount, |
| 539 | nextRemCount, rowM2, nextRowM2, remM2, nextRemM2, numScale); | 539 | nextRemCount, rowM2, nextRowM2, remM2, nextRemM2, numScale); |
| 540 | Add(x1, x1, x2, pregLoop); | 540 | Add(x1, x1, x2, pregLoop); |
| 541 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + i * rLoopStride + aLoopOffset), x1, pregLoop); | 541 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + i * rLoopStride + aLoopOffset), x1, pregLoop); |
| 542 | } | 542 | } |
| 543 | // 剩余的前半部分，一次for循环，处理8行 | 543 | // 剩余的前半部分，一次for循环，处理8行 |
| 544 | for (uint16_t i = 0; i < quotientLoopCount; i++) { | 544 | for (uint16_t i = 0; i < quotientLoopCount; i++) { |
| @@ -552,53 +552,52 @@ private: | |||
| 552 | baseCountOffset + ROW_THREE_OFFSET, saveMean, rem, rowCount, nextRowCount, rowM2, | 552 | baseCountOffset + ROW_THREE_OFFSET, saveMean, rem, rowCount, nextRowCount, rowM2, |
| 553 | remM2, numScale); | 553 | remM2, numScale); |
| 554 | Add(x1, x1, x2, pregLoop); | 554 | Add(x1, x1, x2, pregLoop); |
| 555 | - DataCopy( | 555 | + StoreAlign( |
| 556 | - ((__local_mem__ float*)binaryAddTmpAddr + (remainderLoopCount + i) * rLoopStride + aLoopOffset), | 556 | + ((__ubuf__ float*)binaryAddTmpAddr + (remainderLoopCount + i) * rLoopStride + aLoopOffset), x1, |
| 557 | - x1, pregLoop); | 557 | + pregLoop); |
| 558 | } | 558 | } |
| 559 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 559 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 560 | BinaryAddVF(binaryAddTmpAddr, rLoopStride, binaryAddKLoop, binaryAddInnerLoop, binaryAddLastLoop, | 560 | BinaryAddVF(binaryAddTmpAddr, rLoopStride, binaryAddKLoop, binaryAddInnerLoop, binaryAddLastLoop, |
| 561 | pregLoop, aLoopOffset, x1, x2, x3, x4); | 561 | pregLoop, aLoopOffset, x1, x2, x3, x4); |
| 562 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + aLoopOffset)); | 562 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + aLoopOffset)); |
| 563 | - DataCopy(((__local_mem__ float*)batchVarInUbAddr + aLoopOffset), x1, pregLoop); | 563 | + StoreAlign(((__ubuf__ float*)batchVarInUbAddr + aLoopOffset), x1, pregLoop); |
| 564 | } | 564 | } |
| 565 | } | 565 | } |
| 566 | } | 566 | } |
| 567 | 567 | ||
| 568 | - __aicore__ inline void BinaryAddVF(__local_mem__ float* binaryAddTmpAddr, uint32_t rLoopStride, | 568 | + __aicore__ inline void BinaryAddVF(__ubuf__ float* binaryAddTmpAddr, uint32_t rLoopStride, uint16_t binaryAddKLoop, |
| 569 | - uint16_t binaryAddKLoop, uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop, | 569 | + uint16_t binaryAddInnerLoop, uint16_t binaryAddLastLoop, MaskReg& pregLoop, |
| 570 | - MaskReg& pregLoop, uint32_t offset, RegTensor<float>& x1, RegTensor<float>& x2, | 570 | + uint32_t offset, RegTensor<float>& x1, RegTensor<float>& x2, |
| 571 | RegTensor<float>& x3, RegTensor<float>& x4) | 571 | RegTensor<float>& x3, RegTensor<float>& x4) |
| 572 | { | 572 | { |
| 573 | uint16_t curBinaryAddInnerLoop = binaryAddInnerLoop; | 573 | uint16_t curBinaryAddInnerLoop = binaryAddInnerLoop; |
| 574 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { | 574 | for (uint16_t i = 0; i < binaryAddKLoop; i++) { |
| 575 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET; | 575 | curBinaryAddInnerLoop = curBinaryAddInnerLoop / ROW_FOUR_OFFSET; |
| 576 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { | 576 | for (uint16_t j = 0; j < curBinaryAddInnerLoop; j++) { |
| 577 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset)); | 577 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET) * rLoopStride + offset)); |
| 578 | - DataCopy(x2, | 578 | + LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset)); |
| 579 | - ((__local_mem__ float*)binaryAddTmpAddr + (j * ROW_FOUR_OFFSET + 1) * rLoopStride + offset)); | ||
| 580 | Add(x1, x1, x2, pregLoop); | 579 | Add(x1, x1, x2, pregLoop); |
| 581 | - DataCopy(x3, ((__local_mem__ float*)binaryAddTmpAddr + | 580 | + LoadAlign(x3, ((__ubuf__ float*)binaryAddTmpAddr + |
| 582 | - (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset)); | 581 | + (j * ROW_FOUR_OFFSET + ROW_TWO_OFFSET) * rLoopStride + offset)); |
| 583 | - DataCopy(x4, ((__local_mem__ float*)binaryAddTmpAddr + | 582 | + LoadAlign(x4, ((__ubuf__ float*)binaryAddTmpAddr + |
| 584 | - (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset)); | 583 | + (j * ROW_FOUR_OFFSET + ROW_THREE_OFFSET) * rLoopStride + offset)); |
| 585 | Add(x3, x3, x4, pregLoop); | 584 | Add(x3, x3, x4, pregLoop); |
| 586 | Add(x1, x1, x3, pregLoop); | 585 | Add(x1, x1, x3, pregLoop); |
| 587 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop); | 586 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + j * rLoopStride + offset), x1, pregLoop); |
| 588 | } | 587 | } |
| 589 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 588 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 590 | } | 589 | } |
| 591 | for (uint16_t i = 0; i < binaryAddLastLoop; i++) { | 590 | for (uint16_t i = 0; i < binaryAddLastLoop; i++) { |
| 592 | - DataCopy(x1, ((__local_mem__ float*)binaryAddTmpAddr + offset)); | 591 | + LoadAlign(x1, ((__ubuf__ float*)binaryAddTmpAddr + offset)); |
| 593 | - DataCopy(x2, ((__local_mem__ float*)binaryAddTmpAddr + rLoopStride + offset)); | 592 | + LoadAlign(x2, ((__ubuf__ float*)binaryAddTmpAddr + rLoopStride + offset)); |
| 594 | Add(x1, x1, x2, pregLoop); | 593 | Add(x1, x1, x2, pregLoop); |
| 595 | - DataCopy(((__local_mem__ float*)binaryAddTmpAddr + offset), x1, pregLoop); | 594 | + StoreAlign(((__ubuf__ float*)binaryAddTmpAddr + offset), x1, pregLoop); |
| 596 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); | 595 | LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>(); |
| 597 | } | 596 | } |
| 598 | } | 597 | } |
| 599 | 598 | ||
| 600 | - __aicore__ inline void TwoRowAddForMeanWithTail(RegTensor<float>& dst, __local_mem__ float* input, | 599 | + __aicore__ inline void TwoRowAddForMeanWithTail(RegTensor<float>& dst, __ubuf__ float* input, |
| 601 | - __local_mem__ float* tCount, MaskReg& preg, uint32_t offset1, | 600 | + __ubuf__ float* tCount, MaskReg& preg, uint32_t offset1, |
| 602 | uint32_t offset2, uint32_t offset3, uint32_t offset4, | 601 | uint32_t offset2, uint32_t offset3, uint32_t offset4, |
| 603 | uint32_t offset5, uint32_t offset6, uint32_t offset7, | 602 | uint32_t offset5, uint32_t offset6, uint32_t offset7, |
| 604 | uint32_t offset8, RegTensor<float>& rem, RegTensor<float>& nextRow, | 603 | uint32_t offset8, RegTensor<float>& rem, RegTensor<float>& nextRow, |
| @@ -606,19 +605,19 @@ private: | |||
| 606 | RegTensor<float>& remCount, RegTensor<float>& nextRowCount, | 605 | RegTensor<float>& remCount, RegTensor<float>& nextRowCount, |
| 607 | RegTensor<float>& remNextRowCount, float n) | 606 | RegTensor<float>& remNextRowCount, float n) |
| 608 | { | 607 | { |
| 609 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 608 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 610 | - DataCopy(rem, ((__local_mem__ float*)(input) + (offset2))); | 609 | + LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2))); |
| 611 | - DataCopy<float, LoadDist::DIST_BRC_B32>(dstCount, ((__local_mem__ float*)(tCount) + (offset5))); | 610 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(dstCount, ((__ubuf__ float*)(tCount) + (offset5))); |
| 612 | - DataCopy<float, LoadDist::DIST_BRC_B32>(remCount, ((__local_mem__ float*)(tCount) + (offset6))); | 611 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(remCount, ((__ubuf__ float*)(tCount) + (offset6))); |
| 613 | Mul(dst, dst, dstCount, preg); | 612 | Mul(dst, dst, dstCount, preg); |
| 614 | Mul(rem, rem, remCount, preg); | 613 | Mul(rem, rem, remCount, preg); |
| 615 | Muls(dst, dst, n, preg); | 614 | Muls(dst, dst, n, preg); |
| 616 | Muls(rem, rem, n, preg); | 615 | Muls(rem, rem, n, preg); |
| 617 | Add(dst, dst, rem, preg); | 616 | Add(dst, dst, rem, preg); |
| 618 | - DataCopy(nextRow, ((__local_mem__ float*)(input) + (offset3))); | 617 | + LoadAlign(nextRow, ((__ubuf__ float*)(input) + (offset3))); |
| 619 | - DataCopy(remNextRow, ((__local_mem__ float*)(input) + (offset4))); | 618 | + LoadAlign(remNextRow, ((__ubuf__ float*)(input) + (offset4))); |
| 620 | - DataCopy<float, LoadDist::DIST_BRC_B32>(nextRowCount, ((__local_mem__ float*)(tCount) + (offset7))); | 619 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(nextRowCount, ((__ubuf__ float*)(tCount) + (offset7))); |
| 621 | - DataCopy<float, LoadDist::DIST_BRC_B32>(remNextRowCount, ((__local_mem__ float*)(tCount) + (offset8))); | 620 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(remNextRowCount, ((__ubuf__ float*)(tCount) + (offset8))); |
| 622 | Mul(nextRow, nextRow, nextRowCount, preg); | 621 | Mul(nextRow, nextRow, nextRowCount, preg); |
| 623 | Mul(remNextRow, remNextRow, remNextRowCount, preg); | 622 | Mul(remNextRow, remNextRow, remNextRowCount, preg); |
| 624 | Muls(nextRow, nextRow, n, preg); | 623 | Muls(nextRow, nextRow, n, preg); |
| @@ -627,15 +626,15 @@ private: | |||
| 627 | Add(dst, dst, nextRow, preg); | 626 | Add(dst, dst, nextRow, preg); |
| 628 | } | 627 | } |
| 629 | 628 | ||
| 630 | - __aicore__ inline void TwoRowAddForMean(RegTensor<float>& dst, __local_mem__ float* input, | 629 | + __aicore__ inline void TwoRowAddForMean(RegTensor<float>& dst, __ubuf__ float* input, __ubuf__ float* tCount, |
| 631 | - __local_mem__ float* tCount, MaskReg& preg, uint32_t offset1, | 630 | + MaskReg& preg, uint32_t offset1, uint32_t offset2, uint32_t offset5, |
| 632 | - uint32_t offset2, uint32_t offset5, uint32_t offset6, RegTensor<float>& rem, | 631 | + uint32_t offset6, RegTensor<float>& rem, RegTensor<float>& dstCount, |
| 633 | - RegTensor<float>& dstCount, RegTensor<float>& remCount, float n) | 632 | + RegTensor<float>& remCount, float n) |
| 634 | { | 633 | { |
| 635 | - DataCopy(dst, ((__local_mem__ float*)(input) + (offset1))); | 634 | + LoadAlign(dst, ((__ubuf__ float*)(input) + (offset1))); |
| 636 | - DataCopy(rem, ((__local_mem__ float*)(input) + (offset2))); | 635 | + LoadAlign(rem, ((__ubuf__ float*)(input) + (offset2))); |
| 637 | - DataCopy<float, LoadDist::DIST_BRC_B32>(dstCount, ((__local_mem__ float*)(tCount) + (offset5))); | 636 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(dstCount, ((__ubuf__ float*)(tCount) + (offset5))); |
| 638 | - DataCopy<float, LoadDist::DIST_BRC_B32>(remCount, ((__local_mem__ float*)(tCount) + (offset6))); | 637 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(remCount, ((__ubuf__ float*)(tCount) + (offset6))); |
| 639 | Mul(dst, dst, dstCount, preg); | 638 | Mul(dst, dst, dstCount, preg); |
| 640 | Mul(rem, rem, remCount, preg); | 639 | Mul(rem, rem, remCount, preg); |
| 641 | Muls(dst, dst, n, preg); | 640 | Muls(dst, dst, n, preg); |
| @@ -644,43 +643,43 @@ private: | |||
| 644 | } | 643 | } |
| 645 | 644 | ||
| 646 | __aicore__ inline void TwoRowAddForVarWithTail( | 645 | __aicore__ inline void TwoRowAddForVarWithTail( |
| 647 | - RegTensor<float>& dst, __local_mem__ float* tmpMean, __local_mem__ float* tmpM2, __local_mem__ float* tCount, | 646 | + RegTensor<float>& dst, __ubuf__ float* tmpMean, __ubuf__ float* tmpM2, __ubuf__ float* tCount, MaskReg& preg, |
| 648 | - MaskReg& preg, uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, uint32_t offset5, | 647 | + uint32_t offset1, uint32_t offset2, uint32_t offset3, uint32_t offset4, uint32_t offset5, uint32_t offset6, |
| 649 | - uint32_t offset6, uint32_t offset7, uint32_t offset8, RegTensor<float>& mean, RegTensor<float>& rem, | 648 | + uint32_t offset7, uint32_t offset8, RegTensor<float>& mean, RegTensor<float>& rem, RegTensor<float>& nextRow, |
| 650 | - RegTensor<float>& nextRow, RegTensor<float>& remNextRow, RegTensor<float>& dstCount, RegTensor<float>& remCount, | 649 | + RegTensor<float>& remNextRow, RegTensor<float>& dstCount, RegTensor<float>& remCount, |
| 651 | RegTensor<float>& nextRowCount, RegTensor<float>& remNextRowCount, RegTensor<float>& dstM2, | 650 | RegTensor<float>& nextRowCount, RegTensor<float>& remNextRowCount, RegTensor<float>& dstM2, |
| 652 | RegTensor<float>& remM2, RegTensor<float>& nextRowM2, RegTensor<float>& remNextRowM2, float n) | 651 | RegTensor<float>& remM2, RegTensor<float>& nextRowM2, RegTensor<float>& remNextRowM2, float n) |
| 653 | { | 652 | { |
| 654 | - DataCopy(dst, ((__local_mem__ float*)(tmpMean) + (offset1))); | 653 | + LoadAlign(dst, ((__ubuf__ float*)(tmpMean) + (offset1))); |
| 655 | - DataCopy(rem, ((__local_mem__ float*)(tmpMean) + (offset2))); | 654 | + LoadAlign(rem, ((__ubuf__ float*)(tmpMean) + (offset2))); |
| 656 | - DataCopy<float, LoadDist::DIST_BRC_B32>(dstCount, ((__local_mem__ float*)(tCount) + (offset5))); | 655 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(dstCount, ((__ubuf__ float*)(tCount) + (offset5))); |
| 657 | - DataCopy<float, LoadDist::DIST_BRC_B32>(remCount, ((__local_mem__ float*)(tCount) + (offset6))); | 656 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(remCount, ((__ubuf__ float*)(tCount) + (offset6))); |
| 658 | Sub(dst, dst, mean, preg); | 657 | Sub(dst, dst, mean, preg); |
| 659 | Mul(dst, dst, dst, preg); | 658 | Mul(dst, dst, dst, preg); |
| 660 | Sub(rem, rem, mean, preg); | 659 | Sub(rem, rem, mean, preg); |
| 661 | Mul(rem, rem, rem, preg); | 660 | Mul(rem, rem, rem, preg); |
| 662 | Mul(dst, dst, dstCount, preg); | 661 | Mul(dst, dst, dstCount, preg); |
| 663 | Mul(rem, rem, remCount, preg); | 662 | Mul(rem, rem, remCount, preg); |
| 664 | - DataCopy(dstM2, ((__local_mem__ float*)(tmpM2) + (offset1))); | 663 | + LoadAlign(dstM2, ((__ubuf__ float*)(tmpM2) + (offset1))); |
| 665 | - DataCopy(remM2, ((__local_mem__ float*)(tmpM2) + (offset2))); | 664 | + LoadAlign(remM2, ((__ubuf__ float*)(tmpM2) + (offset2))); |
| 666 | Add(dst, dstM2, dst, preg); | 665 | Add(dst, dstM2, dst, preg); |
| 667 | Muls(dst, dst, n, preg); | 666 | Muls(dst, dst, n, preg); |
| 668 | Add(rem, remM2, rem, preg); | 667 | Add(rem, remM2, rem, preg); |
| 669 | Muls(rem, rem, n, preg); | 668 | Muls(rem, rem, n, preg); |
| 670 | Add(dst, dst, rem, preg); | 669 | Add(dst, dst, rem, preg); |
| 671 | 670 | ||
| 672 | - DataCopy(nextRow, ((__local_mem__ float*)(tmpMean) + (offset3))); | 671 | + LoadAlign(nextRow, ((__ubuf__ float*)(tmpMean) + (offset3))); |
| 673 | - DataCopy(remNextRow, ((__local_mem__ float*)(tmpMean) + (offset4))); | 672 | + LoadAlign(remNextRow, ((__ubuf__ float*)(tmpMean) + (offset4))); |
| 674 | - DataCopy<float, LoadDist::DIST_BRC_B32>(nextRowCount, ((__local_mem__ float*)(tCount) + (offset7))); | 673 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(nextRowCount, ((__ubuf__ float*)(tCount) + (offset7))); |
| 675 | - DataCopy<float, LoadDist::DIST_BRC_B32>(remNextRowCount, ((__local_mem__ float*)(tCount) + (offset8))); | 674 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(remNextRowCount, ((__ubuf__ float*)(tCount) + (offset8))); |
| 676 | Sub(nextRow, nextRow, mean, preg); | 675 | Sub(nextRow, nextRow, mean, preg); |
| 677 | Mul(nextRow, nextRow, nextRow, preg); | 676 | Mul(nextRow, nextRow, nextRow, preg); |
| 678 | Sub(remNextRow, remNextRow, mean, preg); | 677 | Sub(remNextRow, remNextRow, mean, preg); |
| 679 | Mul(remNextRow, remNextRow, remNextRow, preg); | 678 | Mul(remNextRow, remNextRow, remNextRow, preg); |
| 680 | Mul(nextRow, nextRow, nextRowCount, preg); | 679 | Mul(nextRow, nextRow, nextRowCount, preg); |
| 681 | Mul(remNextRow, remNextRow, remNextRowCount, preg); | 680 | Mul(remNextRow, remNextRow, remNextRowCount, preg); |
| 682 | - DataCopy(nextRowM2, ((__local_mem__ float*)(tmpM2) + (offset3))); | 681 | + LoadAlign(nextRowM2, ((__ubuf__ float*)(tmpM2) + (offset3))); |
| 683 | - DataCopy(remNextRowM2, ((__local_mem__ float*)(tmpM2) + (offset4))); | 682 | + LoadAlign(remNextRowM2, ((__ubuf__ float*)(tmpM2) + (offset4))); |
| 684 | Add(nextRow, nextRowM2, nextRow, preg); | 683 | Add(nextRow, nextRowM2, nextRow, preg); |
| 685 | Muls(nextRow, nextRow, n, preg); | 684 | Muls(nextRow, nextRow, n, preg); |
| 686 | Add(remNextRow, remNextRowM2, remNextRow, preg); | 685 | Add(remNextRow, remNextRowM2, remNextRow, preg); |
| @@ -690,25 +689,25 @@ private: | |||
| 690 | Add(dst, dst, nextRow, preg); | 689 | Add(dst, dst, nextRow, preg); |
| 691 | } | 690 | } |
| 692 | 691 | ||
| 693 | - __aicore__ inline void TwoRowAddForVar(RegTensor<float>& dst, __local_mem__ float* tmpMean, | 692 | + __aicore__ inline void TwoRowAddForVar(RegTensor<float>& dst, __ubuf__ float* tmpMean, __ubuf__ float* tmpM2, |
| 694 | - __local_mem__ float* tmpM2, __local_mem__ float* tCount, MaskReg& preg, | 693 | + __ubuf__ float* tCount, MaskReg& preg, uint32_t offset1, uint32_t offset2, |
| 695 | - uint32_t offset1, uint32_t offset2, uint32_t offset5, uint32_t offset6, | 694 | + uint32_t offset5, uint32_t offset6, RegTensor<float>& mean, |
| 696 | - RegTensor<float>& mean, RegTensor<float>& rem, RegTensor<float>& dstCount, | 695 | + RegTensor<float>& rem, RegTensor<float>& dstCount, |
| 697 | RegTensor<float>& remCount, RegTensor<float>& dstM2, RegTensor<float>& remM2, | 696 | RegTensor<float>& remCount, RegTensor<float>& dstM2, RegTensor<float>& remM2, |
| 698 | float n) | 697 | float n) |
| 699 | { | 698 | { |
| 700 | - DataCopy(dst, ((__local_mem__ float*)(tmpMean) + (offset1))); | 699 | + LoadAlign(dst, ((__ubuf__ float*)(tmpMean) + (offset1))); |
| 701 | - DataCopy(rem, ((__local_mem__ float*)(tmpMean) + (offset2))); | 700 | + LoadAlign(rem, ((__ubuf__ float*)(tmpMean) + (offset2))); |
| 702 | - DataCopy<float, LoadDist::DIST_BRC_B32>(dstCount, ((__local_mem__ float*)(tCount) + (offset5))); | 701 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(dstCount, ((__ubuf__ float*)(tCount) + (offset5))); |
| 703 | - DataCopy<float, LoadDist::DIST_BRC_B32>(remCount, ((__local_mem__ float*)(tCount) + (offset6))); | 702 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(remCount, ((__ubuf__ float*)(tCount) + (offset6))); |
| 704 | Sub(dst, dst, mean, preg); | 703 | Sub(dst, dst, mean, preg); |
| 705 | Mul(dst, dst, dst, preg); | 704 | Mul(dst, dst, dst, preg); |
| 706 | Sub(rem, rem, mean, preg); | 705 | Sub(rem, rem, mean, preg); |
| 707 | Mul(rem, rem, rem, preg); | 706 | Mul(rem, rem, rem, preg); |
| 708 | Mul(dst, dst, dstCount, preg); | 707 | Mul(dst, dst, dstCount, preg); |
| 709 | Mul(rem, rem, remCount, preg); | 708 | Mul(rem, rem, remCount, preg); |
| 710 | - DataCopy(dstM2, ((__local_mem__ float*)(tmpM2) + (offset1))); | 709 | + LoadAlign(dstM2, ((__ubuf__ float*)(tmpM2) + (offset1))); |
| 711 | - DataCopy(remM2, ((__local_mem__ float*)(tmpM2) + (offset2))); | 710 | + LoadAlign(remM2, ((__ubuf__ float*)(tmpM2) + (offset2))); |
| 712 | Add(dst, dstM2, dst, preg); | 711 | Add(dst, dstM2, dst, preg); |
| 713 | Muls(dst, dst, n, preg); | 712 | Muls(dst, dst, n, preg); |
| 714 | Add(rem, remM2, rem, preg); | 713 | Add(rem, remM2, rem, preg); |
| @@ -716,8 +715,7 @@ private: | |||
| 716 | Add(dst, dst, rem, preg); | 715 | Add(dst, dst, rem, preg); |
| 717 | } | 716 | } |
| 718 | 717 | ||
| 719 | - __aicore__ inline void ComputeRstd(int64_t currentANum, __local_mem__ float* rstdLocal, | 718 | + __aicore__ inline void ComputeRstd(int64_t currentANum, __ubuf__ float* rstdLocal, __ubuf__ float* batchVarInUbAddr) |
| 720 | - __local_mem__ float* batchVarInUbAddr) | ||
| 721 | { | 719 | { |
| 722 | uint16_t aLoop = currentALoopCount; | 720 | uint16_t aLoop = currentALoopCount; |
| 723 | __VEC_SCOPE__ | 721 | __VEC_SCOPE__ |
| @@ -753,7 +751,7 @@ private: | |||
| 753 | Duplicate(s, float(1.0), pregLoop); | 751 | Duplicate(s, float(1.0), pregLoop); |
| 754 | 752 | ||
| 755 | // rstd | 753 | // rstd |
| 756 | - DataCopy(var, ((__local_mem__ float*)batchVarInUbAddr + a * VL_F32)); | 754 | + LoadAlign(var, ((__ubuf__ float*)batchVarInUbAddr + a * VL_F32)); |
| 757 | Adds(var, var, epsilon, pregLoop); | 755 | Adds(var, var, epsilon, pregLoop); |
| 758 | Div(r, one, var, pregLoop); | 756 | Div(r, one, var, pregLoop); |
| 759 | Sqrt(y, r, pregLoop); | 757 | Sqrt(y, r, pregLoop); |
| @@ -768,22 +766,22 @@ private: | |||
| 768 | Mula(s, var, r, pregLoop); // s + x * t | 766 | Mula(s, var, r, pregLoop); // s + x * t |
| 769 | Mul(s, s, rstd, pregLoop); // e * y | 767 | Mul(s, s, rstd, pregLoop); // e * y |
| 770 | Mula(rstd, s, scalar1, pregLoop); // y + y * e * 0.5 | 768 | Mula(rstd, s, scalar1, pregLoop); // y + y * e * 0.5 |
| 771 | - CompareScalar(cmpRegZero, var, POS_INF, pregLoop); | 769 | + Compares(cmpRegZero, var, POS_INF, pregLoop); |
| 772 | Select(rstd, scalarZero, rstd, cmpRegZero); | 770 | Select(rstd, scalarZero, rstd, cmpRegZero); |
| 773 | - CompareScalar(cmpRegInf, var, float(0.0), pregLoop); | 771 | + Compares(cmpRegInf, var, float(0.0), pregLoop); |
| 774 | Select(rstd, scalarInf, rstd, cmpRegInf); | 772 | Select(rstd, scalarInf, rstd, cmpRegInf); |
| 775 | - DataCopy(((__local_mem__ float*)rstdLocal + a * VL_F32), rstd, pregLoop); | 773 | + StoreAlign(((__ubuf__ float*)rstdLocal + a * VL_F32), rstd, pregLoop); |
| 776 | } | 774 | } |
| 777 | } | 775 | } |
| 778 | } | 776 | } |
| 779 | 777 | ||
| 780 | __aicore__ inline void Normalize(int64_t curA0Idx, int64_t curA1Idx, int64_t currentANum, | 778 | __aicore__ inline void Normalize(int64_t curA0Idx, int64_t curA1Idx, int64_t currentANum, |
| 781 | - __local_mem__ float* batchMeanInUbAddr, __local_mem__ float* rstdLocal) | 779 | + __ubuf__ float* batchMeanInUbAddr, __ubuf__ float* rstdLocal) |
| 782 | { | 780 | { |
| 783 | LocalTensor<T_BETA> betaInUb = betaQueue.template DeQue<T_BETA>(); | 781 | LocalTensor<T_BETA> betaInUb = betaQueue.template DeQue<T_BETA>(); |
| 784 | LocalTensor<T_BETA> gammaInUb = gammaQueue.template DeQue<T_BETA>(); | 782 | LocalTensor<T_BETA> gammaInUb = gammaQueue.template DeQue<T_BETA>(); |
| 785 | - __local_mem__ T_BETA* betaInUbAddr = (__local_mem__ T_BETA*)betaInUb.GetPhyAddr(); | 783 | + __ubuf__ T_BETA* betaInUbAddr = (__ubuf__ T_BETA*)betaInUb.GetPhyAddr(); |
| 786 | - __local_mem__ T_BETA* gammaInUbAddr = (__local_mem__ T_BETA*)gammaInUb.GetPhyAddr(); | 784 | + __ubuf__ T_BETA* gammaInUbAddr = (__ubuf__ T_BETA*)gammaInUb.GetPhyAddr(); |
| 787 | int64_t quotient = (this->r + this->rFactor - 1) / this->rFactor; | 785 | int64_t quotient = (this->r + this->rFactor - 1) / this->rFactor; |
| 788 | for (int64_t rLoopIdx = 0; rLoopIdx < quotient; rLoopIdx++) { | 786 | for (int64_t rLoopIdx = 0; rLoopIdx < quotient; rLoopIdx++) { |
| 789 | int64_t copyXOffset = curA1Idx * this->r * this->a0 + rLoopIdx * this->rFactor * this->a0 + | 787 | int64_t copyXOffset = curA1Idx * this->r * this->a0 + rLoopIdx * this->rFactor * this->a0 + |
| @@ -799,14 +797,14 @@ private: | |||
| 799 | gammaQueue.FreeTensor(gammaInUb); | 797 | gammaQueue.FreeTensor(gammaInUb); |
| 800 | } | 798 | } |
| 801 | 799 | ||
| 802 | - __aicore__ inline void NormalizeVF(int64_t currentR, int64_t currentANum, __local_mem__ float* batchMeanInUbAddr, | 800 | + __aicore__ inline void NormalizeVF(int64_t currentR, int64_t currentANum, __ubuf__ float* batchMeanInUbAddr, |
| 803 | - __local_mem__ float* rstdLocal, __local_mem__ T_BETA* betaInUbAddr, | 801 | + __ubuf__ float* rstdLocal, __ubuf__ T_BETA* betaInUbAddr, |
| 804 | - __local_mem__ T_BETA* gammaInUbAddr) | 802 | + __ubuf__ T_BETA* gammaInUbAddr) |
| 805 | { | 803 | { |
| 806 | LocalTensor<T> xInUb = xQueue.DeQue<T>(); | 804 | LocalTensor<T> xInUb = xQueue.DeQue<T>(); |
| 807 | LocalTensor<T> yInUb = yQueue.AllocTensor<T>(); | 805 | LocalTensor<T> yInUb = yQueue.AllocTensor<T>(); |
| 808 | - __local_mem__ T* xInUbAddr = (__local_mem__ T*)xInUb.GetPhyAddr(); | 806 | + __ubuf__ T* xInUbAddr = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 809 | - __local_mem__ T* yInUbAddr = (__local_mem__ T*)yInUb.GetPhyAddr(); | 807 | + __ubuf__ T* yInUbAddr = (__ubuf__ T*)yInUb.GetPhyAddr(); |
| 810 | 808 | ||
| 811 | uint16_t rLoopCount = currentR; | 809 | uint16_t rLoopCount = currentR; |
| 812 | uint16_t aLoopCount = currentALoopCount; | 810 | uint16_t aLoopCount = currentALoopCount; |
| @@ -830,8 +828,8 @@ private: | |||
| 830 | 828 | ||
| 831 | LoadOneTensorForDtypeT(betaInUbAddr, beta, pregLoop, aLoopOffset); | 829 | LoadOneTensorForDtypeT(betaInUbAddr, beta, pregLoop, aLoopOffset); |
| 832 | LoadOneTensorForDtypeT(gammaInUbAddr, gamma, pregLoop, aLoopOffset); | 830 | LoadOneTensorForDtypeT(gammaInUbAddr, gamma, pregLoop, aLoopOffset); |
| 833 | - DataCopy(mean, (__local_mem__ float*)batchMeanInUbAddr + aLoopOffset); | 831 | + LoadAlign(mean, (__ubuf__ float*)batchMeanInUbAddr + aLoopOffset); |
| 834 | - DataCopy(rstd, (__local_mem__ float*)rstdLocal + aLoopOffset); | 832 | + LoadAlign(rstd, (__ubuf__ float*)rstdLocal + aLoopOffset); |
| 835 | for (uint16_t rIndex = 0; rIndex < rLoopCount; rIndex++) { | 833 | for (uint16_t rIndex = 0; rIndex < rLoopCount; rIndex++) { |
| 836 | LoadOneTensorForDtypeT(xInUbAddr, x2, pregLoop, rIndex * rLoopStride + aLoopOffset); | 834 | LoadOneTensorForDtypeT(xInUbAddr, x2, pregLoop, rIndex * rLoopStride + aLoopOffset); |
| 837 | Sub(x2, x2, mean, pregLoop); | 835 | Sub(x2, x2, mean, pregLoop); |
| @@ -841,16 +839,15 @@ private: | |||
| 841 | if constexpr (IsSameType<T, half>::value) { | 839 | if constexpr (IsSameType<T, half>::value) { |
| 842 | RegTensor<half> yFp16; | 840 | RegTensor<half> yFp16; |
| 843 | Cast<half, float, castTraitB322B16>(yFp16, y2, pregLoop); | 841 | Cast<half, float, castTraitB322B16>(yFp16, y2, pregLoop); |
| 844 | - DataCopy<half, StoreDist::DIST_PACK_B32>( | 842 | + StoreAlign<half, StoreDist::DIST_PACK_B32>( |
| 845 | - ((__local_mem__ half*)yInUbAddr + rIndex * rLoopStride + aLoopOffset), yFp16, pregLoop); | 843 | + ((__ubuf__ half*)yInUbAddr + rIndex * rLoopStride + aLoopOffset), yFp16, pregLoop); |
| 846 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { | 844 | } else if constexpr (IsSameType<T, bfloat16_t>::value) { |
| 847 | RegTensor<bfloat16_t> xBf16; | 845 | RegTensor<bfloat16_t> xBf16; |
| 848 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, y2, pregLoop); | 846 | Cast<bfloat16_t, float, castTraitB322B16>(xBf16, y2, pregLoop); |
| 849 | - DataCopy<bfloat16_t, StoreDist::DIST_PACK_B32>( | 847 | + StoreAlign<bfloat16_t, StoreDist::DIST_PACK_B32>( |
| 850 | - ((__local_mem__ bfloat16_t*)yInUbAddr + rIndex * rLoopStride + aLoopOffset), xBf16, | 848 | + ((__ubuf__ bfloat16_t*)yInUbAddr + rIndex * rLoopStride + aLoopOffset), xBf16, pregLoop); |
| 851 | - pregLoop); | ||
| 852 | } else { | 849 | } else { |
| 853 | - DataCopy(((__local_mem__ float*)yInUbAddr + rIndex * rLoopStride + aLoopOffset), y2, pregLoop); | 850 | + StoreAlign(((__ubuf__ float*)yInUbAddr + rIndex * rLoopStride + aLoopOffset), y2, pregLoop); |
| 854 | } | 851 | } |
| 855 | } | 852 | } |
| 856 | } | 853 | } |
| @@ -862,10 +859,10 @@ private: | |||
| 862 | __aicore__ inline void CastMeanVar(uint64_t currentANum, LocalTensor<float> batchMeanInUb, | 859 | __aicore__ inline void CastMeanVar(uint64_t currentANum, LocalTensor<float> batchMeanInUb, |
| 863 | LocalTensor<float> batchVarInUb) | 860 | LocalTensor<float> batchVarInUb) |
| 864 | { | 861 | { |
| 865 | - __local_mem__ float* batchMeanInAddr = (__local_mem__ float*)batchMeanInUb.GetPhyAddr(); | 862 | + __ubuf__ float* batchMeanInAddr = (__ubuf__ float*)batchMeanInUb.GetPhyAddr(); |
| 866 | - __local_mem__ float* batchVarInAddr = (__local_mem__ float*)batchVarInUb.GetPhyAddr(); | 863 | + __ubuf__ float* batchVarInAddr = (__ubuf__ float*)batchVarInUb.GetPhyAddr(); |
| 867 | - __local_mem__ T_MEAN* batchMeanOutAddr = (__local_mem__ T_MEAN*)batchMeanInUb.GetPhyAddr(); | 864 | + __ubuf__ T_MEAN* batchMeanOutAddr = (__ubuf__ T_MEAN*)batchMeanInUb.GetPhyAddr(); |
| 868 | - __local_mem__ T_MEAN* batchVarOutAddr = (__local_mem__ T_MEAN*)batchVarInUb.GetPhyAddr(); | 865 | + __ubuf__ T_MEAN* batchVarOutAddr = (__ubuf__ T_MEAN*)batchVarInUb.GetPhyAddr(); |
| 869 | 866 | ||
| 870 | uint32_t castCount = static_cast<uint32_t>(currentANum); | 867 | uint32_t castCount = static_cast<uint32_t>(currentANum); |
| 871 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); | 868 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); |
| @@ -878,14 +875,14 @@ private: | |||
| 878 | MicroAPI::MaskReg pregLoop; | 875 | MicroAPI::MaskReg pregLoop; |
| 879 | for (uint16_t i = 0; i < castLoops; i++) { | 876 | for (uint16_t i = 0; i < castLoops; i++) { |
| 880 | pregLoop = MicroAPI::UpdateMask<float>(castCount); | 877 | pregLoop = MicroAPI::UpdateMask<float>(castCount); |
| 881 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); | 878 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); |
| 882 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_var, batchVarInAddr + VL_F32 * i); | 879 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_var, batchVarInAddr + VL_F32 * i); |
| 883 | Cast<T_MEAN, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 880 | Cast<T_MEAN, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 884 | Cast<T_MEAN, float, castTraitB322B16>(output_var, input_var, pregLoop); | 881 | Cast<T_MEAN, float, castTraitB322B16>(output_var, input_var, pregLoop); |
| 885 | - DataCopy<T_MEAN, StoreDist::DIST_PACK_B32>(((__local_mem__ T_MEAN*)batchMeanOutAddr + i * VL_MEAN), | 882 | + StoreAlign<T_MEAN, StoreDist::DIST_PACK_B32>(((__ubuf__ T_MEAN*)batchMeanOutAddr + i * VL_MEAN), |
| 886 | - output_mean, pregLoop); | 883 | + output_mean, pregLoop); |
| 887 | - DataCopy<T_MEAN, StoreDist::DIST_PACK_B32>(((__local_mem__ T_MEAN*)batchVarOutAddr + i * VL_MEAN), | 884 | + StoreAlign<T_MEAN, StoreDist::DIST_PACK_B32>(((__ubuf__ T_MEAN*)batchVarOutAddr + i * VL_MEAN), |
| 888 | - output_var, pregLoop); | 885 | + output_var, pregLoop); |
| 889 | } | 886 | } |
| 890 | } | 887 | } |
| 891 | } | 888 | } |
| @@ -1002,20 +999,6 @@ private: | |||
| 1002 | 999 | ||
| 1003 | static constexpr float POS_INF = 3.40282366920938E+38; | 1000 | static constexpr float POS_INF = 3.40282366920938E+38; |
| 1004 | 1001 | ||
| 1005 | - constexpr static AscendC::MicroAPI::CastTrait castTraitB162B32 = { | ||
| 1006 | - AscendC::MicroAPI::RegLayout::ZERO, | ||
| 1007 | - AscendC::MicroAPI::SatMode::UNKNOWN, | ||
| 1008 | - AscendC::MicroAPI::MaskMergeMode::ZEROING, | ||
| 1009 | - AscendC::RoundMode::UNKNOWN, | ||
| 1010 | - }; | ||
| 1011 | - | ||
| 1012 | - constexpr static AscendC::MicroAPI::CastTrait castTraitB322B16 = { | ||
| 1013 | - AscendC::MicroAPI::RegLayout::ZERO, | ||
| 1014 | - AscendC::MicroAPI::SatMode::NO_SAT, | ||
| 1015 | - AscendC::MicroAPI::MaskMergeMode::ZEROING, | ||
| 1016 | - AscendC::RoundMode::CAST_RINT, | ||
| 1017 | - }; | ||
| 1018 | - | ||
| 1019 | float epsilon = 1e-5; | 1002 | float epsilon = 1e-5; |
| 1020 | float nFactor; | 1003 | float nFactor; |
| 1021 | float nCorrectionFactor; | 1004 | float nCorrectionFactor; |
| @@ -30,6 +30,8 @@ using AscendC::MicroAPI::MemType; | |||
| 30 | using AscendC::MicroAPI::RegTensor; | 30 | using AscendC::MicroAPI::RegTensor; |
| 31 | using AscendC::MicroAPI::StoreDist; | 31 | using AscendC::MicroAPI::StoreDist; |
| 32 | using AscendC::MicroAPI::UpdateMask; | 32 | using AscendC::MicroAPI::UpdateMask; |
| 33 | +using AscendC::Reg::LoadAlign; | ||
| 34 | +using AscendC::Reg::StoreAlign; | ||
| 33 | 35 | ||
| 34 | constexpr uint32_t VL_FP32 = platform::GetVRegSize() / sizeof(float); | 36 | constexpr uint32_t VL_FP32 = platform::GetVRegSize() / sizeof(float); |
| 35 | constexpr uint32_t VL_F32 = VECTOR_REG_WIDTH / sizeof(float); | 37 | constexpr uint32_t VL_F32 = VECTOR_REG_WIDTH / sizeof(float); |
| @@ -52,54 +54,52 @@ constexpr AscendC::MicroAPI::CastTrait castTraitB322B16 = { | |||
| 52 | }; | 54 | }; |
| 53 | 55 | ||
| 54 | template <typename T_IN> | 56 | template <typename T_IN> |
| 55 | -__aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, | 57 | +__aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset) |
| 56 | - uint32_t offset) | ||
| 57 | { | 58 | { |
| 58 | if constexpr (IsSameType<T_IN, float>::value) { | 59 | if constexpr (IsSameType<T_IN, float>::value) { |
| 59 | - DataCopy<float, LoadDist::DIST_NORM>(dst, src + offset); | 60 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, src + offset); |
| 60 | } else { | 61 | } else { |
| 61 | RegTensor<T_IN> xIn; | 62 | RegTensor<T_IN> xIn; |
| 62 | - DataCopy<T_IN, LoadDist::DIST_UNPACK_B16>(xIn, src + offset); | 63 | + LoadAlign<T_IN, LoadDist::DIST_UNPACK_B16>(xIn, src + offset); |
| 63 | Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg); | 64 | Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg); |
| 64 | } | 65 | } |
| 65 | } | 66 | } |
| 66 | 67 | ||
| 67 | template <typename T_IN> | 68 | template <typename T_IN> |
| 68 | -__aicore__ inline void LoadScalarForDtypeTIn(__local_mem__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, | 69 | +__aicore__ inline void LoadScalarForDtypeTIn(__ubuf__ T_IN* src, RegTensor<float>& dst, MaskReg& preg, uint32_t offset) |
| 69 | - uint32_t offset) | ||
| 70 | { | 70 | { |
| 71 | if constexpr (IsSameType<T_IN, float>::value) { | 71 | if constexpr (IsSameType<T_IN, float>::value) { |
| 72 | - DataCopy<float, LoadDist::DIST_BRC_B32>(dst, src + offset); | 72 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(dst, src + offset); |
| 73 | } else { | 73 | } else { |
| 74 | RegTensor<T_IN> xIn; | 74 | RegTensor<T_IN> xIn; |
| 75 | - DataCopy<T_IN, LoadDist::DIST_BRC_B16>(xIn, src + offset); | 75 | + LoadAlign<T_IN, LoadDist::DIST_BRC_B16>(xIn, src + offset); |
| 76 | Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg); | 76 | Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg); |
| 77 | } | 77 | } |
| 78 | } | 78 | } |
| 79 | 79 | ||
| 80 | template <typename T_OUT> | 80 | template <typename T_OUT> |
| 81 | -__aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T_OUT* dst, RegTensor<float>& src, MaskReg& preg, | 81 | +__aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T_OUT* dst, RegTensor<float>& src, MaskReg& preg, |
| 82 | uint32_t offset) | 82 | uint32_t offset) |
| 83 | { | 83 | { |
| 84 | if constexpr (IsSameType<T_OUT, float>::value) { | 84 | if constexpr (IsSameType<T_OUT, float>::value) { |
| 85 | - DataCopy<T_OUT, StoreDist::DIST_NORM>(dst + offset, src, preg); | 85 | + StoreAlign<T_OUT, StoreDist::DIST_NORM>(dst + offset, src, preg); |
| 86 | } else { | 86 | } else { |
| 87 | RegTensor<T_OUT> xOut; | 87 | RegTensor<T_OUT> xOut; |
| 88 | Cast<T_OUT, float, castTraitB322B16>(xOut, src, preg); | 88 | Cast<T_OUT, float, castTraitB322B16>(xOut, src, preg); |
| 89 | - DataCopy<T_OUT, StoreDist::DIST_PACK_B32>(dst + offset, xOut, preg); | 89 | + StoreAlign<T_OUT, StoreDist::DIST_PACK_B32>(dst + offset, xOut, preg); |
| 90 | } | 90 | } |
| 91 | } | 91 | } |
| 92 | 92 | ||
| 93 | template <typename T_OUT> | 93 | template <typename T_OUT> |
| 94 | -__aicore__ inline void StoreOneElementForDtypeTOut(__local_mem__ T_OUT* dst, RegTensor<float>& src, MaskReg& preg, | 94 | +__aicore__ inline void StoreOneElementForDtypeTOut(__ubuf__ T_OUT* dst, RegTensor<float>& src, MaskReg& preg, |
| 95 | uint32_t offset) | 95 | uint32_t offset) |
| 96 | { | 96 | { |
| 97 | if constexpr (IsSameType<T_OUT, float>::value) { | 97 | if constexpr (IsSameType<T_OUT, float>::value) { |
| 98 | - DataCopy<T_OUT, StoreDist::DIST_FIRST_ELEMENT_B32>(dst + offset, src, preg); | 98 | + StoreAlign<T_OUT, StoreDist::DIST_FIRST_ELEMENT_B32>(dst + offset, src, preg); |
| 99 | } else { | 99 | } else { |
| 100 | RegTensor<T_OUT> xOut; | 100 | RegTensor<T_OUT> xOut; |
| 101 | Cast<T_OUT, float, castTraitB322B16>(xOut, src, preg); | 101 | Cast<T_OUT, float, castTraitB322B16>(xOut, src, preg); |
| 102 | - DataCopy<T_OUT, StoreDist::DIST_FIRST_ELEMENT_B16>(dst + offset, xOut, preg); | 102 | + StoreAlign<T_OUT, StoreDist::DIST_FIRST_ELEMENT_B16>(dst + offset, xOut, preg); |
| 103 | } | 103 | } |
| 104 | } | 104 | } |
| 105 | } // namespace InstanceNormOps | 105 | } // namespace InstanceNormOps |
| @@ -34,207 +34,3 @@ protected: | |||
| 34 | static void SetUpTestCase() { cout << " instance_norm_test SetUp\n" << endl; } | 34 | static void SetUpTestCase() { cout << " instance_norm_test SetUp\n" << endl; } |
| 35 | static void TearDownTestCase() { cout << " instance_norm_test TearDown\n" << endl; } | 35 | static void TearDownTestCase() { cout << " instance_norm_test TearDown\n" << endl; } |
| 36 | }; | 36 | }; |
| 37 | - | ||
| 38 | -TEST_F(instance_norm_test, test_case_200000) | ||
| 39 | -{ | ||
| 40 | - int64_t numN = 64; | ||
| 41 | - int64_t numC = 1; | ||
| 42 | - int64_t numR = 4; | ||
| 43 | - size_t xByteSize = numN * numC * numR * sizeof(half); | ||
| 44 | - size_t gammaByteSize = numC * sizeof(half); | ||
| 45 | - | ||
| 46 | - size_t tiling_data_size = sizeof(InstanceNormARFullReduceTilingData); | ||
| 47 | - uint32_t blockDim = 1; | ||
| 48 | - | ||
| 49 | - uint8_t* x = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 50 | - uint8_t* gamma = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 51 | - uint8_t* beta = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 52 | - uint8_t* y = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 53 | - uint8_t* mean_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 54 | - uint8_t* variance_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 55 | - uint8_t* workspace = (uint8_t*)AscendC::GmAlloc(16 * 2); | ||
| 56 | - uint8_t* tiling = (uint8_t*)AscendC::GmAlloc(tiling_data_size); | ||
| 57 | - | ||
| 58 | - InstanceNormARFullReduceTilingData* tilingDatafromBin = reinterpret_cast<InstanceNormARFullReduceTilingData*>( | ||
| 59 | - tiling); | ||
| 60 | - | ||
| 61 | - tilingDatafromBin->numN = 64; | ||
| 62 | - tilingDatafromBin->numC = 1; | ||
| 63 | - tilingDatafromBin->numR = 4; | ||
| 64 | - tilingDatafromBin->rAlign = 16; | ||
| 65 | - tilingDatafromBin->cInner = 1; | ||
| 66 | - tilingDatafromBin->cOuter = 1; | ||
| 67 | - tilingDatafromBin->cTail = 1; | ||
| 68 | - tilingDatafromBin->binaryAddQuotient = 8; | ||
| 69 | - tilingDatafromBin->perCoreCnt = 1; | ||
| 70 | - tilingDatafromBin->epsilon = 0.0001; | ||
| 71 | - tilingDatafromBin->avgFactor = 1.0 / numR; | ||
| 72 | - AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 73 | - ICPU_SET_TILING_KEY(200000); | ||
| 74 | - ICPU_RUN_KF(instance_norm, blockDim, x, gamma, beta, y, mean_out, variance_out, workspace, | ||
| 75 | - (uint8_t*)(tilingDatafromBin)); | ||
| 76 | - | ||
| 77 | - AscendC::GmFree(x); | ||
| 78 | - AscendC::GmFree(gamma); | ||
| 79 | - AscendC::GmFree(beta); | ||
| 80 | - AscendC::GmFree(y); | ||
| 81 | - AscendC::GmFree(mean_out); | ||
| 82 | - AscendC::GmFree(variance_out); | ||
| 83 | - AscendC::GmFree(workspace); | ||
| 84 | - AscendC::GmFree(tiling); | ||
| 85 | -} | ||
| 86 | - | ||
| 87 | -TEST_F(instance_norm_test, test_case_300000) | ||
| 88 | -{ | ||
| 89 | - int64_t numN = 2; | ||
| 90 | - int64_t numC = 1; | ||
| 91 | - int64_t numR = 22499; | ||
| 92 | - size_t xByteSize = numN * numC * numR * sizeof(half); | ||
| 93 | - size_t gammaByteSize = numC * sizeof(half); | ||
| 94 | - | ||
| 95 | - size_t tiling_data_size = sizeof(InstanceNormARWelfordTilingData); | ||
| 96 | - uint32_t blockDim = 1; | ||
| 97 | - | ||
| 98 | - uint8_t* x = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 99 | - uint8_t* gamma = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 100 | - uint8_t* beta = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 101 | - uint8_t* y = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 102 | - uint8_t* mean_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 103 | - uint8_t* variance_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 104 | - uint8_t* workspace = (uint8_t*)AscendC::GmAlloc(16 * 2); | ||
| 105 | - uint8_t* tiling = (uint8_t*)AscendC::GmAlloc(tiling_data_size); | ||
| 106 | - | ||
| 107 | - InstanceNormARWelfordTilingData* tilingDatafromBin = reinterpret_cast<InstanceNormARWelfordTilingData*>(tiling); | ||
| 108 | - | ||
| 109 | - tilingDatafromBin->a1 = 2; | ||
| 110 | - tilingDatafromBin->a0 = 1; | ||
| 111 | - tilingDatafromBin->r = 22499; | ||
| 112 | - tilingDatafromBin->blockNum = 2; | ||
| 113 | - tilingDatafromBin->totalTiles = 2; | ||
| 114 | - tilingDatafromBin->tilesPerCore = 1; | ||
| 115 | - tilingDatafromBin->a0Outer = 1; | ||
| 116 | - tilingDatafromBin->a0Inner = 128; | ||
| 117 | - tilingDatafromBin->a0Tail = 1; | ||
| 118 | - tilingDatafromBin->welfordTileLength = 10304; | ||
| 119 | - tilingDatafromBin->welfordTempSize = 83968; | ||
| 120 | - tilingDatafromBin->welfordUpdateTimes = 2; | ||
| 121 | - tilingDatafromBin->welfordUpdateTail = 1891; | ||
| 122 | - tilingDatafromBin->apiTempBufferSize = 42240; | ||
| 123 | - tilingDatafromBin->epsilon = 0.001; | ||
| 124 | - ICPU_SET_TILING_KEY(300000); | ||
| 125 | - ICPU_RUN_KF(instance_norm, blockDim, x, gamma, beta, y, mean_out, variance_out, workspace, | ||
| 126 | - (uint8_t*)(tilingDatafromBin)); | ||
| 127 | - | ||
| 128 | - AscendC::GmFree(x); | ||
| 129 | - AscendC::GmFree(gamma); | ||
| 130 | - AscendC::GmFree(beta); | ||
| 131 | - AscendC::GmFree(y); | ||
| 132 | - AscendC::GmFree(mean_out); | ||
| 133 | - AscendC::GmFree(variance_out); | ||
| 134 | - AscendC::GmFree(workspace); | ||
| 135 | - AscendC::GmFree(tiling); | ||
| 136 | -} | ||
| 137 | - | ||
| 138 | -TEST_F(instance_norm_test, test_case_400000) | ||
| 139 | -{ | ||
| 140 | - int64_t numN = 33; | ||
| 141 | - int64_t numC = 207; | ||
| 142 | - int64_t numR = 51; | ||
| 143 | - size_t xByteSize = numN * numC * numR * sizeof(half); | ||
| 144 | - size_t gammaByteSize = numC * sizeof(half); | ||
| 145 | - | ||
| 146 | - size_t tiling_data_size = sizeof(InstanceNormARAFullReduceTilingData); | ||
| 147 | - uint32_t blockDim = 1; | ||
| 148 | - | ||
| 149 | - uint8_t* x = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 150 | - uint8_t* gamma = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 151 | - uint8_t* beta = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 152 | - uint8_t* y = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 153 | - uint8_t* mean_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 154 | - uint8_t* variance_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 155 | - uint8_t* workspace = (uint8_t*)AscendC::GmAlloc(16 * 2); | ||
| 156 | - uint8_t* tiling = (uint8_t*)AscendC::GmAlloc(tiling_data_size); | ||
| 157 | - | ||
| 158 | - InstanceNormARAFullReduceTilingData* tilingDatafromBin = reinterpret_cast<InstanceNormARAFullReduceTilingData*>( | ||
| 159 | - tiling); | ||
| 160 | - | ||
| 161 | - tilingDatafromBin->usedCoreNum = 33; | ||
| 162 | - tilingDatafromBin->totalTiles = 66; | ||
| 163 | - tilingDatafromBin->tilesPerCore = 2; | ||
| 164 | - tilingDatafromBin->totalA1Len = 33; | ||
| 165 | - tilingDatafromBin->totalRLen = 51; | ||
| 166 | - tilingDatafromBin->totalA0Len = 207; | ||
| 167 | - tilingDatafromBin->a0Outer = 2; | ||
| 168 | - tilingDatafromBin->tileA0Len = 192; | ||
| 169 | - tilingDatafromBin->tileA0Tail = 15; | ||
| 170 | - tilingDatafromBin->powerOfTwoForR = 32; | ||
| 171 | - tilingDatafromBin->binaryAddQuotient = 32; | ||
| 172 | - tilingDatafromBin->binaryAddK = 1; | ||
| 173 | - tilingDatafromBin->binaryAddLast = 0; | ||
| 174 | - tilingDatafromBin->epsilon = 0.001; | ||
| 175 | - AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 176 | - ICPU_SET_TILING_KEY(400000); | ||
| 177 | - ICPU_RUN_KF(instance_norm, blockDim, x, gamma, beta, y, mean_out, variance_out, workspace, | ||
| 178 | - (uint8_t*)(tilingDatafromBin)); | ||
| 179 | - | ||
| 180 | - AscendC::GmFree(x); | ||
| 181 | - AscendC::GmFree(gamma); | ||
| 182 | - AscendC::GmFree(beta); | ||
| 183 | - AscendC::GmFree(y); | ||
| 184 | - AscendC::GmFree(mean_out); | ||
| 185 | - AscendC::GmFree(variance_out); | ||
| 186 | - AscendC::GmFree(workspace); | ||
| 187 | - AscendC::GmFree(tiling); | ||
| 188 | -} | ||
| 189 | - | ||
| 190 | -TEST_F(instance_norm_test, test_case_500000) | ||
| 191 | -{ | ||
| 192 | - int64_t numN = 2; | ||
| 193 | - int64_t numC = 1000; | ||
| 194 | - int64_t numR = 2; | ||
| 195 | - size_t xByteSize = numN * numC * numR * sizeof(half); | ||
| 196 | - size_t gammaByteSize = numC * sizeof(half); | ||
| 197 | - | ||
| 198 | - size_t tiling_data_size = sizeof(InstanceNormARAWelfordTilingData); | ||
| 199 | - uint32_t blockDim = 1; | ||
| 200 | - | ||
| 201 | - uint8_t* x = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 202 | - uint8_t* gamma = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 203 | - uint8_t* beta = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 204 | - uint8_t* y = (uint8_t*)AscendC::GmAlloc(xByteSize); | ||
| 205 | - uint8_t* mean_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 206 | - uint8_t* variance_out = (uint8_t*)AscendC::GmAlloc(gammaByteSize); | ||
| 207 | - uint8_t* workspace = (uint8_t*)AscendC::GmAlloc(16 * 2); | ||
| 208 | - uint8_t* tiling = (uint8_t*)AscendC::GmAlloc(tiling_data_size); | ||
| 209 | - | ||
| 210 | - InstanceNormARAWelfordTilingData* tilingDatafromBin = reinterpret_cast<InstanceNormARAWelfordTilingData*>(tiling); | ||
| 211 | - | ||
| 212 | - tilingDatafromBin->a1 = 2; | ||
| 213 | - tilingDatafromBin->r = 1000; | ||
| 214 | - tilingDatafromBin->a0 = 2; | ||
| 215 | - tilingDatafromBin->usedCoreNum = 2; | ||
| 216 | - tilingDatafromBin->totalTiles = 2; | ||
| 217 | - tilingDatafromBin->tilesPerCore = 1; | ||
| 218 | - tilingDatafromBin->a0Outer = 1; | ||
| 219 | - tilingDatafromBin->tileA0Len = 128; | ||
| 220 | - tilingDatafromBin->tileA0Tail = 2; | ||
| 221 | - tilingDatafromBin->welfordrFactor = 120; | ||
| 222 | - tilingDatafromBin->binaryAddQuotient = 64; | ||
| 223 | - tilingDatafromBin->binaryAddK = 2; | ||
| 224 | - tilingDatafromBin->binaryAddLast = 0; | ||
| 225 | - tilingDatafromBin->powerOfTwoForR = 1024; | ||
| 226 | - tilingDatafromBin->epsilon = 0.001; | ||
| 227 | - AscendC::SetKernelMode(KernelMode::AIV_MODE); | ||
| 228 | - ICPU_SET_TILING_KEY(500000); | ||
| 229 | - ICPU_RUN_KF(instance_norm, blockDim, x, gamma, beta, y, mean_out, variance_out, workspace, | ||
| 230 | - (uint8_t*)(tilingDatafromBin)); | ||
| 231 | - | ||
| 232 | - AscendC::GmFree(x); | ||
| 233 | - AscendC::GmFree(gamma); | ||
| 234 | - AscendC::GmFree(beta); | ||
| 235 | - AscendC::GmFree(y); | ||
| 236 | - AscendC::GmFree(mean_out); | ||
| 237 | - AscendC::GmFree(variance_out); | ||
| 238 | - AscendC::GmFree(workspace); | ||
| 239 | - AscendC::GmFree(tiling); | ||
| 240 | -} | ||
| @@ -20,6 +20,9 @@ | |||
| 20 | 20 | ||
| 21 | namespace LayerNormGrad { | 21 | namespace LayerNormGrad { |
| 22 | using namespace AscendC; | 22 | using namespace AscendC; |
| 23 | +using AscendC::Reg::LoadAlign; | ||
| 24 | +using AscendC::Reg::Move; | ||
| 25 | +using AscendC::Reg::StoreAlign; | ||
| 23 | 26 | ||
| 24 | constexpr static int64_t BLOCK_SIZE = 32; | 27 | constexpr static int64_t BLOCK_SIZE = 32; |
| 25 | 28 | ||
| @@ -153,19 +156,19 @@ __aicore__ inline void VectorAdd(const LocalTensor<float>& dstTensor, const Loca | |||
| 153 | uint16_t loopTimes = Arith::CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(VREG_SIZE)); | 156 | uint16_t loopTimes = Arith::CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(VREG_SIZE)); |
| 154 | __VEC_SCOPE__ | 157 | __VEC_SCOPE__ |
| 155 | { | 158 | { |
| 156 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 159 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 157 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 160 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 158 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 161 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 159 | uint32_t sreg = static_cast<uint32_t>(count); | 162 | uint32_t sreg = static_cast<uint32_t>(count); |
| 160 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 163 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 161 | AscendC::MicroAPI::MaskReg pMask; | 164 | AscendC::MicroAPI::MaskReg pMask; |
| 162 | for (uint16_t i = 0; i < loopTimes; ++i) { | 165 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 163 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 166 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 164 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32); | 167 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32); |
| 165 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32); | 168 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32); |
| 166 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 169 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 167 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 170 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 168 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask); | 171 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask); |
| 169 | } | 172 | } |
| 170 | } | 173 | } |
| 171 | } | 174 | } |
| @@ -207,20 +210,20 @@ __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const Lo | |||
| 207 | uint32_t innerLoopStride = stride; | 210 | uint32_t innerLoopStride = stride; |
| 208 | __VEC_SCOPE__ | 211 | __VEC_SCOPE__ |
| 209 | { | 212 | { |
| 210 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 213 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 211 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 214 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 212 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 215 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 213 | uint32_t sreg = static_cast<uint32_t>(count); | 216 | uint32_t sreg = static_cast<uint32_t>(count); |
| 214 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 217 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 215 | AscendC::MicroAPI::MaskReg pMask; | 218 | AscendC::MicroAPI::MaskReg pMask; |
| 216 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 219 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 217 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 220 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 218 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 221 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 219 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 222 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 220 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride); | 223 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride); |
| 221 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 224 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 222 | } | 225 | } |
| 223 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 226 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 224 | } | 227 | } |
| 225 | } | 228 | } |
| 226 | } | 229 | } |
| @@ -229,4 +232,4 @@ __aicore__ inline void UpdateCache(const LocalTensor<float>& dstTensor, const Lo | |||
| 229 | 232 | ||
| 230 | } // namespace LayerNormGrad | 233 | } // namespace LayerNormGrad |
| 231 | 234 | ||
| 232 | -#endif // LAYER_NORM_GRAD_API_ | 235 | +#endif // LAYER_NORM_GRAD_API_ |
| @@ -34,6 +34,10 @@ using AscendC::MicroAPI::RegTensor; | |||
| 34 | using namespace NormCommon; | 34 | using namespace NormCommon; |
| 35 | using namespace NormCommon::NormCommonRegbase; | 35 | using namespace NormCommon::NormCommonRegbase; |
| 36 | using namespace LayerNormGrad::Arith; | 36 | using namespace LayerNormGrad::Arith; |
| 37 | +using AscendC::Reg::LoadAlign; | ||
| 38 | +using AscendC::Reg::Move; | ||
| 39 | +using AscendC::Reg::Reduce; | ||
| 40 | +using AscendC::Reg::StoreAlign; | ||
| 37 | 41 | ||
| 38 | constexpr static AscendC::MicroAPI::CastTrait castTraitB162B32 = { | 42 | constexpr static AscendC::MicroAPI::CastTrait castTraitB162B32 = { |
| 39 | AscendC::MicroAPI::RegLayout::ZERO, | 43 | AscendC::MicroAPI::RegLayout::ZERO, |
| @@ -122,7 +126,7 @@ public: | |||
| 122 | const LocalTensor<float>& meanTensor, const LocalTensor<float>& varTensor, | 126 | const LocalTensor<float>& meanTensor, const LocalTensor<float>& varTensor, |
| 123 | const int64_t rowSize, const int64_t colSize, const float epsilon); | 127 | const int64_t rowSize, const int64_t colSize, const float epsilon); |
| 124 | template <typename T> | 128 | template <typename T> |
| 125 | - __aicore__ inline static void StoreTensorForDtypeT(__local_mem__ T* dst, AscendC::MicroAPI::RegTensor<float>& src, | 129 | + __aicore__ inline static void StoreTensorForDtypeT(__ubuf__ T* dst, AscendC::MicroAPI::RegTensor<float>& src, |
| 126 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); | 130 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); |
| 127 | 131 | ||
| 128 | protected: | 132 | protected: |
| @@ -176,8 +180,8 @@ __aicore__ inline void LayerNormGradBase::CastToFp32From(const LocalTensor<float | |||
| 176 | if (innerLoopTimes == 1) { | 180 | if (innerLoopTimes == 1) { |
| 177 | __VEC_SCOPE__ | 181 | __VEC_SCOPE__ |
| 178 | { | 182 | { |
| 179 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 183 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 180 | - __local_mem__ T* src = (__local_mem__ T*)srcTensor.GetPhyAddr(); | 184 | + __ubuf__ T* src = (__ubuf__ T*)srcTensor.GetPhyAddr(); |
| 181 | uint32_t count; | 185 | uint32_t count; |
| 182 | AscendC::MicroAPI::RegTensor<float> fp32Reg; | 186 | AscendC::MicroAPI::RegTensor<float> fp32Reg; |
| 183 | AscendC::MicroAPI::RegTensor<T> b16Reg; | 187 | AscendC::MicroAPI::RegTensor<T> b16Reg; |
| @@ -185,17 +189,17 @@ __aicore__ inline void LayerNormGradBase::CastToFp32From(const LocalTensor<float | |||
| 185 | count = static_cast<uint32_t>(colSize); | 189 | count = static_cast<uint32_t>(colSize); |
| 186 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 190 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 187 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 191 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 188 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( | 192 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( |
| 189 | - b16Reg, (__local_mem__ T*)src + i * outerLoopSrcStride + 0 * innerLoopStride); | 193 | + b16Reg, (__ubuf__ T*)src + i * outerLoopSrcStride + 0 * innerLoopStride); |
| 190 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); | 194 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); |
| 191 | - DataCopy((__local_mem__ float*)dst + i * outerLoopDstStride + 0 * innerLoopStride, fp32Reg, pMask); | 195 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopDstStride + 0 * innerLoopStride, fp32Reg, pMask); |
| 192 | } | 196 | } |
| 193 | } | 197 | } |
| 194 | } else { | 198 | } else { |
| 195 | __VEC_SCOPE__ | 199 | __VEC_SCOPE__ |
| 196 | { | 200 | { |
| 197 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 201 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 198 | - __local_mem__ T* src = (__local_mem__ T*)srcTensor.GetPhyAddr(); | 202 | + __ubuf__ T* src = (__ubuf__ T*)srcTensor.GetPhyAddr(); |
| 199 | uint32_t count; | 203 | uint32_t count; |
| 200 | AscendC::MicroAPI::RegTensor<float> fp32Reg; | 204 | AscendC::MicroAPI::RegTensor<float> fp32Reg; |
| 201 | AscendC::MicroAPI::RegTensor<T> b16Reg; | 205 | AscendC::MicroAPI::RegTensor<T> b16Reg; |
| @@ -204,10 +208,10 @@ __aicore__ inline void LayerNormGradBase::CastToFp32From(const LocalTensor<float | |||
| 204 | count = static_cast<uint32_t>(colSize); | 208 | count = static_cast<uint32_t>(colSize); |
| 205 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 209 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 206 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 210 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 207 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( | 211 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( |
| 208 | - b16Reg, (__local_mem__ T*)src + i * outerLoopSrcStride + j * innerLoopStride); | 212 | + b16Reg, (__ubuf__ T*)src + i * outerLoopSrcStride + j * innerLoopStride); |
| 209 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); | 213 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); |
| 210 | - DataCopy((__local_mem__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask); | 214 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask); |
| 211 | } | 215 | } |
| 212 | } | 216 | } |
| 213 | } | 217 | } |
| @@ -287,8 +291,8 @@ __aicore__ inline void LayerNormGradBase::CopyUB2UBWithCast(const LocalTensor<T> | |||
| 287 | if constexpr (IsSameType<T, float>::value) { | 291 | if constexpr (IsSameType<T, float>::value) { |
| 288 | CopyUB2UB(dstTensor, srcTensor, count); | 292 | CopyUB2UB(dstTensor, srcTensor, count); |
| 289 | } else { | 293 | } else { |
| 290 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 294 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 291 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 295 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 292 | 296 | ||
| 293 | uint32_t cnt = count; | 297 | uint32_t cnt = count; |
| 294 | uint16_t loopNum = CeilDiv(cnt, VL_FP32); | 298 | uint16_t loopNum = CeilDiv(cnt, VL_FP32); |
| @@ -301,10 +305,10 @@ __aicore__ inline void LayerNormGradBase::CopyUB2UBWithCast(const LocalTensor<T> | |||
| 301 | for (uint16_t k = 0; k < loopNum; k++) { | 305 | for (uint16_t k = 0; k < loopNum; k++) { |
| 302 | pregMask = UpdateMask<float>(sreg); | 306 | pregMask = UpdateMask<float>(sreg); |
| 303 | uint32_t offset = k * VL_FP32; | 307 | uint32_t offset = k * VL_FP32; |
| 304 | - DataCopy<float, LoadDist::DIST_NORM>(srcReg, (__local_mem__ float*)src + offset); | 308 | + LoadAlign<float, LoadDist::DIST_NORM>(srcReg, (__ubuf__ float*)src + offset); |
| 305 | 309 | ||
| 306 | Cast<T, float, castTraitB322B16>(xFp16, srcReg, pregMask); | 310 | Cast<T, float, castTraitB322B16>(xFp16, srcReg, pregMask); |
| 307 | - DataCopy<T, StoreDist::DIST_PACK_B32>(((__local_mem__ T*)dst) + offset, xFp16, pregMask); | 311 | + StoreAlign<T, StoreDist::DIST_PACK_B32>(((__ubuf__ T*)dst) + offset, xFp16, pregMask); |
| 308 | } | 312 | } |
| 309 | } | 313 | } |
| 310 | } | 314 | } |
| @@ -321,19 +325,19 @@ __aicore__ inline void LayerNormGradBase::VectorAdd(const LocalTensor<float>& ds | |||
| 321 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); | 325 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); |
| 322 | __VEC_SCOPE__ | 326 | __VEC_SCOPE__ |
| 323 | { | 327 | { |
| 324 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 328 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 325 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 329 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 326 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 330 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 327 | uint32_t sreg = static_cast<uint32_t>(count); | 331 | uint32_t sreg = static_cast<uint32_t>(count); |
| 328 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 332 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 329 | AscendC::MicroAPI::MaskReg pMask; | 333 | AscendC::MicroAPI::MaskReg pMask; |
| 330 | for (uint16_t i = 0; i < loopTimes; ++i) { | 334 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 331 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 335 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 332 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32); | 336 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32); |
| 333 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32); | 337 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32); |
| 334 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 338 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 335 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 339 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 336 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask); | 340 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask); |
| 337 | } | 341 | } |
| 338 | } | 342 | } |
| 339 | } | 343 | } |
| @@ -351,38 +355,38 @@ __aicore__ inline void LayerNormGradBase::VectorAdd(const LocalTensor<float>& ds | |||
| 351 | if (innerLoopTimes == 1) { | 355 | if (innerLoopTimes == 1) { |
| 352 | __VEC_SCOPE__ | 356 | __VEC_SCOPE__ |
| 353 | { | 357 | { |
| 354 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 358 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 355 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 359 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 356 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 360 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 357 | uint32_t count = nSize; | 361 | uint32_t count = nSize; |
| 358 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 362 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 359 | AscendC::MicroAPI::MaskReg pMask; | 363 | AscendC::MicroAPI::MaskReg pMask; |
| 360 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 364 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 361 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 365 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 362 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); | 366 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); |
| 363 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride + 0 * innerLoopStride); | 367 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride + 0 * innerLoopStride); |
| 364 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 368 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 365 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 369 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 366 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, aReg, pMask); | 370 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, aReg, pMask); |
| 367 | } | 371 | } |
| 368 | } | 372 | } |
| 369 | } else { | 373 | } else { |
| 370 | __VEC_SCOPE__ | 374 | __VEC_SCOPE__ |
| 371 | { | 375 | { |
| 372 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 376 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 373 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 377 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 374 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 378 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 375 | uint32_t count = nSize; | 379 | uint32_t count = nSize; |
| 376 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 380 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 377 | AscendC::MicroAPI::MaskReg pMask; | 381 | AscendC::MicroAPI::MaskReg pMask; |
| 378 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 382 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 379 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 383 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 380 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 384 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 381 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride); | 385 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride); |
| 382 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride + j * innerLoopStride); | 386 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride + j * innerLoopStride); |
| 383 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 387 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 384 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 388 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 385 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask); | 389 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask); |
| 386 | } | 390 | } |
| 387 | } | 391 | } |
| 388 | } | 392 | } |
| @@ -400,19 +404,19 @@ __aicore__ inline void LayerNormGradBase::VectorMul(const LocalTensor<float>& ds | |||
| 400 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); | 404 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); |
| 401 | __VEC_SCOPE__ | 405 | __VEC_SCOPE__ |
| 402 | { | 406 | { |
| 403 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 407 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 404 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 408 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 405 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 409 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 406 | uint32_t sreg = static_cast<uint32_t>(count); | 410 | uint32_t sreg = static_cast<uint32_t>(count); |
| 407 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 411 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 408 | AscendC::MicroAPI::MaskReg pMask; | 412 | AscendC::MicroAPI::MaskReg pMask; |
| 409 | 413 | ||
| 410 | for (uint16_t i = 0; i < loopTimes; ++i) { | 414 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 411 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 415 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 412 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32); | 416 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32); |
| 413 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32); | 417 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32); |
| 414 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 418 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 415 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, cReg, pMask); | 419 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, cReg, pMask); |
| 416 | } | 420 | } |
| 417 | } | 421 | } |
| 418 | } | 422 | } |
| @@ -436,36 +440,36 @@ __aicore__ inline void LayerNormGradBase::NlastBroadcastMul(const LocalTensor<fl | |||
| 436 | if (innerLoopTimes == 1) { | 440 | if (innerLoopTimes == 1) { |
| 437 | __VEC_SCOPE__ | 441 | __VEC_SCOPE__ |
| 438 | { | 442 | { |
| 439 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 443 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 440 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 444 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 441 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 445 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 442 | uint32_t count = static_cast<uint32_t>(aSize); | 446 | uint32_t count = static_cast<uint32_t>(aSize); |
| 443 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 447 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 444 | AscendC::MicroAPI::MaskReg pMask; | 448 | AscendC::MicroAPI::MaskReg pMask; |
| 445 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 449 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 446 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 450 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 447 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride); | 451 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride); |
| 448 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); | 452 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); |
| 449 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 453 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 450 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); | 454 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); |
| 451 | } | 455 | } |
| 452 | } | 456 | } |
| 453 | } else { | 457 | } else { |
| 454 | __VEC_SCOPE__ | 458 | __VEC_SCOPE__ |
| 455 | { | 459 | { |
| 456 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 460 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 457 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 461 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 458 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 462 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 459 | uint32_t count = static_cast<uint32_t>(aSize); | 463 | uint32_t count = static_cast<uint32_t>(aSize); |
| 460 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 464 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 461 | AscendC::MicroAPI::MaskReg pMask; | 465 | AscendC::MicroAPI::MaskReg pMask; |
| 462 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 466 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 463 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 467 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 464 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride); | 468 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride); |
| 465 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 469 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 466 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride); | 470 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride); |
| 467 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 471 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 468 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); | 472 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); |
| 469 | } | 473 | } |
| 470 | } | 474 | } |
| 471 | } | 475 | } |
| @@ -491,40 +495,40 @@ __aicore__ inline void LayerNormGradBase::LastReduceSumSmallR(const LocalTensor< | |||
| 491 | if (rSize <= VL_FP32) { | 495 | if (rSize <= VL_FP32) { |
| 492 | __VEC_SCOPE__ | 496 | __VEC_SCOPE__ |
| 493 | { | 497 | { |
| 494 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 498 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 495 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 499 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 496 | uint32_t count = static_cast<uint32_t>(rSize); | 500 | uint32_t count = static_cast<uint32_t>(rSize); |
| 497 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 501 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 498 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 502 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 499 | - AscendC::MicroAPI::UnalignReg UReg; | 503 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 500 | for (uint16_t i = 0; i < loopTimes; ++i) { | 504 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 501 | - DataCopy(aReg, (__local_mem__ float*)src + i * stride); | 505 | + LoadAlign(aReg, (__ubuf__ float*)src + i * stride); |
| 502 | - ReduceSum(bReg, aReg, pMask); | 506 | + Reduce<ReduceType::SUM>(bReg, aReg, pMask); |
| 503 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 507 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 504 | } | 508 | } |
| 505 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 509 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 506 | } | 510 | } |
| 507 | } else { | 511 | } else { |
| 508 | __VEC_SCOPE__ | 512 | __VEC_SCOPE__ |
| 509 | { | 513 | { |
| 510 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 514 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 511 | - __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 515 | + __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 512 | - __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32; | 516 | + __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32; |
| 513 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); | 517 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); |
| 514 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 518 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 515 | - AscendC::MicroAPI::UnalignReg UReg; | 519 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 516 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 520 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 517 | AscendC::MicroAPI::MaskReg | 521 | AscendC::MicroAPI::MaskReg |
| 518 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 522 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 519 | for (uint16_t i = 0; i < loopTimes; ++i) { | 523 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 520 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * stride); | 524 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * stride); |
| 521 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * stride); | 525 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * stride); |
| 522 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 526 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 523 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 527 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 524 | - ReduceSum(bReg, aReg, pFull); | 528 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 525 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 529 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 526 | } | 530 | } |
| 527 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 531 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 528 | } | 532 | } |
| 529 | } | 533 | } |
| 530 | } | 534 | } |
| @@ -567,43 +571,43 @@ __aicore__ inline void LayerNormGradBase::LastReduceSum(const LocalTensor<float> | |||
| 567 | 571 | ||
| 568 | __VEC_SCOPE__ | 572 | __VEC_SCOPE__ |
| 569 | { | 573 | { |
| 570 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr(); | 574 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr(); |
| 571 | - __local_mem__ float* foldSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 575 | + __ubuf__ float* foldSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 572 | - __local_mem__ float* foldSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; | 576 | + __ubuf__ float* foldSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; |
| 573 | - __local_mem__ float* tailSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; | 577 | + __ubuf__ float* tailSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; |
| 574 | - __local_mem__ float* tailSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; | 578 | + __ubuf__ float* tailSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; |
| 575 | - __local_mem__ float* unFoldSrc = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; | 579 | + __ubuf__ float* unFoldSrc = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; |
| 576 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 580 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 577 | - AscendC::MicroAPI::UnalignReg UReg; | 581 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 578 | 582 | ||
| 579 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 583 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 580 | - dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; | 584 | + dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; |
| 581 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { | 585 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { |
| 582 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; | 586 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; |
| 583 | - DataCopy(aReg, (__local_mem__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); | 587 | + LoadAlign(aReg, (__ubuf__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); |
| 584 | - DataCopy(bReg, (__local_mem__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); | 588 | + LoadAlign(bReg, (__ubuf__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); |
| 585 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); | 589 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); |
| 586 | - ReduceSum(dReg, cReg, pFull); | 590 | + Reduce<ReduceType::SUM>(dReg, cReg, pFull); |
| 587 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, dReg, UReg, 1); | 591 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, dReg, UReg, 1); |
| 588 | } | 592 | } |
| 589 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { | 593 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { |
| 590 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); | 594 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); |
| 591 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 595 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 592 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 596 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 593 | - DataCopy(aReg, (__local_mem__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); | 597 | + LoadAlign(aReg, (__ubuf__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); |
| 594 | - DataCopy(bReg, (__local_mem__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); | 598 | + LoadAlign(bReg, (__ubuf__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); |
| 595 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 599 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 596 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 600 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 597 | - ReduceSum(bReg, aReg, pFull); | 601 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 598 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 602 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 599 | } | 603 | } |
| 600 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { | 604 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { |
| 601 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 605 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 602 | - DataCopy(aReg, (__local_mem__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); | 606 | + LoadAlign(aReg, (__ubuf__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); |
| 603 | - ReduceSum(bReg, aReg, pFull); | 607 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 604 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 608 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 605 | } | 609 | } |
| 606 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 610 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 607 | } | 611 | } |
| 608 | } | 612 | } |
| 609 | LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); | 613 | LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); |
| @@ -611,13 +615,13 @@ __aicore__ inline void LayerNormGradBase::LastReduceSum(const LocalTensor<float> | |||
| 611 | 615 | ||
| 612 | template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1> | 616 | template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1> |
| 613 | struct NlastDichotomyAdd { | 617 | struct NlastDichotomyAdd { |
| 614 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 618 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 615 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 619 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 616 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 620 | + uint32_t stride) |
| 617 | { | 621 | { |
| 618 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 622 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 619 | - __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO; | 623 | + __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO; |
| 620 | - __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO; | 624 | + __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO; |
| 621 | if constexpr (TailCount <= 0) { | 625 | if constexpr (TailCount <= 0) { |
| 622 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, | 626 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, |
| 623 | stride * CONST_TWO); | 627 | stride * CONST_TWO); |
| @@ -625,13 +629,13 @@ struct NlastDichotomyAdd { | |||
| 625 | } | 629 | } |
| 626 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 630 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 627 | } | 631 | } |
| 628 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 632 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 629 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 633 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 630 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 634 | + uint32_t stride, uint32_t offset) |
| 631 | { | 635 | { |
| 632 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 636 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 633 | - __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO; | 637 | + __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO; |
| 634 | - __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO; | 638 | + __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO; |
| 635 | if constexpr (TailCount <= 0) { | 639 | if constexpr (TailCount <= 0) { |
| 636 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, | 640 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, |
| 637 | stride * CONST_TWO, offset); | 641 | stride * CONST_TWO, offset); |
| @@ -649,49 +653,49 @@ struct NlastDichotomyAdd { | |||
| 649 | 653 | ||
| 650 | template <int32_t TailCount, int32_t Index, int32_t Depth> | 654 | template <int32_t TailCount, int32_t Index, int32_t Depth> |
| 651 | struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { | 655 | struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { |
| 652 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 656 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 653 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 657 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 654 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 658 | + uint32_t stride) |
| 655 | { | 659 | { |
| 656 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 660 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 657 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 661 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 658 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 662 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 659 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 663 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 660 | } | 664 | } |
| 661 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 665 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 662 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 666 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 663 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 667 | + uint32_t stride, uint32_t offset) |
| 664 | { | 668 | { |
| 665 | if constexpr (TailCount <= 0) { | 669 | if constexpr (TailCount <= 0) { |
| 666 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 670 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 667 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 671 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 668 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 672 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 669 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 673 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 670 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 674 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 671 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 675 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 672 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 676 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 673 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 677 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 674 | } else { | 678 | } else { |
| 675 | if constexpr (Index + Depth < TailCount) { | 679 | if constexpr (Index + Depth < TailCount) { |
| 676 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 680 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 677 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 681 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 678 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 682 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 679 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 683 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 680 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 684 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 681 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 685 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 682 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 686 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 683 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 687 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 684 | } else if constexpr (Index < TailCount) { | 688 | } else if constexpr (Index < TailCount) { |
| 685 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 689 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 686 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 690 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 687 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 691 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 688 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 692 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 689 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 693 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 690 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 694 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 691 | } else { | 695 | } else { |
| 692 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 696 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 693 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 697 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 694 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 698 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 695 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 699 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 696 | } | 700 | } |
| 697 | } | 701 | } |
| @@ -700,25 +704,25 @@ struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { | |||
| 700 | 704 | ||
| 701 | template <> | 705 | template <> |
| 702 | struct NlastDichotomyAdd<CONST_TWO> { | 706 | struct NlastDichotomyAdd<CONST_TWO> { |
| 703 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 707 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 704 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 708 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 705 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 709 | + uint32_t stride) |
| 706 | { | 710 | { |
| 707 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 711 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 708 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 712 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 709 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 713 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 710 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 714 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 711 | } | 715 | } |
| 712 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 716 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 713 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 717 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 714 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 718 | + uint32_t stride, uint32_t offset) |
| 715 | { | 719 | { |
| 716 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 720 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 717 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 721 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 718 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 722 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 719 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 723 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 720 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 724 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 721 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 725 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 722 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 726 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 723 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 727 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 724 | } | 728 | } |
| @@ -726,11 +730,11 @@ struct NlastDichotomyAdd<CONST_TWO> { | |||
| 726 | 730 | ||
| 727 | template <> | 731 | template <> |
| 728 | struct NlastDichotomyAdd<1> { | 732 | struct NlastDichotomyAdd<1> { |
| 729 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 733 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 730 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 734 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 731 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 735 | + uint32_t stride) |
| 732 | { | 736 | { |
| 733 | - DataCopy(acc, (__local_mem__ float*)srcA); | 737 | + LoadAlign(acc, (__ubuf__ float*)srcA); |
| 734 | } | 738 | } |
| 735 | }; | 739 | }; |
| 736 | 740 | ||
| @@ -757,56 +761,56 @@ __aicore__ inline void LayerNormGradBase::UpdateCache(const LocalTensor<float>& | |||
| 757 | if (innerLoopTimes == 1) { | 761 | if (innerLoopTimes == 1) { |
| 758 | __VEC_SCOPE__ | 762 | __VEC_SCOPE__ |
| 759 | { | 763 | { |
| 760 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 764 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 761 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 765 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 762 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 766 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 763 | uint32_t sreg = static_cast<uint32_t>(count); | 767 | uint32_t sreg = static_cast<uint32_t>(count); |
| 764 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 768 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 765 | AscendC::MicroAPI::MaskReg pMask; | 769 | AscendC::MicroAPI::MaskReg pMask; |
| 766 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 770 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 767 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 771 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 768 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 772 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 769 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); | 773 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); |
| 770 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 774 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 771 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 775 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 772 | } | 776 | } |
| 773 | } | 777 | } |
| 774 | } else if (innerLoopTimes == CONST_TWO) { | 778 | } else if (innerLoopTimes == CONST_TWO) { |
| 775 | __VEC_SCOPE__ | 779 | __VEC_SCOPE__ |
| 776 | { | 780 | { |
| 777 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 781 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 778 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 782 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 779 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 783 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 780 | uint32_t sreg = static_cast<uint32_t>(count); | 784 | uint32_t sreg = static_cast<uint32_t>(count); |
| 781 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 785 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 782 | AscendC::MicroAPI::MaskReg pMask; | 786 | AscendC::MicroAPI::MaskReg pMask; |
| 783 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 787 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 784 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 788 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 785 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 789 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 786 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); | 790 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); |
| 787 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 791 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 788 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + 1 * innerLoopStride); | 792 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + 1 * innerLoopStride); |
| 789 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 793 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 790 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 794 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 791 | } | 795 | } |
| 792 | } | 796 | } |
| 793 | } else { | 797 | } else { |
| 794 | __VEC_SCOPE__ | 798 | __VEC_SCOPE__ |
| 795 | { | 799 | { |
| 796 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 800 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 797 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 801 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 798 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 802 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 799 | uint32_t sreg = static_cast<uint32_t>(count); | 803 | uint32_t sreg = static_cast<uint32_t>(count); |
| 800 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 804 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 801 | AscendC::MicroAPI::MaskReg pMask; | 805 | AscendC::MicroAPI::MaskReg pMask; |
| 802 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 806 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 803 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 807 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 804 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 808 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 805 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 809 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 806 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride); | 810 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride); |
| 807 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 811 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 808 | } | 812 | } |
| 809 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 813 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 810 | } | 814 | } |
| 811 | } | 815 | } |
| 812 | } | 816 | } |
| @@ -827,10 +831,10 @@ __aicore__ inline void LayerNormGradBase::Normalize(const LocalTensor<float>& ds | |||
| 827 | if (innerLoopTimes == 1) { | 831 | if (innerLoopTimes == 1) { |
| 828 | __VEC_SCOPE__ | 832 | __VEC_SCOPE__ |
| 829 | { | 833 | { |
| 830 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 834 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 831 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 835 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 832 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 836 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 833 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 837 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 834 | uint32_t count; | 838 | uint32_t count; |
| 835 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 839 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 836 | AscendC::MicroAPI::RegTensor<float> meanReg, varReg, rstdReg; | 840 | AscendC::MicroAPI::RegTensor<float> meanReg, varReg, rstdReg; |
| @@ -838,41 +842,41 @@ __aicore__ inline void LayerNormGradBase::Normalize(const LocalTensor<float>& ds | |||
| 838 | count = static_cast<uint32_t>(colSize); | 842 | count = static_cast<uint32_t>(colSize); |
| 839 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 843 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 840 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 844 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 841 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 845 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 842 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 846 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 843 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride + 0 * innerLoopStride); | 847 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride + 0 * innerLoopStride); |
| 844 | AscendC::MicroAPI::MaskReg | 848 | AscendC::MicroAPI::MaskReg |
| 845 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 849 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 846 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilon); | 850 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilon); |
| 847 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); | 851 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); |
| 848 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); | 852 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); |
| 849 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); | 853 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); |
| 850 | } | 854 | } |
| 851 | } | 855 | } |
| 852 | } else { | 856 | } else { |
| 853 | __VEC_SCOPE__ | 857 | __VEC_SCOPE__ |
| 854 | { | 858 | { |
| 855 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 859 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 856 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 860 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 857 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 861 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 858 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 862 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 859 | uint32_t count; | 863 | uint32_t count; |
| 860 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 864 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 861 | AscendC::MicroAPI::RegTensor<float> meanReg, varReg, rstdReg; | 865 | AscendC::MicroAPI::RegTensor<float> meanReg, varReg, rstdReg; |
| 862 | AscendC::MicroAPI::MaskReg pMask; | 866 | AscendC::MicroAPI::MaskReg pMask; |
| 863 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 867 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 864 | count = static_cast<uint32_t>(colSize); | 868 | count = static_cast<uint32_t>(colSize); |
| 865 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 869 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 866 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 870 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 867 | AscendC::MicroAPI::MaskReg | 871 | AscendC::MicroAPI::MaskReg |
| 868 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 872 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 869 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilon); | 873 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilon); |
| 870 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 874 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 871 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 875 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 872 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride + j * innerLoopStride); | 876 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride + j * innerLoopStride); |
| 873 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); | 877 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); |
| 874 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); | 878 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); |
| 875 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); | 879 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); |
| 876 | } | 880 | } |
| 877 | } | 881 | } |
| 878 | } | 882 | } |
| @@ -880,16 +884,16 @@ __aicore__ inline void LayerNormGradBase::Normalize(const LocalTensor<float>& ds | |||
| 880 | } | 884 | } |
| 881 | 885 | ||
| 882 | template <typename T> | 886 | template <typename T> |
| 883 | -__aicore__ inline void LayerNormGradBase::StoreTensorForDtypeT(__local_mem__ T* dst, | 887 | +__aicore__ inline void LayerNormGradBase::StoreTensorForDtypeT(__ubuf__ T* dst, |
| 884 | AscendC::MicroAPI::RegTensor<float>& src, | 888 | AscendC::MicroAPI::RegTensor<float>& src, |
| 885 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset) | 889 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset) |
| 886 | { | 890 | { |
| 887 | if constexpr (IsSameType<T, float>::value) { | 891 | if constexpr (IsSameType<T, float>::value) { |
| 888 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); | 892 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); |
| 889 | } else { | 893 | } else { |
| 890 | AscendC::MicroAPI::RegTensor<T> xFp16; | 894 | AscendC::MicroAPI::RegTensor<T> xFp16; |
| 891 | Cast<T, float, castTraitB322B16>(xFp16, src, preg); | 895 | Cast<T, float, castTraitB322B16>(xFp16, src, preg); |
| 892 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); | 896 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); |
| 893 | } | 897 | } |
| 894 | } | 898 | } |
| 895 | 899 | ||
| @@ -21,6 +21,8 @@ | |||
| 21 | 21 | ||
| 22 | namespace LayerNormGrad { | 22 | namespace LayerNormGrad { |
| 23 | using namespace AscendC; | 23 | using namespace AscendC; |
| 24 | +using AscendC::Reg::LoadAlign; | ||
| 25 | +using AscendC::Reg::StoreAlign; | ||
| 24 | 26 | ||
| 25 | template <typename T, typename PD_GAMMA_TYPE> | 27 | template <typename T, typename PD_GAMMA_TYPE> |
| 26 | __aicore__ inline void LayerNormGradGroupedReduceBigMGammaBeta<T, PD_GAMMA_TYPE>::Init( | 28 | __aicore__ inline void LayerNormGradGroupedReduceBigMGammaBeta<T, PD_GAMMA_TYPE>::Init( |
| @@ -387,47 +389,47 @@ __aicore__ inline void LayerNormGradGroupedReduceBigMGammaBeta<T, PD_GAMMA_TYPE> | |||
| 387 | if (innerLoopTimes == 1) { | 389 | if (innerLoopTimes == 1) { |
| 388 | __VEC_SCOPE__ | 390 | __VEC_SCOPE__ |
| 389 | { | 391 | { |
| 390 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 392 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 391 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 393 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 392 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 394 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 393 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 395 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 394 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 396 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 395 | uint32_t count = static_cast<uint32_t>(colSize); | 397 | uint32_t count = static_cast<uint32_t>(colSize); |
| 396 | AscendC::MicroAPI::MaskReg pMask; | 398 | AscendC::MicroAPI::MaskReg pMask; |
| 397 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 399 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 398 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 400 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 399 | AscendC::MicroAPI::RegTensor<float> meanReg; | 401 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 400 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; | 402 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; |
| 401 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 403 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 402 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 404 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 403 | AscendC::MicroAPI::MaskReg | 405 | AscendC::MicroAPI::MaskReg |
| 404 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 406 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 405 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilonTmp); | 407 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilonTmp); |
| 406 | 408 | ||
| 407 | AscendC::MicroAPI::RegTensor<float> xReg; | 409 | AscendC::MicroAPI::RegTensor<float> xReg; |
| 408 | AscendC::MicroAPI::RegTensor<float> dyReg; | 410 | AscendC::MicroAPI::RegTensor<float> dyReg; |
| 409 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + 0 * innerStride); | 411 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + 0 * innerStride); |
| 410 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 412 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 411 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 413 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 412 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + 0 * innerStride); | 414 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + 0 * innerStride); |
| 413 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 415 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 414 | - DataCopy((__local_mem__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); | 416 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); |
| 415 | } | 417 | } |
| 416 | } | 418 | } |
| 417 | } else { | 419 | } else { |
| 418 | __VEC_SCOPE__ | 420 | __VEC_SCOPE__ |
| 419 | { | 421 | { |
| 420 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 422 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 421 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 423 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 422 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 424 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 423 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 425 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 424 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 426 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 425 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 427 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 426 | uint32_t count = static_cast<uint32_t>(colSize); | 428 | uint32_t count = static_cast<uint32_t>(colSize); |
| 427 | AscendC::MicroAPI::RegTensor<float> meanReg; | 429 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 428 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; | 430 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; |
| 429 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 431 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 430 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 432 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 431 | AscendC::MicroAPI::MaskReg | 433 | AscendC::MicroAPI::MaskReg |
| 432 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 434 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 433 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilonTmp); | 435 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilonTmp); |
| @@ -437,12 +439,12 @@ __aicore__ inline void LayerNormGradGroupedReduceBigMGammaBeta<T, PD_GAMMA_TYPE> | |||
| 437 | AscendC::MicroAPI::MaskReg pMask; | 439 | AscendC::MicroAPI::MaskReg pMask; |
| 438 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 440 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 439 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 441 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 440 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + j * innerStride); | 442 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + j * innerStride); |
| 441 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 443 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 442 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 444 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 443 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + j * innerStride); | 445 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + j * innerStride); |
| 444 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 446 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 445 | - DataCopy((__local_mem__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); | 447 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); |
| 446 | } | 448 | } |
| 447 | } | 449 | } |
| 448 | } | 450 | } |
| @@ -845,13 +847,13 @@ __aicore__ inline void LayerNormGradGroupedReduceBigMBackward<T, U>::ComputeDx( | |||
| 845 | if (innerLoopTimes == 1) { | 847 | if (innerLoopTimes == 1) { |
| 846 | __VEC_SCOPE__ | 848 | __VEC_SCOPE__ |
| 847 | { | 849 | { |
| 848 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 850 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 849 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 851 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 850 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 852 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 851 | - __local_mem__ float* gamma = (__local_mem__ float*)gammaTensor.GetPhyAddr(); | 853 | + __ubuf__ float* gamma = (__ubuf__ float*)gammaTensor.GetPhyAddr(); |
| 852 | - __local_mem__ float* sum1 = (__local_mem__ float*)sum1Tensor.GetPhyAddr(); | 854 | + __ubuf__ float* sum1 = (__ubuf__ float*)sum1Tensor.GetPhyAddr(); |
| 853 | - __local_mem__ float* sum2 = (__local_mem__ float*)sum2Tensor.GetPhyAddr(); | 855 | + __ubuf__ float* sum2 = (__ubuf__ float*)sum2Tensor.GetPhyAddr(); |
| 854 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 856 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 855 | uint32_t count; | 857 | uint32_t count; |
| 856 | 858 | ||
| 857 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; | 859 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; |
| @@ -862,15 +864,15 @@ __aicore__ inline void LayerNormGradGroupedReduceBigMBackward<T, U>::ComputeDx( | |||
| 862 | count = static_cast<uint32_t>(colSize); | 864 | count = static_cast<uint32_t>(colSize); |
| 863 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 865 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 864 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 866 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 865 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__local_mem__ float*)sum1 + i); | 867 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__ubuf__ float*)sum1 + i); |
| 866 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__local_mem__ float*)sum2 + i); | 868 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__ubuf__ float*)sum2 + i); |
| 867 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 869 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 868 | AscendC::MicroAPI::MaskReg | 870 | AscendC::MicroAPI::MaskReg |
| 869 | pregRstdAll3 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 871 | pregRstdAll3 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 870 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll3, epsilonTmp); | 872 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll3, epsilonTmp); |
| 871 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerLoopStride + 0 * innerLoopStride); | 873 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerLoopStride + 0 * innerLoopStride); |
| 872 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerLoopStride + 0 * innerLoopStride); | 874 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerLoopStride + 0 * innerLoopStride); |
| 873 | - DataCopy(gammaReg, (__local_mem__ float*)gamma + 0 * innerLoopStride); | 875 | + LoadAlign(gammaReg, (__ubuf__ float*)gamma + 0 * innerLoopStride); |
| 874 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); | 876 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); |
| 875 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); | 877 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); |
| 876 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); | 878 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); |
| @@ -884,13 +886,13 @@ __aicore__ inline void LayerNormGradGroupedReduceBigMBackward<T, U>::ComputeDx( | |||
| 884 | } else { | 886 | } else { |
| 885 | __VEC_SCOPE__ | 887 | __VEC_SCOPE__ |
| 886 | { | 888 | { |
| 887 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 889 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 888 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 890 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 889 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 891 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 890 | - __local_mem__ float* gamma = (__local_mem__ float*)gammaTensor.GetPhyAddr(); | 892 | + __ubuf__ float* gamma = (__ubuf__ float*)gammaTensor.GetPhyAddr(); |
| 891 | - __local_mem__ float* sum1 = (__local_mem__ float*)sum1Tensor.GetPhyAddr(); | 893 | + __ubuf__ float* sum1 = (__ubuf__ float*)sum1Tensor.GetPhyAddr(); |
| 892 | - __local_mem__ float* sum2 = (__local_mem__ float*)sum2Tensor.GetPhyAddr(); | 894 | + __ubuf__ float* sum2 = (__ubuf__ float*)sum2Tensor.GetPhyAddr(); |
| 893 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 895 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 894 | uint32_t count; | 896 | uint32_t count; |
| 895 | 897 | ||
| 896 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; | 898 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; |
| @@ -900,17 +902,17 @@ __aicore__ inline void LayerNormGradGroupedReduceBigMBackward<T, U>::ComputeDx( | |||
| 900 | AscendC::MicroAPI::MaskReg pMask; | 902 | AscendC::MicroAPI::MaskReg pMask; |
| 901 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 903 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 902 | count = static_cast<uint32_t>(colSize); | 904 | count = static_cast<uint32_t>(colSize); |
| 903 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__local_mem__ float*)sum1 + i); | 905 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__ubuf__ float*)sum1 + i); |
| 904 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__local_mem__ float*)sum2 + i); | 906 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__ubuf__ float*)sum2 + i); |
| 905 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 907 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 906 | AscendC::MicroAPI::MaskReg | 908 | AscendC::MicroAPI::MaskReg |
| 907 | pregRstdAll4 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 909 | pregRstdAll4 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 908 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll4, epsilonTmp); | 910 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll4, epsilonTmp); |
| 909 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 911 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 910 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 912 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 911 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerLoopStride + j * innerLoopStride); | 913 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerLoopStride + j * innerLoopStride); |
| 912 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerLoopStride + j * innerLoopStride); | 914 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerLoopStride + j * innerLoopStride); |
| 913 | - DataCopy(gammaReg, (__local_mem__ float*)gamma + j * innerLoopStride); | 915 | + LoadAlign(gammaReg, (__ubuf__ float*)gamma + j * innerLoopStride); |
| 914 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); | 916 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); |
| 915 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); | 917 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); |
| 916 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); | 918 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); |
| @@ -934,4 +936,4 @@ __aicore__ inline void LayerNormGradGroupedReduceBigMBackward<T, U>::Epilogue() | |||
| 934 | } | 936 | } |
| 935 | 937 | ||
| 936 | } // namespace LayerNormGrad | 938 | } // namespace LayerNormGrad |
| 937 | -#endif // LAYER_NORM_GRAD_GROUPED_REDUCE_BIG_M_IMPL_ | 939 | +#endif // LAYER_NORM_GRAD_GROUPED_REDUCE_BIG_M_IMPL_ |
| @@ -20,6 +20,8 @@ | |||
| 20 | 20 | ||
| 21 | namespace LayerNormGrad { | 21 | namespace LayerNormGrad { |
| 22 | using namespace AscendC; | 22 | using namespace AscendC; |
| 23 | +using AscendC::Reg::LoadAlign; | ||
| 24 | +using AscendC::Reg::StoreAlign; | ||
| 23 | 25 | ||
| 24 | template <typename T, typename PD_GAMMA_TYPE> | 26 | template <typename T, typename PD_GAMMA_TYPE> |
| 25 | __aicore__ inline void LayerNormGradGroupedReduceBigNGammaBeta<T, PD_GAMMA_TYPE>::Init( | 27 | __aicore__ inline void LayerNormGradGroupedReduceBigNGammaBeta<T, PD_GAMMA_TYPE>::Init( |
| @@ -285,47 +287,47 @@ __aicore__ inline void LayerNormGradGroupedReduceBigNGammaBeta<T, PD_GAMMA_TYPE> | |||
| 285 | if (innerLoopTimes == 1) { | 287 | if (innerLoopTimes == 1) { |
| 286 | __VEC_SCOPE__ | 288 | __VEC_SCOPE__ |
| 287 | { | 289 | { |
| 288 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 290 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 289 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 291 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 290 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 292 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 291 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 293 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 292 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 294 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 293 | uint32_t count = static_cast<uint32_t>(colSize); | 295 | uint32_t count = static_cast<uint32_t>(colSize); |
| 294 | AscendC::MicroAPI::MaskReg pMask; | 296 | AscendC::MicroAPI::MaskReg pMask; |
| 295 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 297 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 296 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 298 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 297 | AscendC::MicroAPI::RegTensor<float> meanReg; | 299 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 298 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; | 300 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; |
| 299 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 301 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 300 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 302 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 301 | AscendC::MicroAPI::MaskReg | 303 | AscendC::MicroAPI::MaskReg |
| 302 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 304 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 303 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilonTmp); | 305 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilonTmp); |
| 304 | 306 | ||
| 305 | AscendC::MicroAPI::RegTensor<float> xReg; | 307 | AscendC::MicroAPI::RegTensor<float> xReg; |
| 306 | AscendC::MicroAPI::RegTensor<float> dyReg; | 308 | AscendC::MicroAPI::RegTensor<float> dyReg; |
| 307 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + 0 * innerStride); | 309 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + 0 * innerStride); |
| 308 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 310 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 309 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 311 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 310 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + 0 * innerStride); | 312 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + 0 * innerStride); |
| 311 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 313 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 312 | - DataCopy((__local_mem__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); | 314 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); |
| 313 | } | 315 | } |
| 314 | } | 316 | } |
| 315 | } else { | 317 | } else { |
| 316 | __VEC_SCOPE__ | 318 | __VEC_SCOPE__ |
| 317 | { | 319 | { |
| 318 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 320 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 319 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 321 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 320 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 322 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 321 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 323 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 322 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 324 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 323 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 325 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 324 | uint32_t count = static_cast<uint32_t>(colSize); | 326 | uint32_t count = static_cast<uint32_t>(colSize); |
| 325 | AscendC::MicroAPI::RegTensor<float> meanReg; | 327 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 326 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; | 328 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; |
| 327 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 329 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 328 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 330 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 329 | AscendC::MicroAPI::MaskReg | 331 | AscendC::MicroAPI::MaskReg |
| 330 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 332 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 331 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilonTmp); | 333 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilonTmp); |
| @@ -335,12 +337,12 @@ __aicore__ inline void LayerNormGradGroupedReduceBigNGammaBeta<T, PD_GAMMA_TYPE> | |||
| 335 | AscendC::MicroAPI::MaskReg pMask; | 337 | AscendC::MicroAPI::MaskReg pMask; |
| 336 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 338 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 337 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 339 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 338 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + j * innerStride); | 340 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + j * innerStride); |
| 339 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 341 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 340 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 342 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 341 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + j * innerStride); | 343 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + j * innerStride); |
| 342 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 344 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 343 | - DataCopy((__local_mem__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); | 345 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); |
| 344 | } | 346 | } |
| 345 | } | 347 | } |
| 346 | } | 348 | } |
| @@ -718,13 +720,13 @@ __aicore__ inline void LayerNormGradGroupedReduceBigNBackward<T, U>::ComputeDx( | |||
| 718 | if (innerLoopTimes == 1) { | 720 | if (innerLoopTimes == 1) { |
| 719 | __VEC_SCOPE__ | 721 | __VEC_SCOPE__ |
| 720 | { | 722 | { |
| 721 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 723 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 722 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 724 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 723 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 725 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 724 | - __local_mem__ float* gamma = (__local_mem__ float*)gammaTensor.GetPhyAddr(); | 726 | + __ubuf__ float* gamma = (__ubuf__ float*)gammaTensor.GetPhyAddr(); |
| 725 | - __local_mem__ float* sum1 = (__local_mem__ float*)sum1Tensor.GetPhyAddr(); | 727 | + __ubuf__ float* sum1 = (__ubuf__ float*)sum1Tensor.GetPhyAddr(); |
| 726 | - __local_mem__ float* sum2 = (__local_mem__ float*)sum2Tensor.GetPhyAddr(); | 728 | + __ubuf__ float* sum2 = (__ubuf__ float*)sum2Tensor.GetPhyAddr(); |
| 727 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 729 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 728 | uint32_t count; | 730 | uint32_t count; |
| 729 | 731 | ||
| 730 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; | 732 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; |
| @@ -735,15 +737,15 @@ __aicore__ inline void LayerNormGradGroupedReduceBigNBackward<T, U>::ComputeDx( | |||
| 735 | count = static_cast<uint32_t>(colSize); | 737 | count = static_cast<uint32_t>(colSize); |
| 736 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 738 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 737 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 739 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 738 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__local_mem__ float*)sum1 + i); | 740 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__ubuf__ float*)sum1 + i); |
| 739 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__local_mem__ float*)sum2 + i); | 741 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__ubuf__ float*)sum2 + i); |
| 740 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 742 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 741 | AscendC::MicroAPI::MaskReg | 743 | AscendC::MicroAPI::MaskReg |
| 742 | pregRstdAll3 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 744 | pregRstdAll3 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 743 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll3, epsilonTmp); | 745 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll3, epsilonTmp); |
| 744 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerLoopStride + 0 * innerLoopStride); | 746 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerLoopStride + 0 * innerLoopStride); |
| 745 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerLoopStride + 0 * innerLoopStride); | 747 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerLoopStride + 0 * innerLoopStride); |
| 746 | - DataCopy(gammaReg, (__local_mem__ float*)gamma + 0 * innerLoopStride); | 748 | + LoadAlign(gammaReg, (__ubuf__ float*)gamma + 0 * innerLoopStride); |
| 747 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); | 749 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); |
| 748 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); | 750 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); |
| 749 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); | 751 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); |
| @@ -757,13 +759,13 @@ __aicore__ inline void LayerNormGradGroupedReduceBigNBackward<T, U>::ComputeDx( | |||
| 757 | } else { | 759 | } else { |
| 758 | __VEC_SCOPE__ | 760 | __VEC_SCOPE__ |
| 759 | { | 761 | { |
| 760 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 762 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 761 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 763 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 762 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 764 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 763 | - __local_mem__ float* gamma = (__local_mem__ float*)gammaTensor.GetPhyAddr(); | 765 | + __ubuf__ float* gamma = (__ubuf__ float*)gammaTensor.GetPhyAddr(); |
| 764 | - __local_mem__ float* sum1 = (__local_mem__ float*)sum1Tensor.GetPhyAddr(); | 766 | + __ubuf__ float* sum1 = (__ubuf__ float*)sum1Tensor.GetPhyAddr(); |
| 765 | - __local_mem__ float* sum2 = (__local_mem__ float*)sum2Tensor.GetPhyAddr(); | 767 | + __ubuf__ float* sum2 = (__ubuf__ float*)sum2Tensor.GetPhyAddr(); |
| 766 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 768 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 767 | uint32_t count; | 769 | uint32_t count; |
| 768 | 770 | ||
| 769 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; | 771 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; |
| @@ -773,17 +775,17 @@ __aicore__ inline void LayerNormGradGroupedReduceBigNBackward<T, U>::ComputeDx( | |||
| 773 | AscendC::MicroAPI::MaskReg pMask; | 775 | AscendC::MicroAPI::MaskReg pMask; |
| 774 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 776 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 775 | count = static_cast<uint32_t>(colSize); | 777 | count = static_cast<uint32_t>(colSize); |
| 776 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__local_mem__ float*)sum1 + i); | 778 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__ubuf__ float*)sum1 + i); |
| 777 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__local_mem__ float*)sum2 + i); | 779 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__ubuf__ float*)sum2 + i); |
| 778 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 780 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 779 | AscendC::MicroAPI::MaskReg | 781 | AscendC::MicroAPI::MaskReg |
| 780 | pregRstdAll4 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 782 | pregRstdAll4 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 781 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll4, epsilonTmp); | 783 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll4, epsilonTmp); |
| 782 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 784 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 783 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 785 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 784 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerLoopStride + j * innerLoopStride); | 786 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerLoopStride + j * innerLoopStride); |
| 785 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerLoopStride + j * innerLoopStride); | 787 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerLoopStride + j * innerLoopStride); |
| 786 | - DataCopy(gammaReg, (__local_mem__ float*)gamma + j * innerLoopStride); | 788 | + LoadAlign(gammaReg, (__ubuf__ float*)gamma + j * innerLoopStride); |
| 787 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); | 789 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); |
| 788 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); | 790 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); |
| 789 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); | 791 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); |
| @@ -21,6 +21,8 @@ | |||
| 21 | 21 | ||
| 22 | namespace LayerNormGrad { | 22 | namespace LayerNormGrad { |
| 23 | using namespace AscendC; | 23 | using namespace AscendC; |
| 24 | +using AscendC::Reg::LoadAlign; | ||
| 25 | +using AscendC::Reg::StoreAlign; | ||
| 24 | template <typename T, typename PD_GAMMA_TYPE> | 26 | template <typename T, typename PD_GAMMA_TYPE> |
| 25 | __aicore__ inline void LayerNormGradRecomputeGammaBeta<T, PD_GAMMA_TYPE>::Init( | 27 | __aicore__ inline void LayerNormGradRecomputeGammaBeta<T, PD_GAMMA_TYPE>::Init( |
| 26 | GM_ADDR dy, GM_ADDR x, GM_ADDR var, GM_ADDR mean, GM_ADDR pdGamma, GM_ADDR pdBeta, GM_ADDR workspace, | 28 | GM_ADDR dy, GM_ADDR x, GM_ADDR var, GM_ADDR mean, GM_ADDR pdGamma, GM_ADDR pdBeta, GM_ADDR workspace, |
| @@ -291,47 +293,47 @@ __aicore__ inline void LayerNormGradRecomputeGammaBeta<T, PD_GAMMA_TYPE>::Comput | |||
| 291 | if (innerLoopTimes == 1) { | 293 | if (innerLoopTimes == 1) { |
| 292 | __VEC_SCOPE__ | 294 | __VEC_SCOPE__ |
| 293 | { | 295 | { |
| 294 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 296 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 295 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 297 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 296 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 298 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 297 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 299 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 298 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 300 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 299 | uint32_t count = static_cast<uint32_t>(colSize); | 301 | uint32_t count = static_cast<uint32_t>(colSize); |
| 300 | AscendC::MicroAPI::MaskReg pMask; | 302 | AscendC::MicroAPI::MaskReg pMask; |
| 301 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 303 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 302 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 304 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 303 | AscendC::MicroAPI::RegTensor<float> meanReg; | 305 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 304 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; | 306 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; |
| 305 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 307 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 306 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 308 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 307 | AscendC::MicroAPI::MaskReg | 309 | AscendC::MicroAPI::MaskReg |
| 308 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 310 | pregRstdAll1 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 309 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilonTmp); | 311 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll1, epsilonTmp); |
| 310 | 312 | ||
| 311 | AscendC::MicroAPI::RegTensor<float> xReg; | 313 | AscendC::MicroAPI::RegTensor<float> xReg; |
| 312 | AscendC::MicroAPI::RegTensor<float> dyReg; | 314 | AscendC::MicroAPI::RegTensor<float> dyReg; |
| 313 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + 0 * innerStride); | 315 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + 0 * innerStride); |
| 314 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 316 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 315 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 317 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 316 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + 0 * innerStride); | 318 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + 0 * innerStride); |
| 317 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 319 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 318 | - DataCopy((__local_mem__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); | 320 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); |
| 319 | } | 321 | } |
| 320 | } | 322 | } |
| 321 | } else { | 323 | } else { |
| 322 | __VEC_SCOPE__ | 324 | __VEC_SCOPE__ |
| 323 | { | 325 | { |
| 324 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 326 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 325 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 327 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 326 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 328 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 327 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 329 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 328 | - __local_mem__ float* var = (__local_mem__ float*)varTensor.GetPhyAddr(); | 330 | + __ubuf__ float* var = (__ubuf__ float*)varTensor.GetPhyAddr(); |
| 329 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 331 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 330 | uint32_t count = static_cast<uint32_t>(colSize); | 332 | uint32_t count = static_cast<uint32_t>(colSize); |
| 331 | AscendC::MicroAPI::RegTensor<float> meanReg; | 333 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 332 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; | 334 | AscendC::MicroAPI::RegTensor<float> varReg, rstdReg; |
| 333 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 335 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 334 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__local_mem__ float*)var + i); | 336 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(varReg, (__ubuf__ float*)var + i); |
| 335 | AscendC::MicroAPI::MaskReg | 337 | AscendC::MicroAPI::MaskReg |
| 336 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 338 | pregRstdAll2 = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 337 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilonTmp); | 339 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdReg, pregRstdAll2, epsilonTmp); |
| @@ -341,12 +343,12 @@ __aicore__ inline void LayerNormGradRecomputeGammaBeta<T, PD_GAMMA_TYPE>::Comput | |||
| 341 | AscendC::MicroAPI::MaskReg pMask; | 343 | AscendC::MicroAPI::MaskReg pMask; |
| 342 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 344 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 343 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 345 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 344 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + j * innerStride); | 346 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + j * innerStride); |
| 345 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 347 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 346 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 348 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 347 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + j * innerStride); | 349 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + j * innerStride); |
| 348 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 350 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 349 | - DataCopy((__local_mem__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); | 351 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); |
| 350 | } | 352 | } |
| 351 | } | 353 | } |
| 352 | } | 354 | } |
| @@ -4,10 +4,10 @@ | |||
| 4 | # CANN Open Software License Agreement Version 2.0 (the "License"). | 4 | # CANN Open Software License Agreement Version 2.0 (the "License"). |
| 5 | # Please refer to the License for details. You may not use this file except in compliance with the License. | 5 | # Please refer to the License for details. You may not use this file except in compliance with the License. |
| 6 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, | 6 | # THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, |
| 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 | 10 | ||
| 11 | if ((UT_TEST_ALL OR OP_KERNEL_UT) AND NOT UT_DONE) | 11 | if ((UT_TEST_ALL OR OP_KERNEL_UT) AND NOT UT_DONE) |
| 12 | - AddOpTestCase(layer_norm_grad "ascend950pr_9599" "-DDTYPE_DY=float -DDTYPE_GAMMA=float -DDTYPE_PD_GAMMA=float" ) | 12 | + AddOpTestCase(layer_norm_grad "ascend950pr_9599" "-DDTYPE_DY=float -DDTYPE_GAMMA=float -DDTYPE_PD_GAMMA=float" norm_common) |
| 13 | endif() | 13 | endif() |
| @@ -31,6 +31,10 @@ using namespace AscendC; | |||
| 31 | using namespace NormCommon; | 31 | using namespace NormCommon; |
| 32 | using namespace NormCommon::NormCommonRegbase; | 32 | using namespace NormCommon::NormCommonRegbase; |
| 33 | using namespace LayerNormGradV3::Arith; | 33 | using namespace LayerNormGradV3::Arith; |
| 34 | +using AscendC::Reg::LoadAlign; | ||
| 35 | +using AscendC::Reg::Move; | ||
| 36 | +using AscendC::Reg::Reduce; | ||
| 37 | +using AscendC::Reg::StoreAlign; | ||
| 34 | 38 | ||
| 35 | constexpr static AscendC::MicroAPI::CastTrait castTraitB162B32 = { | 39 | constexpr static AscendC::MicroAPI::CastTrait castTraitB162B32 = { |
| 36 | AscendC::MicroAPI::RegLayout::ZERO, | 40 | AscendC::MicroAPI::RegLayout::ZERO, |
| @@ -162,7 +166,7 @@ public: | |||
| 162 | const LocalTensor<float>& sum2Tensor, const LocalTensor<float>& rstdTensor, const int64_t rowSize, | 166 | const LocalTensor<float>& sum2Tensor, const LocalTensor<float>& rstdTensor, const int64_t rowSize, |
| 163 | const int64_t colSize, const int64_t stride, const int64_t fullColSize); | 167 | const int64_t colSize, const int64_t stride, const int64_t fullColSize); |
| 164 | template <typename T> | 168 | template <typename T> |
| 165 | - __aicore__ inline static void StoreTensorForDtypeT(__local_mem__ T* dst, AscendC::MicroAPI::RegTensor<float>& src, | 169 | + __aicore__ inline static void StoreTensorForDtypeT(__ubuf__ T* dst, AscendC::MicroAPI::RegTensor<float>& src, |
| 166 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); | 170 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset); |
| 167 | 171 | ||
| 168 | protected: | 172 | protected: |
| @@ -214,8 +218,8 @@ __aicore__ inline void LayerNormGradV3Base::CastToFp32From(const LocalTensor<flo | |||
| 214 | if (innerLoopTimes == 1) { | 218 | if (innerLoopTimes == 1) { |
| 215 | __VEC_SCOPE__ | 219 | __VEC_SCOPE__ |
| 216 | { | 220 | { |
| 217 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 221 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 218 | - __local_mem__ T* src = (__local_mem__ T*)srcTensor.GetPhyAddr(); | 222 | + __ubuf__ T* src = (__ubuf__ T*)srcTensor.GetPhyAddr(); |
| 219 | uint32_t count; | 223 | uint32_t count; |
| 220 | AscendC::MicroAPI::RegTensor<float> fp32Reg; | 224 | AscendC::MicroAPI::RegTensor<float> fp32Reg; |
| 221 | AscendC::MicroAPI::RegTensor<T> b16Reg; | 225 | AscendC::MicroAPI::RegTensor<T> b16Reg; |
| @@ -223,17 +227,17 @@ __aicore__ inline void LayerNormGradV3Base::CastToFp32From(const LocalTensor<flo | |||
| 223 | count = static_cast<uint32_t>(colSize); | 227 | count = static_cast<uint32_t>(colSize); |
| 224 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 228 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 225 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 229 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 226 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( | 230 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( |
| 227 | - b16Reg, (__local_mem__ T*)src + i * outerLoopSrcStride + 0 * innerLoopStride); | 231 | + b16Reg, (__ubuf__ T*)src + i * outerLoopSrcStride + 0 * innerLoopStride); |
| 228 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); | 232 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); |
| 229 | - DataCopy((__local_mem__ float*)dst + i * outerLoopDstStride + 0 * innerLoopStride, fp32Reg, pMask); | 233 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopDstStride + 0 * innerLoopStride, fp32Reg, pMask); |
| 230 | } | 234 | } |
| 231 | } | 235 | } |
| 232 | } else { | 236 | } else { |
| 233 | __VEC_SCOPE__ | 237 | __VEC_SCOPE__ |
| 234 | { | 238 | { |
| 235 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 239 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 236 | - __local_mem__ T* src = (__local_mem__ T*)srcTensor.GetPhyAddr(); | 240 | + __ubuf__ T* src = (__ubuf__ T*)srcTensor.GetPhyAddr(); |
| 237 | uint32_t count; | 241 | uint32_t count; |
| 238 | AscendC::MicroAPI::RegTensor<float> fp32Reg; | 242 | AscendC::MicroAPI::RegTensor<float> fp32Reg; |
| 239 | AscendC::MicroAPI::RegTensor<T> b16Reg; | 243 | AscendC::MicroAPI::RegTensor<T> b16Reg; |
| @@ -242,11 +246,10 @@ __aicore__ inline void LayerNormGradV3Base::CastToFp32From(const LocalTensor<flo | |||
| 242 | count = static_cast<uint32_t>(colSize); | 246 | count = static_cast<uint32_t>(colSize); |
| 243 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 247 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 244 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 248 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 245 | - DataCopy<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( | 249 | + LoadAlign<T, AscendC::MicroAPI::LoadDist::DIST_UNPACK_B16>( |
| 246 | - b16Reg, (__local_mem__ T*)src + i * outerLoopSrcStride + j * innerLoopStride); | 250 | + b16Reg, (__ubuf__ T*)src + i * outerLoopSrcStride + j * innerLoopStride); |
| 247 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); | 251 | Cast<float, T, castTraitB162B32>(fp32Reg, b16Reg, pMask); |
| 248 | - DataCopy((__local_mem__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, | 252 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopDstStride + j * innerLoopStride, fp32Reg, pMask); |
| 249 | - pMask); | ||
| 250 | } | 253 | } |
| 251 | } | 254 | } |
| 252 | } | 255 | } |
| @@ -327,8 +330,8 @@ __aicore__ inline void LayerNormGradV3Base::CopyUB2UBWithCast(const LocalTensor< | |||
| 327 | if constexpr (IsSameType<T, float>::value) { | 330 | if constexpr (IsSameType<T, float>::value) { |
| 328 | CopyUB2UB(dstTensor, srcTensor, count); | 331 | CopyUB2UB(dstTensor, srcTensor, count); |
| 329 | } else { | 332 | } else { |
| 330 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 333 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 331 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 334 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 332 | 335 | ||
| 333 | uint32_t cnt = count; | 336 | uint32_t cnt = count; |
| 334 | uint16_t loopNum = CeilDiv(cnt, VL_FP32); | 337 | uint16_t loopNum = CeilDiv(cnt, VL_FP32); |
| @@ -341,10 +344,10 @@ __aicore__ inline void LayerNormGradV3Base::CopyUB2UBWithCast(const LocalTensor< | |||
| 341 | for (uint16_t k = 0; k < loopNum; k++) { | 344 | for (uint16_t k = 0; k < loopNum; k++) { |
| 342 | pregMask = UpdateMask<float>(sreg); | 345 | pregMask = UpdateMask<float>(sreg); |
| 343 | uint32_t offset = k * VL_FP32; | 346 | uint32_t offset = k * VL_FP32; |
| 344 | - DataCopy<float, LoadDist::DIST_NORM>(srcReg, (__local_mem__ float*)src + offset); | 347 | + LoadAlign<float, LoadDist::DIST_NORM>(srcReg, (__ubuf__ float*)src + offset); |
| 345 | 348 | ||
| 346 | Cast<T, float, castTraitB322B16>(xFp16, srcReg, pregMask); | 349 | Cast<T, float, castTraitB322B16>(xFp16, srcReg, pregMask); |
| 347 | - DataCopy<T, StoreDist::DIST_PACK_B32>(((__local_mem__ T*)dst) + offset, xFp16, pregMask); | 350 | + StoreAlign<T, StoreDist::DIST_PACK_B32>(((__ubuf__ T*)dst) + offset, xFp16, pregMask); |
| 348 | } | 351 | } |
| 349 | } | 352 | } |
| 350 | } | 353 | } |
| @@ -361,19 +364,19 @@ __aicore__ inline void LayerNormGradV3Base::VectorAdd(const LocalTensor<float>& | |||
| 361 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); | 364 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); |
| 362 | __VEC_SCOPE__ | 365 | __VEC_SCOPE__ |
| 363 | { | 366 | { |
| 364 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 367 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 365 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 368 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 366 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 369 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 367 | uint32_t sreg = static_cast<uint32_t>(count); | 370 | uint32_t sreg = static_cast<uint32_t>(count); |
| 368 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 371 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 369 | AscendC::MicroAPI::MaskReg pMask; | 372 | AscendC::MicroAPI::MaskReg pMask; |
| 370 | for (uint16_t i = 0; i < loopTimes; ++i) { | 373 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 371 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 374 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 372 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32); | 375 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32); |
| 373 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32); | 376 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32); |
| 374 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 377 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 375 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 378 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 376 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, aReg, pMask); | 379 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, aReg, pMask); |
| 377 | } | 380 | } |
| 378 | } | 381 | } |
| 379 | } | 382 | } |
| @@ -391,38 +394,38 @@ __aicore__ inline void LayerNormGradV3Base::VectorAdd(const LocalTensor<float>& | |||
| 391 | if (innerLoopTimes == 1) { | 394 | if (innerLoopTimes == 1) { |
| 392 | __VEC_SCOPE__ | 395 | __VEC_SCOPE__ |
| 393 | { | 396 | { |
| 394 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 397 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 395 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 398 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 396 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 399 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 397 | uint32_t count = nSize; | 400 | uint32_t count = nSize; |
| 398 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 401 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 399 | AscendC::MicroAPI::MaskReg pMask; | 402 | AscendC::MicroAPI::MaskReg pMask; |
| 400 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 403 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 401 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 404 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 402 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); | 405 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); |
| 403 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride + 0 * innerLoopStride); | 406 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride + 0 * innerLoopStride); |
| 404 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 407 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 405 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 408 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 406 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, aReg, pMask); | 409 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, aReg, pMask); |
| 407 | } | 410 | } |
| 408 | } | 411 | } |
| 409 | } else { | 412 | } else { |
| 410 | __VEC_SCOPE__ | 413 | __VEC_SCOPE__ |
| 411 | { | 414 | { |
| 412 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 415 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 413 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 416 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 414 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 417 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 415 | uint32_t count = nSize; | 418 | uint32_t count = nSize; |
| 416 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 419 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 417 | AscendC::MicroAPI::MaskReg pMask; | 420 | AscendC::MicroAPI::MaskReg pMask; |
| 418 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 421 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 419 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 422 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 420 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 423 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 421 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride); | 424 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride); |
| 422 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride + j * innerLoopStride); | 425 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride + j * innerLoopStride); |
| 423 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 426 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 424 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 427 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 425 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask); | 428 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, aReg, pMask); |
| 426 | } | 429 | } |
| 427 | } | 430 | } |
| 428 | } | 431 | } |
| @@ -440,19 +443,19 @@ __aicore__ inline void LayerNormGradV3Base::VectorMul(const LocalTensor<float>& | |||
| 440 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); | 443 | uint16_t loopTimes = CeilDiv(static_cast<int64_t>(count * sizeof(float)), static_cast<int64_t>(GetVRegSize())); |
| 441 | __VEC_SCOPE__ | 444 | __VEC_SCOPE__ |
| 442 | { | 445 | { |
| 443 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 446 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 444 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 447 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 445 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 448 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 446 | uint32_t sreg = static_cast<uint32_t>(count); | 449 | uint32_t sreg = static_cast<uint32_t>(count); |
| 447 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, mulReg; | 450 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, mulReg; |
| 448 | AscendC::MicroAPI::MaskReg pMask; | 451 | AscendC::MicroAPI::MaskReg pMask; |
| 449 | 452 | ||
| 450 | for (uint16_t i = 0; i < loopTimes; ++i) { | 453 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 451 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 454 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 452 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * VL_FP32); | 455 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * VL_FP32); |
| 453 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * VL_FP32); | 456 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * VL_FP32); |
| 454 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(mulReg, aReg, bReg, pMask); | 457 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(mulReg, aReg, bReg, pMask); |
| 455 | - DataCopy((__local_mem__ float*)dst + i * VL_FP32, mulReg, pMask); | 458 | + StoreAlign((__ubuf__ float*)dst + i * VL_FP32, mulReg, pMask); |
| 456 | } | 459 | } |
| 457 | } | 460 | } |
| 458 | } | 461 | } |
| @@ -476,36 +479,36 @@ __aicore__ inline void LayerNormGradV3Base::NlastBroadcastMul(const LocalTensor< | |||
| 476 | if (innerLoopTimes == 1) { | 479 | if (innerLoopTimes == 1) { |
| 477 | __VEC_SCOPE__ | 480 | __VEC_SCOPE__ |
| 478 | { | 481 | { |
| 479 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 482 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 480 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 483 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 481 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 484 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 482 | uint32_t count = static_cast<uint32_t>(aSize); | 485 | uint32_t count = static_cast<uint32_t>(aSize); |
| 483 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 486 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 484 | AscendC::MicroAPI::MaskReg pMask; | 487 | AscendC::MicroAPI::MaskReg pMask; |
| 485 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 488 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 486 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 489 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 487 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride); | 490 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride); |
| 488 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); | 491 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + 0 * innerLoopStride); |
| 489 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 492 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 490 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); | 493 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); |
| 491 | } | 494 | } |
| 492 | } | 495 | } |
| 493 | } else { | 496 | } else { |
| 494 | __VEC_SCOPE__ | 497 | __VEC_SCOPE__ |
| 495 | { | 498 | { |
| 496 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 499 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 497 | - __local_mem__ float* src0 = (__local_mem__ float*)src0Tensor.GetPhyAddr(); | 500 | + __ubuf__ float* src0 = (__ubuf__ float*)src0Tensor.GetPhyAddr(); |
| 498 | - __local_mem__ float* src1 = (__local_mem__ float*)src1Tensor.GetPhyAddr(); | 501 | + __ubuf__ float* src1 = (__ubuf__ float*)src1Tensor.GetPhyAddr(); |
| 499 | uint32_t count = static_cast<uint32_t>(aSize); | 502 | uint32_t count = static_cast<uint32_t>(aSize); |
| 500 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 503 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 501 | AscendC::MicroAPI::MaskReg pMask; | 504 | AscendC::MicroAPI::MaskReg pMask; |
| 502 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 505 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 503 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 506 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 504 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * outerLoopStride); | 507 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * outerLoopStride); |
| 505 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 508 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 506 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * outerLoopStride + j * innerLoopStride); | 509 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * outerLoopStride + j * innerLoopStride); |
| 507 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 510 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 508 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); | 511 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); |
| 509 | } | 512 | } |
| 510 | } | 513 | } |
| 511 | } | 514 | } |
| @@ -532,40 +535,40 @@ __aicore__ inline void LayerNormGradV3Base::LastReduceSumSmallR(const LocalTenso | |||
| 532 | if (rSize <= VL_FP32) { | 535 | if (rSize <= VL_FP32) { |
| 533 | __VEC_SCOPE__ | 536 | __VEC_SCOPE__ |
| 534 | { | 537 | { |
| 535 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 538 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 536 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 539 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 537 | uint32_t count = static_cast<uint32_t>(rSize); | 540 | uint32_t count = static_cast<uint32_t>(rSize); |
| 538 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 541 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 539 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 542 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 540 | - AscendC::MicroAPI::UnalignReg UReg; | 543 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 541 | for (uint16_t i = 0; i < loopTimes; ++i) { | 544 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 542 | - DataCopy(aReg, (__local_mem__ float*)src + i * stride); | 545 | + LoadAlign(aReg, (__ubuf__ float*)src + i * stride); |
| 543 | - ReduceSum(bReg, aReg, pMask); | 546 | + Reduce<ReduceType::SUM>(bReg, aReg, pMask); |
| 544 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 547 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 545 | } | 548 | } |
| 546 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 549 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 547 | } | 550 | } |
| 548 | } else { | 551 | } else { |
| 549 | __VEC_SCOPE__ | 552 | __VEC_SCOPE__ |
| 550 | { | 553 | { |
| 551 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 554 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 552 | - __local_mem__ float* src0 = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 555 | + __ubuf__ float* src0 = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 553 | - __local_mem__ float* src1 = (__local_mem__ float*)srcTensor.GetPhyAddr() + VL_FP32; | 556 | + __ubuf__ float* src1 = (__ubuf__ float*)srcTensor.GetPhyAddr() + VL_FP32; |
| 554 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); | 557 | uint32_t count = static_cast<uint32_t>(rSize - VL_FP32); |
| 555 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 558 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 556 | - AscendC::MicroAPI::UnalignReg UReg; | 559 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 557 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 560 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 558 | AscendC::MicroAPI::MaskReg | 561 | AscendC::MicroAPI::MaskReg |
| 559 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 562 | pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 560 | for (uint16_t i = 0; i < loopTimes; ++i) { | 563 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 561 | - DataCopy(aReg, (__local_mem__ float*)src0 + i * stride); | 564 | + LoadAlign(aReg, (__ubuf__ float*)src0 + i * stride); |
| 562 | - DataCopy(bReg, (__local_mem__ float*)src1 + i * stride); | 565 | + LoadAlign(bReg, (__ubuf__ float*)src1 + i * stride); |
| 563 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 566 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 564 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 567 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 565 | - ReduceSum(bReg, aReg, pFull); | 568 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 566 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 569 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 567 | } | 570 | } |
| 568 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 571 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 569 | } | 572 | } |
| 570 | } | 573 | } |
| 571 | } | 574 | } |
| @@ -609,43 +612,43 @@ __aicore__ inline void LayerNormGradV3Base::LastReduceSum(const LocalTensor<floa | |||
| 609 | 612 | ||
| 610 | __VEC_SCOPE__ | 613 | __VEC_SCOPE__ |
| 611 | { | 614 | { |
| 612 | - __local_mem__ float* dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr(); | 615 | + __ubuf__ float* dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr(); |
| 613 | - __local_mem__ float* foldSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 616 | + __ubuf__ float* foldSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 614 | - __local_mem__ float* foldSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; | 617 | + __ubuf__ float* foldSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + foldSrcBOffset; |
| 615 | - __local_mem__ float* tailSrcA = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; | 618 | + __ubuf__ float* tailSrcA = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcAOffset; |
| 616 | - __local_mem__ float* tailSrcB = (__local_mem__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; | 619 | + __ubuf__ float* tailSrcB = (__ubuf__ float*)srcTensor.GetPhyAddr() + tailSrcBOffset; |
| 617 | - __local_mem__ float* unFoldSrc = (__local_mem__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; | 620 | + __ubuf__ float* unFoldSrc = (__ubuf__ float*)srcTensor.GetPhyAddr() + unFoldSrcOffset; |
| 618 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); | 621 | AscendC::MicroAPI::MaskReg pFull = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>(); |
| 619 | - AscendC::MicroAPI::UnalignReg UReg; | 622 | + AscendC::MicroAPI::UnalignRegForStore UReg; |
| 620 | 623 | ||
| 621 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 624 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 622 | - dst = (__local_mem__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; | 625 | + dst = (__ubuf__ float*)reduceSumTempTensor.GetPhyAddr() + i * outerLoopDstStride; |
| 623 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { | 626 | for (uint16_t j = 0; j < mainFoldLoopTimes; ++j) { |
| 624 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; | 627 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg, dReg; |
| 625 | - DataCopy(aReg, (__local_mem__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); | 628 | + LoadAlign(aReg, (__ubuf__ float*)foldSrcA + i * outerLoopStride + j * innerLoopStride); |
| 626 | - DataCopy(bReg, (__local_mem__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); | 629 | + LoadAlign(bReg, (__ubuf__ float*)foldSrcB + i * outerLoopStride + j * innerLoopStride); |
| 627 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); | 630 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pFull); |
| 628 | - ReduceSum(dReg, cReg, pFull); | 631 | + Reduce<ReduceType::SUM>(dReg, cReg, pFull); |
| 629 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, dReg, UReg, 1); | 632 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, dReg, UReg, 1); |
| 630 | } | 633 | } |
| 631 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { | 634 | for (uint16_t j = 0; j < tailFoldLoopTimes; ++j) { |
| 632 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); | 635 | uint32_t count = static_cast<uint32_t>(tailFoldElemCount); |
| 633 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 636 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 634 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 637 | AscendC::MicroAPI::MaskReg pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 635 | - DataCopy(aReg, (__local_mem__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); | 638 | + LoadAlign(aReg, (__ubuf__ float*)tailSrcA + i * outerLoopStride + j * innerLoopStride); |
| 636 | - DataCopy(bReg, (__local_mem__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); | 639 | + LoadAlign(bReg, (__ubuf__ float*)tailSrcB + i * outerLoopStride + j * innerLoopStride); |
| 637 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); | 640 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, aReg, bReg, pMask); |
| 638 | - Copy<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); | 641 | + Move<float, AscendC::MicroAPI::MaskMergeMode::MERGING>(aReg, cReg, pMask); |
| 639 | - ReduceSum(bReg, aReg, pFull); | 642 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 640 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 643 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 641 | } | 644 | } |
| 642 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { | 645 | for (uint16_t j = 0; j < unFoldLoopTimes; ++j) { |
| 643 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 646 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 644 | - DataCopy(aReg, (__local_mem__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); | 647 | + LoadAlign(aReg, (__ubuf__ float*)unFoldSrc + i * outerLoopStride + j * innerLoopStride); |
| 645 | - ReduceSum(bReg, aReg, pFull); | 648 | + Reduce<ReduceType::SUM>(bReg, aReg, pFull); |
| 646 | - AscendC::MicroAPI::DataCopyUnAlign((__local_mem__ float*&)dst, bReg, UReg, 1); | 649 | + AscendC::MicroAPI::StoreUnAlign((__ubuf__ float*&)dst, bReg, UReg, 1); |
| 647 | } | 650 | } |
| 648 | - AscendC::MicroAPI::DataCopyUnAlignPost((__local_mem__ float*&)dst, UReg, 0); | 651 | + AscendC::MicroAPI::StoreUnAlignPost((__ubuf__ float*&)dst, UReg, 0); |
| 649 | } | 652 | } |
| 650 | } | 653 | } |
| 651 | LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); | 654 | LastReduceSumSmallR(dstTensor, reduceSumTempTensor, aSize, foldPoint, outerLoopDstStride); |
| @@ -653,13 +656,13 @@ __aicore__ inline void LayerNormGradV3Base::LastReduceSum(const LocalTensor<floa | |||
| 653 | 656 | ||
| 654 | template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1> | 657 | template <uint32_t RSize, int32_t TailCount = -1, int32_t Index = 0, int32_t Depth = 1> |
| 655 | struct NlastDichotomyAdd { | 658 | struct NlastDichotomyAdd { |
| 656 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 659 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 657 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 660 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 658 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 661 | + uint32_t stride) |
| 659 | { | 662 | { |
| 660 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 663 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 661 | - __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO; | 664 | + __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO; |
| 662 | - __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO; | 665 | + __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO; |
| 663 | if constexpr (TailCount <= 0) { | 666 | if constexpr (TailCount <= 0) { |
| 664 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, | 667 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, |
| 665 | stride * CONST_TWO); | 668 | stride * CONST_TWO); |
| @@ -667,13 +670,13 @@ struct NlastDichotomyAdd { | |||
| 667 | } | 670 | } |
| 668 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 671 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 669 | } | 672 | } |
| 670 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 673 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 671 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 674 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 672 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 675 | + uint32_t stride, uint32_t offset) |
| 673 | { | 676 | { |
| 674 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 677 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 675 | - __local_mem__ float* srcAOffset = srcA + stride * CONST_TWO; | 678 | + __ubuf__ float* srcAOffset = srcA + stride * CONST_TWO; |
| 676 | - __local_mem__ float* srcBOffset = srcB + stride * CONST_TWO; | 679 | + __ubuf__ float* srcBOffset = srcB + stride * CONST_TWO; |
| 677 | if constexpr (TailCount <= 0) { | 680 | if constexpr (TailCount <= 0) { |
| 678 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, | 681 | NlastDichotomyAdd<(RSize + 1) / CONST_TWO>::LoadAndAccumulate(aReg, srcA, srcAOffset, pMask, |
| 679 | stride * CONST_TWO, offset); | 682 | stride * CONST_TWO, offset); |
| @@ -691,49 +694,49 @@ struct NlastDichotomyAdd { | |||
| 691 | 694 | ||
| 692 | template <int32_t TailCount, int32_t Index, int32_t Depth> | 695 | template <int32_t TailCount, int32_t Index, int32_t Depth> |
| 693 | struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { | 696 | struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { |
| 694 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 697 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 695 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 698 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 696 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 699 | + uint32_t stride) |
| 697 | { | 700 | { |
| 698 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 701 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 699 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 702 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 700 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 703 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 701 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 704 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 702 | } | 705 | } |
| 703 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 706 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 704 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 707 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 705 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 708 | + uint32_t stride, uint32_t offset) |
| 706 | { | 709 | { |
| 707 | if constexpr (TailCount <= 0) { | 710 | if constexpr (TailCount <= 0) { |
| 708 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 711 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 709 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 712 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 710 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 713 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 711 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 714 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 712 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 715 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 713 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 716 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 714 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 717 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 715 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 718 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 716 | } else { | 719 | } else { |
| 717 | if constexpr (Index + Depth < TailCount) { | 720 | if constexpr (Index + Depth < TailCount) { |
| 718 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 721 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 719 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 722 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 720 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 723 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 721 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 724 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 722 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 725 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 723 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 726 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 724 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 727 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 725 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 728 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 726 | } else if constexpr (Index < TailCount) { | 729 | } else if constexpr (Index < TailCount) { |
| 727 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 730 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 728 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 731 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 729 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 732 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 730 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 733 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 731 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 734 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 732 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 735 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 733 | } else { | 736 | } else { |
| 734 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 737 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 735 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 738 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 736 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 739 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 737 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 740 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 738 | } | 741 | } |
| 739 | } | 742 | } |
| @@ -742,25 +745,25 @@ struct NlastDichotomyAdd<CONST_TWO, TailCount, Index, Depth> { | |||
| 742 | 745 | ||
| 743 | template <> | 746 | template <> |
| 744 | struct NlastDichotomyAdd<CONST_TWO> { | 747 | struct NlastDichotomyAdd<CONST_TWO> { |
| 745 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 748 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 746 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 749 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 747 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 750 | + uint32_t stride) |
| 748 | { | 751 | { |
| 749 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 752 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 750 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 753 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 751 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 754 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 752 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 755 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 753 | } | 756 | } |
| 754 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 757 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 755 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 758 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 756 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride, uint32_t offset) | 759 | + uint32_t stride, uint32_t offset) |
| 757 | { | 760 | { |
| 758 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 761 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 759 | - DataCopy(aReg, (__local_mem__ float*)srcA); | 762 | + LoadAlign(aReg, (__ubuf__ float*)srcA); |
| 760 | - DataCopy(bReg, (__local_mem__ float*)srcA + offset); | 763 | + LoadAlign(bReg, (__ubuf__ float*)srcA + offset); |
| 761 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 764 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 762 | - DataCopy(bReg, (__local_mem__ float*)srcB); | 765 | + LoadAlign(bReg, (__ubuf__ float*)srcB); |
| 763 | - DataCopy(cReg, (__local_mem__ float*)srcB + offset); | 766 | + LoadAlign(cReg, (__ubuf__ float*)srcB + offset); |
| 764 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); | 767 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, bReg, cReg, pMask); |
| 765 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); | 768 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(acc, aReg, bReg, pMask); |
| 766 | } | 769 | } |
| @@ -768,11 +771,11 @@ struct NlastDichotomyAdd<CONST_TWO> { | |||
| 768 | 771 | ||
| 769 | template <> | 772 | template <> |
| 770 | struct NlastDichotomyAdd<1> { | 773 | struct NlastDichotomyAdd<1> { |
| 771 | - __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, | 774 | + __aicore__ static inline void LoadAndAccumulate(AscendC::MicroAPI::RegTensor<float>& acc, __ubuf__ float*& srcA, |
| 772 | - __local_mem__ float*& srcA, __local_mem__ float*& srcB, | 775 | + __ubuf__ float*& srcB, AscendC::MicroAPI::MaskReg& pMask, |
| 773 | - AscendC::MicroAPI::MaskReg& pMask, uint32_t stride) | 776 | + uint32_t stride) |
| 774 | { | 777 | { |
| 775 | - DataCopy(acc, (__local_mem__ float*)srcA); | 778 | + LoadAlign(acc, (__ubuf__ float*)srcA); |
| 776 | } | 779 | } |
| 777 | }; | 780 | }; |
| 778 | 781 | ||
| @@ -800,56 +803,56 @@ __aicore__ inline void LayerNormGradV3Base::UpdateCache(const LocalTensor<float> | |||
| 800 | if (innerLoopTimes == 1) { | 803 | if (innerLoopTimes == 1) { |
| 801 | __VEC_SCOPE__ | 804 | __VEC_SCOPE__ |
| 802 | { | 805 | { |
| 803 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 806 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 804 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 807 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 805 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 808 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 806 | uint32_t sreg = static_cast<uint32_t>(count); | 809 | uint32_t sreg = static_cast<uint32_t>(count); |
| 807 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 810 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 808 | AscendC::MicroAPI::MaskReg pMask; | 811 | AscendC::MicroAPI::MaskReg pMask; |
| 809 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 812 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 810 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 813 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 811 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 814 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 812 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); | 815 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); |
| 813 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 816 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 814 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 817 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 815 | } | 818 | } |
| 816 | } | 819 | } |
| 817 | } else if (innerLoopTimes == CONST_TWO) { | 820 | } else if (innerLoopTimes == CONST_TWO) { |
| 818 | __VEC_SCOPE__ | 821 | __VEC_SCOPE__ |
| 819 | { | 822 | { |
| 820 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 823 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 821 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 824 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 822 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 825 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 823 | uint32_t sreg = static_cast<uint32_t>(count); | 826 | uint32_t sreg = static_cast<uint32_t>(count); |
| 824 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 827 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 825 | AscendC::MicroAPI::MaskReg pMask; | 828 | AscendC::MicroAPI::MaskReg pMask; |
| 826 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 829 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 827 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 830 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 828 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 831 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 829 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); | 832 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride); |
| 830 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 833 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 831 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + 1 * innerLoopStride); | 834 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + 1 * innerLoopStride); |
| 832 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 835 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 833 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 836 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 834 | } | 837 | } |
| 835 | } | 838 | } |
| 836 | } else { | 839 | } else { |
| 837 | __VEC_SCOPE__ | 840 | __VEC_SCOPE__ |
| 838 | { | 841 | { |
| 839 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 842 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 840 | - __local_mem__ float* cah = (__local_mem__ float*)dstTensor.GetPhyAddr() + cacheID * stride; | 843 | + __ubuf__ float* cah = (__ubuf__ float*)dstTensor.GetPhyAddr() + cacheID * stride; |
| 841 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 844 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 842 | uint32_t sreg = static_cast<uint32_t>(count); | 845 | uint32_t sreg = static_cast<uint32_t>(count); |
| 843 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; | 846 | AscendC::MicroAPI::RegTensor<float> aReg, bReg; |
| 844 | AscendC::MicroAPI::MaskReg pMask; | 847 | AscendC::MicroAPI::MaskReg pMask; |
| 845 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 848 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 846 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); | 849 | pMask = AscendC::MicroAPI::UpdateMask<float>(sreg); |
| 847 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride); | 850 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride); |
| 848 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 851 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 849 | - DataCopy(bReg, (__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride); | 852 | + LoadAlign(bReg, (__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride); |
| 850 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); | 853 | Add<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(aReg, aReg, bReg, pMask); |
| 851 | } | 854 | } |
| 852 | - DataCopy((__local_mem__ float*)cah + i * outerLoopStride, aReg, pMask); | 855 | + StoreAlign((__ubuf__ float*)cah + i * outerLoopStride, aReg, pMask); |
| 853 | } | 856 | } |
| 854 | } | 857 | } |
| 855 | } | 858 | } |
| @@ -870,10 +873,10 @@ __aicore__ inline void LayerNormGradV3Base::Normalize(const LocalTensor<float>& | |||
| 870 | if (innerLoopTimes == 1) { | 873 | if (innerLoopTimes == 1) { |
| 871 | __VEC_SCOPE__ | 874 | __VEC_SCOPE__ |
| 872 | { | 875 | { |
| 873 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 876 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 874 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 877 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 875 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 878 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 876 | - __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 879 | + __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 877 | uint32_t count; | 880 | uint32_t count; |
| 878 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 881 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 879 | AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg; | 882 | AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg; |
| @@ -881,35 +884,35 @@ __aicore__ inline void LayerNormGradV3Base::Normalize(const LocalTensor<float>& | |||
| 881 | count = static_cast<uint32_t>(colSize); | 884 | count = static_cast<uint32_t>(colSize); |
| 882 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 885 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 883 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 886 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 884 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 887 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 885 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i); | 888 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i); |
| 886 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride + 0 * innerLoopStride); | 889 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride + 0 * innerLoopStride); |
| 887 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); | 890 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); |
| 888 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); | 891 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); |
| 889 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); | 892 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + 0 * innerLoopStride, cReg, pMask); |
| 890 | } | 893 | } |
| 891 | } | 894 | } |
| 892 | } else { | 895 | } else { |
| 893 | __VEC_SCOPE__ | 896 | __VEC_SCOPE__ |
| 894 | { | 897 | { |
| 895 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 898 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 896 | - __local_mem__ float* src = (__local_mem__ float*)srcTensor.GetPhyAddr(); | 899 | + __ubuf__ float* src = (__ubuf__ float*)srcTensor.GetPhyAddr(); |
| 897 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 900 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 898 | - __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 901 | + __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 899 | uint32_t count; | 902 | uint32_t count; |
| 900 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; | 903 | AscendC::MicroAPI::RegTensor<float> aReg, bReg, cReg; |
| 901 | AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg; | 904 | AscendC::MicroAPI::RegTensor<float> meanReg, rstdReg; |
| 902 | AscendC::MicroAPI::MaskReg pMask; | 905 | AscendC::MicroAPI::MaskReg pMask; |
| 903 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 906 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 904 | count = static_cast<uint32_t>(colSize); | 907 | count = static_cast<uint32_t>(colSize); |
| 905 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 908 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 906 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i); | 909 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i); |
| 907 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 910 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 908 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 911 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 909 | - DataCopy(aReg, (__local_mem__ float*)src + i * outerLoopStride + j * innerLoopStride); | 912 | + LoadAlign(aReg, (__ubuf__ float*)src + i * outerLoopStride + j * innerLoopStride); |
| 910 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); | 913 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(bReg, aReg, meanReg, pMask); |
| 911 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); | 914 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(cReg, bReg, rstdReg, pMask); |
| 912 | - DataCopy((__local_mem__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); | 915 | + StoreAlign((__ubuf__ float*)dst + i * outerLoopStride + j * innerLoopStride, cReg, pMask); |
| 913 | } | 916 | } |
| 914 | } | 917 | } |
| 915 | } | 918 | } |
| @@ -928,56 +931,56 @@ __aicore__ inline void LayerNormGradV3Base::ComputeGammaCommon( | |||
| 928 | if (innerLoopTimes == 1) { | 931 | if (innerLoopTimes == 1) { |
| 929 | __VEC_SCOPE__ | 932 | __VEC_SCOPE__ |
| 930 | { | 933 | { |
| 931 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 934 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 932 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 935 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 933 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 936 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 934 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 937 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 935 | - __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 938 | + __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 936 | uint32_t count = static_cast<uint32_t>(colSize); | 939 | uint32_t count = static_cast<uint32_t>(colSize); |
| 937 | AscendC::MicroAPI::MaskReg pMask; | 940 | AscendC::MicroAPI::MaskReg pMask; |
| 938 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 941 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 939 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 942 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 940 | AscendC::MicroAPI::RegTensor<float> meanReg; | 943 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 941 | AscendC::MicroAPI::RegTensor<float> rstdReg; | 944 | AscendC::MicroAPI::RegTensor<float> rstdReg; |
| 942 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 945 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 943 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i); | 946 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i); |
| 944 | 947 | ||
| 945 | AscendC::MicroAPI::RegTensor<float> xReg; | 948 | AscendC::MicroAPI::RegTensor<float> xReg; |
| 946 | AscendC::MicroAPI::RegTensor<float> dyReg; | 949 | AscendC::MicroAPI::RegTensor<float> dyReg; |
| 947 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + 0 * innerStride); | 950 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + 0 * innerStride); |
| 948 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 951 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 949 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 952 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 950 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + 0 * innerStride); | 953 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + 0 * innerStride); |
| 951 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 954 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 952 | - DataCopy((__local_mem__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); | 955 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + 0 * innerStride, xReg, pMask); |
| 953 | } | 956 | } |
| 954 | } | 957 | } |
| 955 | } else { | 958 | } else { |
| 956 | __VEC_SCOPE__ | 959 | __VEC_SCOPE__ |
| 957 | { | 960 | { |
| 958 | - __local_mem__ float* dst = (__local_mem__ float*)dstTensor.GetPhyAddr(); | 961 | + __ubuf__ float* dst = (__ubuf__ float*)dstTensor.GetPhyAddr(); |
| 959 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 962 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 960 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 963 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 961 | - __local_mem__ float* mean = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 964 | + __ubuf__ float* mean = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 962 | - __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 965 | + __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 963 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 966 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 964 | uint32_t count = static_cast<uint32_t>(colSize); | 967 | uint32_t count = static_cast<uint32_t>(colSize); |
| 965 | AscendC::MicroAPI::RegTensor<float> meanReg; | 968 | AscendC::MicroAPI::RegTensor<float> meanReg; |
| 966 | AscendC::MicroAPI::RegTensor<float> rstdReg; | 969 | AscendC::MicroAPI::RegTensor<float> rstdReg; |
| 967 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__local_mem__ float*)mean + i); | 970 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(meanReg, (__ubuf__ float*)mean + i); |
| 968 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i); | 971 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i); |
| 969 | 972 | ||
| 970 | AscendC::MicroAPI::RegTensor<float> xReg; | 973 | AscendC::MicroAPI::RegTensor<float> xReg; |
| 971 | AscendC::MicroAPI::RegTensor<float> dyReg; | 974 | AscendC::MicroAPI::RegTensor<float> dyReg; |
| 972 | AscendC::MicroAPI::MaskReg pMask; | 975 | AscendC::MicroAPI::MaskReg pMask; |
| 973 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 976 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 974 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 977 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 975 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerStride + j * innerStride); | 978 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerStride + j * innerStride); |
| 976 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); | 979 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, meanReg, pMask); |
| 977 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); | 980 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, rstdReg, pMask); |
| 978 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerStride + j * innerStride); | 981 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerStride + j * innerStride); |
| 979 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); | 982 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(xReg, xReg, dyReg, pMask); |
| 980 | - DataCopy((__local_mem__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); | 983 | + StoreAlign((__ubuf__ float*)dst + i * outerStride + j * innerStride, xReg, pMask); |
| 981 | } | 984 | } |
| 982 | } | 985 | } |
| 983 | } | 986 | } |
| @@ -1184,13 +1187,13 @@ __aicore__ inline void LayerNormGradV3Base::ComputeDxCommon( | |||
| 1184 | if (innerLoopTimes == 1) { | 1187 | if (innerLoopTimes == 1) { |
| 1185 | __VEC_SCOPE__ | 1188 | __VEC_SCOPE__ |
| 1186 | { | 1189 | { |
| 1187 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 1190 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 1188 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 1191 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 1189 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 1192 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 1190 | - __local_mem__ float* gamma = (__local_mem__ float*)gammaTensor.GetPhyAddr(); | 1193 | + __ubuf__ float* gamma = (__ubuf__ float*)gammaTensor.GetPhyAddr(); |
| 1191 | - __local_mem__ float* sum1 = (__local_mem__ float*)sum1Tensor.GetPhyAddr(); | 1194 | + __ubuf__ float* sum1 = (__ubuf__ float*)sum1Tensor.GetPhyAddr(); |
| 1192 | - __local_mem__ float* sum2 = (__local_mem__ float*)sum2Tensor.GetPhyAddr(); | 1195 | + __ubuf__ float* sum2 = (__ubuf__ float*)sum2Tensor.GetPhyAddr(); |
| 1193 | - __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 1196 | + __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 1194 | uint32_t count; | 1197 | uint32_t count; |
| 1195 | 1198 | ||
| 1196 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; | 1199 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; |
| @@ -1201,12 +1204,12 @@ __aicore__ inline void LayerNormGradV3Base::ComputeDxCommon( | |||
| 1201 | count = static_cast<uint32_t>(colSize); | 1204 | count = static_cast<uint32_t>(colSize); |
| 1202 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 1205 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 1203 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 1206 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 1204 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__local_mem__ float*)sum1 + i); | 1207 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__ubuf__ float*)sum1 + i); |
| 1205 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__local_mem__ float*)sum2 + i); | 1208 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__ubuf__ float*)sum2 + i); |
| 1206 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i); | 1209 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i); |
| 1207 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerLoopStride + 0 * innerLoopStride); | 1210 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerLoopStride + 0 * innerLoopStride); |
| 1208 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerLoopStride + 0 * innerLoopStride); | 1211 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerLoopStride + 0 * innerLoopStride); |
| 1209 | - DataCopy(gammaReg, (__local_mem__ float*)gamma + 0 * innerLoopStride); | 1212 | + LoadAlign(gammaReg, (__ubuf__ float*)gamma + 0 * innerLoopStride); |
| 1210 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); | 1213 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); |
| 1211 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); | 1214 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); |
| 1212 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); | 1215 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); |
| @@ -1220,13 +1223,13 @@ __aicore__ inline void LayerNormGradV3Base::ComputeDxCommon( | |||
| 1220 | } else { | 1223 | } else { |
| 1221 | __VEC_SCOPE__ | 1224 | __VEC_SCOPE__ |
| 1222 | { | 1225 | { |
| 1223 | - __local_mem__ T* dst = (__local_mem__ T*)dstTensor.GetPhyAddr(); | 1226 | + __ubuf__ T* dst = (__ubuf__ T*)dstTensor.GetPhyAddr(); |
| 1224 | - __local_mem__ float* dy = (__local_mem__ float*)dyTensor.GetPhyAddr(); | 1227 | + __ubuf__ float* dy = (__ubuf__ float*)dyTensor.GetPhyAddr(); |
| 1225 | - __local_mem__ float* x = (__local_mem__ float*)xTensor.GetPhyAddr(); | 1228 | + __ubuf__ float* x = (__ubuf__ float*)xTensor.GetPhyAddr(); |
| 1226 | - __local_mem__ float* gamma = (__local_mem__ float*)gammaTensor.GetPhyAddr(); | 1229 | + __ubuf__ float* gamma = (__ubuf__ float*)gammaTensor.GetPhyAddr(); |
| 1227 | - __local_mem__ float* sum1 = (__local_mem__ float*)sum1Tensor.GetPhyAddr(); | 1230 | + __ubuf__ float* sum1 = (__ubuf__ float*)sum1Tensor.GetPhyAddr(); |
| 1228 | - __local_mem__ float* sum2 = (__local_mem__ float*)sum2Tensor.GetPhyAddr(); | 1231 | + __ubuf__ float* sum2 = (__ubuf__ float*)sum2Tensor.GetPhyAddr(); |
| 1229 | - __local_mem__ float* rstd = (__local_mem__ float*)rstdTensor.GetPhyAddr(); | 1232 | + __ubuf__ float* rstd = (__ubuf__ float*)rstdTensor.GetPhyAddr(); |
| 1230 | uint32_t count; | 1233 | uint32_t count; |
| 1231 | 1234 | ||
| 1232 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; | 1235 | AscendC::MicroAPI::RegTensor<float> xReg, dyReg, dxReg; |
| @@ -1236,14 +1239,14 @@ __aicore__ inline void LayerNormGradV3Base::ComputeDxCommon( | |||
| 1236 | AscendC::MicroAPI::MaskReg pMask; | 1239 | AscendC::MicroAPI::MaskReg pMask; |
| 1237 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { | 1240 | for (uint16_t i = 0; i < outerLoopTimes; ++i) { |
| 1238 | count = static_cast<uint32_t>(colSize); | 1241 | count = static_cast<uint32_t>(colSize); |
| 1239 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__local_mem__ float*)sum1 + i); | 1242 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum1Reg, (__ubuf__ float*)sum1 + i); |
| 1240 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__local_mem__ float*)sum2 + i); | 1243 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(sum2Reg, (__ubuf__ float*)sum2 + i); |
| 1241 | - DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__local_mem__ float*)rstd + i); | 1244 | + LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_BRC_B32>(rstdReg, (__ubuf__ float*)rstd + i); |
| 1242 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { | 1245 | for (uint16_t j = 0; j < innerLoopTimes; ++j) { |
| 1243 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 1246 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 1244 | - DataCopy(dyReg, (__local_mem__ float*)dy + i * outerLoopStride + j * innerLoopStride); | 1247 | + LoadAlign(dyReg, (__ubuf__ float*)dy + i * outerLoopStride + j * innerLoopStride); |
| 1245 | - DataCopy(xReg, (__local_mem__ float*)x + i * outerLoopStride + j * innerLoopStride); | 1248 | + LoadAlign(xReg, (__ubuf__ float*)x + i * outerLoopStride + j * innerLoopStride); |
| 1246 | - DataCopy(gammaReg, (__local_mem__ float*)gamma + j * innerLoopStride); | 1249 | + LoadAlign(gammaReg, (__ubuf__ float*)gamma + j * innerLoopStride); |
| 1247 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); | 1250 | Mul<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg0, dyReg, gammaReg, pMask); |
| 1248 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); | 1251 | Muls<float, float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg1, Reg0, floatN, pMask); |
| 1249 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); | 1252 | Sub<float, AscendC::MicroAPI::MaskMergeMode::ZEROING>(Reg2, Reg1, sum1Reg, pMask); |
| @@ -1259,16 +1262,16 @@ __aicore__ inline void LayerNormGradV3Base::ComputeDxCommon( | |||
| 1259 | } | 1262 | } |
| 1260 | 1263 | ||
| 1261 | template <typename T> | 1264 | template <typename T> |
| 1262 | -__aicore__ inline void LayerNormGradV3Base::StoreTensorForDtypeT(__local_mem__ T* dst, | 1265 | +__aicore__ inline void LayerNormGradV3Base::StoreTensorForDtypeT(__ubuf__ T* dst, |
| 1263 | AscendC::MicroAPI::RegTensor<float>& src, | 1266 | AscendC::MicroAPI::RegTensor<float>& src, |
| 1264 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset) | 1267 | AscendC::MicroAPI::MaskReg& preg, uint32_t offset) |
| 1265 | { | 1268 | { |
| 1266 | if constexpr (IsSameType<T, float>::value) { | 1269 | if constexpr (IsSameType<T, float>::value) { |
| 1267 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); | 1270 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_NORM>(dst + offset, src, preg); |
| 1268 | } else { | 1271 | } else { |
| 1269 | AscendC::MicroAPI::RegTensor<T> xFp16; | 1272 | AscendC::MicroAPI::RegTensor<T> xFp16; |
| 1270 | Cast<T, float, castTraitB322B16>(xFp16, src, preg); | 1273 | Cast<T, float, castTraitB322B16>(xFp16, src, preg); |
| 1271 | - DataCopy<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); | 1274 | + StoreAlign<T, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(dst + offset, xFp16, preg); |
| 1272 | } | 1275 | } |
| 1273 | } | 1276 | } |
| 1274 | } // namespace LayerNormGradV3 | 1277 | } // namespace LayerNormGradV3 |
| @@ -22,6 +22,8 @@ | |||
| 22 | 22 | ||
| 23 | namespace LayerNormV3 { | 23 | namespace LayerNormV3 { |
| 24 | using namespace AscendC; | 24 | using namespace AscendC; |
| 25 | +using AscendC::Reg::LoadAlign; | ||
| 26 | +using AscendC::Reg::StoreAlign; | ||
| 25 | 27 | ||
| 26 | constexpr static int64_t BLOCK_SIZE = 32; | 28 | constexpr static int64_t BLOCK_SIZE = 32; |
| 27 | constexpr static uint32_t FLOAT_BYTES = 4; | 29 | constexpr static uint32_t FLOAT_BYTES = 4; |
| @@ -79,43 +81,43 @@ __aicore__ inline void WelfordInitialize(const LocalTensor<float>& mean, const L | |||
| 79 | uint16_t loopTimes = (elemCnt + VL_B32 - 1) / VL_B32; | 81 | uint16_t loopTimes = (elemCnt + VL_B32 - 1) / VL_B32; |
| 80 | __VEC_SCOPE__ | 82 | __VEC_SCOPE__ |
| 81 | { | 83 | { |
| 82 | - __local_mem__ float* meanPtr = (__local_mem__ float*)mean.GetPhyAddr(); | 84 | + __ubuf__ float* meanPtr = (__ubuf__ float*)mean.GetPhyAddr(); |
| 83 | - __local_mem__ float* variancePtr = (__local_mem__ float*)variance.GetPhyAddr(); | 85 | + __ubuf__ float* variancePtr = (__ubuf__ float*)variance.GetPhyAddr(); |
| 84 | uint32_t count = static_cast<uint32_t>(elemCnt); | 86 | uint32_t count = static_cast<uint32_t>(elemCnt); |
| 85 | MicroAPI::RegTensor<float> xReg; | 87 | MicroAPI::RegTensor<float> xReg; |
| 86 | MicroAPI::MaskReg pMask; | 88 | MicroAPI::MaskReg pMask; |
| 87 | Duplicate(xReg, 0.0f); | 89 | Duplicate(xReg, 0.0f); |
| 88 | for (uint16_t i = 0; i < loopTimes; ++i) { | 90 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 89 | pMask = MicroAPI::UpdateMask<float>(count); | 91 | pMask = MicroAPI::UpdateMask<float>(count); |
| 90 | - DataCopy(meanPtr + i * VL_B32, xReg, pMask); | 92 | + StoreAlign(meanPtr + i * VL_B32, xReg, pMask); |
| 91 | - DataCopy(variancePtr + i * VL_B32, xReg, pMask); | 93 | + StoreAlign(variancePtr + i * VL_B32, xReg, pMask); |
| 92 | } | 94 | } |
| 93 | } | 95 | } |
| 94 | } | 96 | } |
| 95 | 97 | ||
| 96 | template <typename T_IN> | 98 | template <typename T_IN> |
| 97 | -__aicore__ inline void LoadTensorForDtypeTIn(__local_mem__ T_IN* src, AscendC::Reg::RegTensor<float>& dst, | 99 | +__aicore__ inline void LoadTensorForDtypeTIn(__ubuf__ T_IN* src, AscendC::Reg::RegTensor<float>& dst, |
| 98 | AscendC::Reg::MaskReg& preg, AscendC::MicroAPI::AddrReg& addrReg) | 100 | AscendC::Reg::MaskReg& preg, AscendC::MicroAPI::AddrReg& addrReg) |
| 99 | { | 101 | { |
| 100 | if constexpr (IsSameType<T_IN, float>::value) { | 102 | if constexpr (IsSameType<T_IN, float>::value) { |
| 101 | - DataCopy<float, AscendC::Reg::LoadDist::DIST_NORM>(dst, src, addrReg); | 103 | + LoadAlign<float, AscendC::Reg::LoadDist::DIST_NORM>(dst, src, addrReg); |
| 102 | } else { | 104 | } else { |
| 103 | AscendC::Reg::RegTensor<T_IN> xIn; | 105 | AscendC::Reg::RegTensor<T_IN> xIn; |
| 104 | - DataCopy<T_IN, AscendC::Reg::LoadDist::DIST_UNPACK_B16>(xIn, src, addrReg); | 106 | + LoadAlign<T_IN, AscendC::Reg::LoadDist::DIST_UNPACK_B16>(xIn, src, addrReg); |
| 105 | Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg); | 107 | Cast<float, T_IN, castTraitB162B32>(dst, xIn, preg); |
| 106 | } | 108 | } |
| 107 | } | 109 | } |
| 108 | 110 | ||
| 109 | template <typename T_OUT> | 111 | template <typename T_OUT> |
| 110 | -__aicore__ inline void StoreTensorForDtypeTOut(__local_mem__ T_OUT* dst, AscendC::Reg::RegTensor<float>& src, | 112 | +__aicore__ inline void StoreTensorForDtypeTOut(__ubuf__ T_OUT* dst, AscendC::Reg::RegTensor<float>& src, |
| 111 | AscendC::Reg::MaskReg& preg, AscendC::MicroAPI::AddrReg& addrReg) | 113 | AscendC::Reg::MaskReg& preg, AscendC::MicroAPI::AddrReg& addrReg) |
| 112 | { | 114 | { |
| 113 | if constexpr (IsSameType<T_OUT, float>::value) { | 115 | if constexpr (IsSameType<T_OUT, float>::value) { |
| 114 | - DataCopy<T_OUT, AscendC::Reg::StoreDist::DIST_NORM>(dst, src, addrReg, preg); | 116 | + StoreAlign<T_OUT, AscendC::Reg::StoreDist::DIST_NORM>(dst, src, addrReg, preg); |
| 115 | } else { | 117 | } else { |
| 116 | AscendC::Reg::RegTensor<T_OUT> xOut; | 118 | AscendC::Reg::RegTensor<T_OUT> xOut; |
| 117 | Cast<T_OUT, float, castTraitB322B16>(xOut, src, preg); | 119 | Cast<T_OUT, float, castTraitB322B16>(xOut, src, preg); |
| 118 | - DataCopy<T_OUT, AscendC::Reg::StoreDist::DIST_PACK_B32>(dst, xOut, addrReg, preg); | 120 | + StoreAlign<T_OUT, AscendC::Reg::StoreDist::DIST_PACK_B32>(dst, xOut, addrReg, preg); |
| 119 | } | 121 | } |
| 120 | } | 122 | } |
| 121 | 123 | ||
| @@ -148,10 +150,10 @@ __aicore__ inline void CastBatchMeanLastout(LocalTensor<float>& meanTensor, Loca | |||
| 148 | constexpr static uint32_t VL_F32 = VECTOR_REG_WIDTH / sizeof(float); | 150 | constexpr static uint32_t VL_F32 = VECTOR_REG_WIDTH / sizeof(float); |
| 149 | constexpr static uint32_t VL_MEAN = VECTOR_REG_WIDTH / sizeof(M); | 151 | constexpr static uint32_t VL_MEAN = VECTOR_REG_WIDTH / sizeof(M); |
| 150 | 152 | ||
| 151 | - __local_mem__ float* batchMeanInAddr = (__local_mem__ float*)meanTensor.GetPhyAddr(); | 153 | + __ubuf__ float* batchMeanInAddr = (__ubuf__ float*)meanTensor.GetPhyAddr(); |
| 152 | - __local_mem__ float* batchLastoutInAddr = (__local_mem__ float*)lastoutTensor.GetPhyAddr(); | 154 | + __ubuf__ float* batchLastoutInAddr = (__ubuf__ float*)lastoutTensor.GetPhyAddr(); |
| 153 | - __local_mem__ M* batchMeanOutAddr = (__local_mem__ M*)meanTensor.GetPhyAddr(); | 155 | + __ubuf__ M* batchMeanOutAddr = (__ubuf__ M*)meanTensor.GetPhyAddr(); |
| 154 | - __local_mem__ M* batchLastoutOutAddr = (__local_mem__ M*)lastoutTensor.GetPhyAddr(); | 156 | + __ubuf__ M* batchLastoutOutAddr = (__ubuf__ M*)lastoutTensor.GetPhyAddr(); |
| 155 | 157 | ||
| 156 | uint32_t castCount = static_cast<uint32_t>(currentANum); | 158 | uint32_t castCount = static_cast<uint32_t>(currentANum); |
| 157 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); | 159 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); |
| @@ -164,14 +166,14 @@ __aicore__ inline void CastBatchMeanLastout(LocalTensor<float>& meanTensor, Loca | |||
| 164 | MicroAPI::MaskReg pregLoop; | 166 | MicroAPI::MaskReg pregLoop; |
| 165 | for (uint16_t i = 0; i < castLoops; i++) { | 167 | for (uint16_t i = 0; i < castLoops; i++) { |
| 166 | pregLoop = MicroAPI::UpdateMask<float>(castCount); | 168 | pregLoop = MicroAPI::UpdateMask<float>(castCount); |
| 167 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); | 169 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); |
| 168 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_lastout, batchLastoutInAddr + VL_F32 * i); | 170 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_lastout, batchLastoutInAddr + VL_F32 * i); |
| 169 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 171 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 170 | Cast<M, float, castTraitB322B16>(output_lastout, input_lastout, pregLoop); | 172 | Cast<M, float, castTraitB322B16>(output_lastout, input_lastout, pregLoop); |
| 171 | - MicroAPI::DataCopy<M, MicroAPI::StoreDist::DIST_PACK_B32>( | 173 | + MicroAPI::StoreAlign<M, MicroAPI::StoreDist::DIST_PACK_B32>(((__ubuf__ M*)batchMeanOutAddr + i * VL_MEAN), |
| 172 | - ((__local_mem__ M*)batchMeanOutAddr + i * VL_MEAN), output_mean, pregLoop); | 174 | + output_mean, pregLoop); |
| 173 | - MicroAPI::DataCopy<M, MicroAPI::StoreDist::DIST_PACK_B32>( | 175 | + MicroAPI::StoreAlign<M, MicroAPI::StoreDist::DIST_PACK_B32>( |
| 174 | - ((__local_mem__ M*)batchLastoutOutAddr + i * VL_MEAN), output_lastout, pregLoop); | 176 | + ((__ubuf__ M*)batchLastoutOutAddr + i * VL_MEAN), output_lastout, pregLoop); |
| 175 | } | 177 | } |
| 176 | } | 178 | } |
| 177 | } | 179 | } |
| @@ -30,6 +30,7 @@ using AscendC::MicroAPI::MemType; | |||
| 30 | using AscendC::MicroAPI::RegTensor; | 30 | using AscendC::MicroAPI::RegTensor; |
| 31 | using AscendC::MicroAPI::StoreDist; | 31 | using AscendC::MicroAPI::StoreDist; |
| 32 | using AscendC::MicroAPI::UpdateMask; | 32 | using AscendC::MicroAPI::UpdateMask; |
| 33 | +using AscendC::Reg::LoadAlign; | ||
| 33 | using NormCommon::NormCommonRegbase::LoadRegForDtype; | 34 | using NormCommon::NormCommonRegbase::LoadRegForDtype; |
| 34 | using NormCommon::NormCommonRegbase::StoreRegForDtype; | 35 | using NormCommon::NormCommonRegbase::StoreRegForDtype; |
| 35 | 36 | ||
| @@ -128,17 +129,17 @@ private: | |||
| 128 | rstdOutUb_ = rstdQueue_.AllocTensor<M>(); | 129 | rstdOutUb_ = rstdQueue_.AllocTensor<M>(); |
| 129 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); | 130 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); |
| 130 | 131 | ||
| 131 | - __local_mem__ T* xInUbAddr = (__local_mem__ T*)xInUb.GetPhyAddr(); | 132 | + __ubuf__ T* xInUbAddr = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 132 | - __local_mem__ M* meanOutUbAddr = (__local_mem__ M*)meanOutUb_.GetPhyAddr(); | 133 | + __ubuf__ M* meanOutUbAddr = (__ubuf__ M*)meanOutUb_.GetPhyAddr(); |
| 133 | - __local_mem__ M* rstdOutUbAddr = (__local_mem__ M*)rstdOutUb_.GetPhyAddr(); | 134 | + __ubuf__ M* rstdOutUbAddr = (__ubuf__ M*)rstdOutUb_.GetPhyAddr(); |
| 134 | - __local_mem__ float* tmpUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr(); | 135 | + __ubuf__ float* tmpUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr(); |
| 135 | CalculateMeanVar(xInUbAddr, meanOutUbAddr, tmpUbAddr, currentANum); | 136 | CalculateMeanVar(xInUbAddr, meanOutUbAddr, tmpUbAddr, currentANum); |
| 136 | CalculateRstd(rstdOutUbAddr, tmpUbAddr, currentANum); | 137 | CalculateRstd(rstdOutUbAddr, tmpUbAddr, currentANum); |
| 137 | 138 | ||
| 138 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); | 139 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); |
| 139 | - __local_mem__ U* gammaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr(); | 140 | + __ubuf__ U* gammaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr(); |
| 140 | - __local_mem__ U* betaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr() + BLOCK_SIZE / sizeof(U); | 141 | + __ubuf__ U* betaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr() + BLOCK_SIZE / sizeof(U); |
| 141 | - __local_mem__ T* yOutUbAddr = (__local_mem__ T*)yOutUb.GetPhyAddr(); | 142 | + __ubuf__ T* yOutUbAddr = (__ubuf__ T*)yOutUb.GetPhyAddr(); |
| 142 | 143 | ||
| 143 | if (hasGamma_ && hasBeta_) { | 144 | if (hasGamma_ && hasBeta_) { |
| 144 | CalculateY<true, true>(xInUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, tmpUbAddr, currentANum); | 145 | CalculateY<true, true>(xInUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, tmpUbAddr, currentANum); |
| @@ -153,8 +154,8 @@ private: | |||
| 153 | yQueue_.EnQue(yOutUb); | 154 | yQueue_.EnQue(yOutUb); |
| 154 | } | 155 | } |
| 155 | 156 | ||
| 156 | - __aicore__ inline void CalculateMeanVar(__local_mem__ T* xInUb, __local_mem__ M* meanInUb, | 157 | + __aicore__ inline void CalculateMeanVar(__ubuf__ T* xInUb, __ubuf__ M* meanInUb, __ubuf__ float* tmpUb, |
| 157 | - __local_mem__ float* tmpUb, uint64_t currentANum) | 158 | + uint64_t currentANum) |
| 158 | { | 159 | { |
| 159 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); | 160 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); |
| 160 | uint32_t sreg = static_cast<uint32_t>(currentANum); | 161 | uint32_t sreg = static_cast<uint32_t>(currentANum); |
| @@ -190,7 +191,7 @@ private: | |||
| 190 | meanQueue_.FreeTensor(meanInUb); | 191 | meanQueue_.FreeTensor(meanInUb); |
| 191 | rstdQueue_.FreeTensor(rstdInUb); | 192 | rstdQueue_.FreeTensor(rstdInUb); |
| 192 | } | 193 | } |
| 193 | - __aicore__ inline void CalculateRstd(__local_mem__ M* rstdOutUb, __local_mem__ float* tmpUb, int64_t currentANum) | 194 | + __aicore__ inline void CalculateRstd(__ubuf__ M* rstdOutUb, __ubuf__ float* tmpUb, int64_t currentANum) |
| 194 | { | 195 | { |
| 195 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); | 196 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); |
| 196 | uint32_t sreg = static_cast<uint32_t>(currentANum); | 197 | uint32_t sreg = static_cast<uint32_t>(currentANum); |
| @@ -216,8 +217,8 @@ private: | |||
| 216 | } | 217 | } |
| 217 | } | 218 | } |
| 218 | template <bool hasGammaFlag, bool hasBetaFlag> | 219 | template <bool hasGammaFlag, bool hasBetaFlag> |
| 219 | - __aicore__ inline void CalculateY(__local_mem__ T* xInUb, __local_mem__ U* betaInUb, __local_mem__ U* gammaInUb, | 220 | + __aicore__ inline void CalculateY(__ubuf__ T* xInUb, __ubuf__ U* betaInUb, __ubuf__ U* gammaInUb, |
| 220 | - __local_mem__ T* yOutUb, __local_mem__ float* tmpUb, int64_t currentANum) | 221 | + __ubuf__ T* yOutUb, __ubuf__ float* tmpUb, int64_t currentANum) |
| 221 | { | 222 | { |
| 222 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); | 223 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); |
| 223 | uint32_t sreg = static_cast<uint32_t>(currentANum); | 224 | uint32_t sreg = static_cast<uint32_t>(currentANum); |
| @@ -243,7 +244,7 @@ private: | |||
| 243 | LoadRegForDtype(tmpUb, rstdReg, pregLoop, (a * VL_B32)); | 244 | LoadRegForDtype(tmpUb, rstdReg, pregLoop, (a * VL_B32)); |
| 244 | Mul(yReg, subReg, rstdReg, pregLoop); | 245 | Mul(yReg, subReg, rstdReg, pregLoop); |
| 245 | if constexpr (hasGammaFlag && hasBetaFlag) { | 246 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 246 | - FusedMulDstAdd(yReg, gammaReg, betaReg, pregLoop); | 247 | + MulDstAdd(yReg, gammaReg, betaReg, pregLoop); |
| 247 | } else { | 248 | } else { |
| 248 | if constexpr (hasGammaFlag) { | 249 | if constexpr (hasGammaFlag) { |
| 249 | Mul(yReg, yReg, gammaReg, pregLoop); | 250 | Mul(yReg, yReg, gammaReg, pregLoop); |
| @@ -269,14 +270,14 @@ private: | |||
| 269 | } | 270 | } |
| 270 | 271 | ||
| 271 | template <typename H> | 272 | template <typename H> |
| 272 | - __aicore__ inline void LoadsTensorForDtypeT(const __local_mem__ void* src, MicroAPI::RegTensor<float>& dst, | 273 | + __aicore__ inline void LoadsTensorForDtypeT(const __ubuf__ void* src, MicroAPI::RegTensor<float>& dst, |
| 273 | MicroAPI::MaskReg& preg, uint32_t offset) | 274 | MicroAPI::MaskReg& preg, uint32_t offset) |
| 274 | { | 275 | { |
| 275 | if constexpr (IsSameType<H, float>::value) { | 276 | if constexpr (IsSameType<H, float>::value) { |
| 276 | - DataCopy<float, LoadDist::DIST_BRC_B32>(dst, (__local_mem__ float*)src + offset); | 277 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(dst, (__ubuf__ float*)src + offset); |
| 277 | } else { // fp16、bf16 | 278 | } else { // fp16、bf16 |
| 278 | RegTensor<H> xFp16; | 279 | RegTensor<H> xFp16; |
| 279 | - DataCopy<H, LoadDist::DIST_BRC_B16>(xFp16, ((__local_mem__ H*)src + offset)); | 280 | + LoadAlign<H, LoadDist::DIST_BRC_B16>(xFp16, ((__ubuf__ H*)src + offset)); |
| 280 | Cast<float, H, castTraitB162B32>(dst, xFp16, preg); | 281 | Cast<float, H, castTraitB162B32>(dst, xFp16, preg); |
| 281 | } | 282 | } |
| 282 | } | 283 | } |
| @@ -313,4 +314,4 @@ private: | |||
| 313 | }; | 314 | }; |
| 314 | } // namespace LayerNormV3 | 315 | } // namespace LayerNormV3 |
| 315 | 316 | ||
| 316 | -#endif // LAYER_NORM_V3_NO_REDUCE_H | 317 | +#endif // LAYER_NORM_V3_NO_REDUCE_H |
| @@ -30,6 +30,9 @@ using AscendC::MicroAPI::MemType; | |||
| 30 | using AscendC::MicroAPI::RegTensor; | 30 | using AscendC::MicroAPI::RegTensor; |
| 31 | using AscendC::MicroAPI::StoreDist; | 31 | using AscendC::MicroAPI::StoreDist; |
| 32 | using AscendC::MicroAPI::UpdateMask; | 32 | using AscendC::MicroAPI::UpdateMask; |
| 33 | +using AscendC::Reg::LoadAlign; | ||
| 34 | +using AscendC::Reg::Reduce; | ||
| 35 | +using AscendC::Reg::StoreAlign; | ||
| 33 | using NormCommon::NormCommonRegbase::LoadRegForDtype; | 36 | using NormCommon::NormCommonRegbase::LoadRegForDtype; |
| 34 | using NormCommon::NormCommonRegbase::StoreRegForDtype; | 37 | using NormCommon::NormCommonRegbase::StoreRegForDtype; |
| 35 | 38 | ||
| @@ -161,9 +164,9 @@ private: | |||
| 161 | int64_t r2 = tl_->r / tl_->r1; | 164 | int64_t r2 = tl_->r / tl_->r1; |
| 162 | int64_t r1 = tl_->r1; | 165 | int64_t r1 = tl_->r1; |
| 163 | 166 | ||
| 164 | - static constexpr AscendC::MultiCopyConfig copyConfig = {false, 0, 0, false}; | 167 | + static constexpr AscendC::NdDmaConfig copyConfig = {false, 0, 0, false}; |
| 165 | constexpr int64_t MULTI_COPY_DIM = 2; | 168 | constexpr int64_t MULTI_COPY_DIM = 2; |
| 166 | - MultiCopyLoopInfo<MULTI_COPY_DIM> multiCopyParams; | 169 | + NdDmaLoopInfo<MULTI_COPY_DIM> multiCopyParams; |
| 167 | multiCopyParams.loopSrcStride[0] = 1; | 170 | multiCopyParams.loopSrcStride[0] = 1; |
| 168 | multiCopyParams.loopSrcStride[1] = 0; | 171 | multiCopyParams.loopSrcStride[1] = 0; |
| 169 | 172 | ||
| @@ -174,7 +177,7 @@ private: | |||
| 174 | multiCopyParams.loopSize[1] = r1; | 177 | multiCopyParams.loopSize[1] = r1; |
| 175 | 178 | ||
| 176 | U constValue = 0; | 179 | U constValue = 0; |
| 177 | - AscendC::MultiCopyParams<U, MULTI_COPY_DIM> copyParams = {multiCopyParams, constValue}; | 180 | + AscendC::NdDmaParams<U, MULTI_COPY_DIM> copyParams = {multiCopyParams, constValue}; |
| 178 | 181 | ||
| 179 | if (hasGamma_) { | 182 | if (hasGamma_) { |
| 180 | AscendC::DataCopy<U, MULTI_COPY_DIM, copyConfig>(gammaBetaInUb_, gammaGm_, copyParams); | 183 | AscendC::DataCopy<U, MULTI_COPY_DIM, copyConfig>(gammaBetaInUb_, gammaGm_, copyParams); |
| @@ -336,11 +339,11 @@ private: | |||
| 336 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); | 339 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); |
| 337 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); | 340 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); |
| 338 | 341 | ||
| 339 | - __local_mem__ T* xInUbAddr = (__local_mem__ T*)xInUb.GetPhyAddr(); | 342 | + __ubuf__ T* xInUbAddr = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 340 | - __local_mem__ float* meanOutUbAddr = (__local_mem__ float*)meanOutUb_.GetPhyAddr(); | 343 | + __ubuf__ float* meanOutUbAddr = (__ubuf__ float*)meanOutUb_.GetPhyAddr(); |
| 341 | - __local_mem__ float* rstdOutUbAddr = (__local_mem__ float*)rstdOutUb_.GetPhyAddr(); | 344 | + __ubuf__ float* rstdOutUbAddr = (__ubuf__ float*)rstdOutUb_.GetPhyAddr(); |
| 342 | - __local_mem__ float* xSubMeanUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr(); | 345 | + __ubuf__ float* xSubMeanUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr(); |
| 343 | - __local_mem__ float* tmpUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr() + elemNum_; | 346 | + __ubuf__ float* tmpUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr() + elemNum_; |
| 344 | 347 | ||
| 345 | if (tl_->rAlign <= VL_B32) { | 348 | if (tl_->rAlign <= VL_B32) { |
| 346 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); | 349 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); |
| @@ -354,9 +357,9 @@ private: | |||
| 354 | } | 357 | } |
| 355 | 358 | ||
| 356 | LocalTensor<float> rstdTmpTensor = rstdTmpBuf_.Get<float>(); | 359 | LocalTensor<float> rstdTmpTensor = rstdTmpBuf_.Get<float>(); |
| 357 | - __local_mem__ float* rstdTmpUbAddr = (__local_mem__ float*)rstdTmpTensor.GetPhyAddr(); | 360 | + __ubuf__ float* rstdTmpUbAddr = (__ubuf__ float*)rstdTmpTensor.GetPhyAddr(); |
| 358 | CalculateRstdVF(rstdOutUbAddr, rstdTmpUbAddr, currentANum); | 361 | CalculateRstdVF(rstdOutUbAddr, rstdTmpUbAddr, currentANum); |
| 359 | - __local_mem__ float* rstdForNorm; | 362 | + __ubuf__ float* rstdForNorm; |
| 360 | if constexpr (IsOutRstd) { | 363 | if constexpr (IsOutRstd) { |
| 361 | rstdForNorm = rstdOutUbAddr; | 364 | rstdForNorm = rstdOutUbAddr; |
| 362 | } else { | 365 | } else { |
| @@ -364,9 +367,9 @@ private: | |||
| 364 | } | 367 | } |
| 365 | 368 | ||
| 366 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); | 369 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); |
| 367 | - __local_mem__ U* gammaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr(); | 370 | + __ubuf__ U* gammaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr(); |
| 368 | - __local_mem__ U* betaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAlign; | 371 | + __ubuf__ U* betaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAlign; |
| 369 | - __local_mem__ T* yOutUbAddr = (__local_mem__ T*)yOutUb.GetPhyAddr(); | 372 | + __ubuf__ T* yOutUbAddr = (__ubuf__ T*)yOutUb.GetPhyAddr(); |
| 370 | if (hasGamma_ && hasBeta_) { | 373 | if (hasGamma_ && hasBeta_) { |
| 371 | CalculateNormalizeVF<true, true>(xSubMeanUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, rstdForNorm, | 374 | CalculateNormalizeVF<true, true>(xSubMeanUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, rstdForNorm, |
| 372 | currentANum); | 375 | currentANum); |
| @@ -392,11 +395,11 @@ private: | |||
| 392 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); | 395 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); |
| 393 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); | 396 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); |
| 394 | 397 | ||
| 395 | - __local_mem__ T* xInUbAddr = (__local_mem__ T*)xInUb.GetPhyAddr(); | 398 | + __ubuf__ T* xInUbAddr = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 396 | - __local_mem__ float* meanOutUbAddr = (__local_mem__ float*)meanOutUb_.GetPhyAddr(); | 399 | + __ubuf__ float* meanOutUbAddr = (__ubuf__ float*)meanOutUb_.GetPhyAddr(); |
| 397 | - __local_mem__ float* rstdOutUbAddr = (__local_mem__ float*)rstdOutUb_.GetPhyAddr(); | 400 | + __ubuf__ float* rstdOutUbAddr = (__ubuf__ float*)rstdOutUb_.GetPhyAddr(); |
| 398 | - __local_mem__ float* xSubMeanUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr(); | 401 | + __ubuf__ float* xSubMeanUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr(); |
| 399 | - __local_mem__ float* tmpUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr() + elemNum_; | 402 | + __ubuf__ float* tmpUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr() + elemNum_; |
| 400 | 403 | ||
| 401 | if (tl_->rAlign <= VL_B32) { | 404 | if (tl_->rAlign <= VL_B32) { |
| 402 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); | 405 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); |
| @@ -410,9 +413,9 @@ private: | |||
| 410 | } | 413 | } |
| 411 | 414 | ||
| 412 | LocalTensor<float> rstdTmpTensor = rstdTmpBuf_.Get<float>(); | 415 | LocalTensor<float> rstdTmpTensor = rstdTmpBuf_.Get<float>(); |
| 413 | - __local_mem__ float* rstdTmpUbAddr = (__local_mem__ float*)rstdTmpTensor.GetPhyAddr(); | 416 | + __ubuf__ float* rstdTmpUbAddr = (__ubuf__ float*)rstdTmpTensor.GetPhyAddr(); |
| 414 | CalculateRstdVF(rstdOutUbAddr, rstdTmpUbAddr, currentANum); | 417 | CalculateRstdVF(rstdOutUbAddr, rstdTmpUbAddr, currentANum); |
| 415 | - __local_mem__ float* rstdForNorm; | 418 | + __ubuf__ float* rstdForNorm; |
| 416 | if constexpr (IsOutRstd) { | 419 | if constexpr (IsOutRstd) { |
| 417 | rstdForNorm = rstdOutUbAddr; | 420 | rstdForNorm = rstdOutUbAddr; |
| 418 | } else { | 421 | } else { |
| @@ -420,9 +423,9 @@ private: | |||
| 420 | } | 423 | } |
| 421 | 424 | ||
| 422 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); | 425 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); |
| 423 | - __local_mem__ T* yOutUbAddr = (__local_mem__ T*)yOutUb.GetPhyAddr(); | 426 | + __ubuf__ T* yOutUbAddr = (__ubuf__ T*)yOutUb.GetPhyAddr(); |
| 424 | - __local_mem__ U* gammaBaseAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr(); | 427 | + __ubuf__ U* gammaBaseAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr(); |
| 425 | - __local_mem__ U* betaBaseAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->b * tl_->rAlign; | 428 | + __ubuf__ U* betaBaseAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->b * tl_->rAlign; |
| 426 | if (hasGamma_ && hasBeta_) { | 429 | if (hasGamma_ && hasBeta_) { |
| 427 | CalculateNormalizeVFFullB<true, true>(xSubMeanUbAddr, betaBaseAddr, gammaBaseAddr, yOutUbAddr, rstdForNorm, | 430 | CalculateNormalizeVFFullB<true, true>(xSubMeanUbAddr, betaBaseAddr, gammaBaseAddr, yOutUbAddr, rstdForNorm, |
| 428 | currentANum, aOffset); | 431 | currentANum, aOffset); |
| @@ -448,11 +451,11 @@ private: | |||
| 448 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); | 451 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); |
| 449 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); | 452 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); |
| 450 | 453 | ||
| 451 | - __local_mem__ T* xInUbAddr = (__local_mem__ T*)xInUb.GetPhyAddr(); | 454 | + __ubuf__ T* xInUbAddr = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 452 | - __local_mem__ float* meanOutUbAddr = (__local_mem__ float*)meanOutUb_.GetPhyAddr(); | 455 | + __ubuf__ float* meanOutUbAddr = (__ubuf__ float*)meanOutUb_.GetPhyAddr(); |
| 453 | - __local_mem__ float* rstdOutUbAddr = (__local_mem__ float*)rstdOutUb_.GetPhyAddr(); | 456 | + __ubuf__ float* rstdOutUbAddr = (__ubuf__ float*)rstdOutUb_.GetPhyAddr(); |
| 454 | - __local_mem__ float* xSubMeanUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr(); | 457 | + __ubuf__ float* xSubMeanUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr(); |
| 455 | - __local_mem__ float* tmpUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr() + elemNum_; | 458 | + __ubuf__ float* tmpUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr() + elemNum_; |
| 456 | 459 | ||
| 457 | if (tl_->rAlign <= VL_B32) { | 460 | if (tl_->rAlign <= VL_B32) { |
| 458 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); | 461 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); |
| @@ -466,9 +469,9 @@ private: | |||
| 466 | } | 469 | } |
| 467 | 470 | ||
| 468 | LocalTensor<float> rstdTmpTensor = rstdTmpBuf_.Get<float>(); | 471 | LocalTensor<float> rstdTmpTensor = rstdTmpBuf_.Get<float>(); |
| 469 | - __local_mem__ float* rstdTmpUbAddr = (__local_mem__ float*)rstdTmpTensor.GetPhyAddr(); | 472 | + __ubuf__ float* rstdTmpUbAddr = (__ubuf__ float*)rstdTmpTensor.GetPhyAddr(); |
| 470 | CalculateRstdVF(rstdOutUbAddr, rstdTmpUbAddr, currentANum); | 473 | CalculateRstdVF(rstdOutUbAddr, rstdTmpUbAddr, currentANum); |
| 471 | - __local_mem__ float* rstdForNorm; | 474 | + __ubuf__ float* rstdForNorm; |
| 472 | if constexpr (IsOutRstd) { | 475 | if constexpr (IsOutRstd) { |
| 473 | rstdForNorm = rstdOutUbAddr; | 476 | rstdForNorm = rstdOutUbAddr; |
| 474 | } else { | 477 | } else { |
| @@ -476,9 +479,9 @@ private: | |||
| 476 | } | 479 | } |
| 477 | 480 | ||
| 478 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); | 481 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); |
| 479 | - __local_mem__ T* yOutUbAddr = (__local_mem__ T*)yOutUb.GetPhyAddr(); | 482 | + __ubuf__ T* yOutUbAddr = (__ubuf__ T*)yOutUb.GetPhyAddr(); |
| 480 | - __local_mem__ U* gammaBaseAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr(); | 483 | + __ubuf__ U* gammaBaseAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr(); |
| 481 | - __local_mem__ U* betaBaseAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAxisCount * tl_->rAlign; | 484 | + __ubuf__ U* betaBaseAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAxisCount * tl_->rAlign; |
| 482 | if (hasGamma_ && hasBeta_) { | 485 | if (hasGamma_ && hasBeta_) { |
| 483 | CalculateNormalizeVFNotFullB<true, true>(xSubMeanUbAddr, betaBaseAddr, gammaBaseAddr, yOutUbAddr, | 486 | CalculateNormalizeVFNotFullB<true, true>(xSubMeanUbAddr, betaBaseAddr, gammaBaseAddr, yOutUbAddr, |
| 484 | rstdForNorm, currentANum); | 487 | rstdForNorm, currentANum); |
| @@ -496,8 +499,8 @@ private: | |||
| 496 | yQueue_.EnQue(yOutUb); | 499 | yQueue_.EnQue(yOutUb); |
| 497 | } | 500 | } |
| 498 | 501 | ||
| 499 | - __aicore__ inline void CalculateMeanVarRLessThanVL(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 502 | + __aicore__ inline void CalculateMeanVarRLessThanVL(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 500 | - __local_mem__ float* rstdInUb, __local_mem__ float* xSubMeanUb, | 503 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 501 | uint16_t currentANum) | 504 | uint16_t currentANum) |
| 502 | { | 505 | { |
| 503 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 506 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| @@ -522,24 +525,24 @@ private: | |||
| 522 | for (uint16_t a = 0; a < currentANum; a++) { | 525 | for (uint16_t a = 0; a < currentANum; a++) { |
| 523 | LoadRegForDtype(xInUb, x, pregLoop, (a * aStride)); | 526 | LoadRegForDtype(xInUb, x, pregLoop, (a * aStride)); |
| 524 | Muls(meanSum, x, n, pregLoop); | 527 | Muls(meanSum, x, n, pregLoop); |
| 525 | - ReduceSum(mean, meanSum, pregLoop); | 528 | + Reduce<ReduceType::SUM>(mean, meanSum, pregLoop); |
| 526 | Muls(mean, mean, nCorrectionFactor, pregOne); | 529 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 527 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 530 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 528 | 531 | ||
| 529 | Duplicate(meanDup, mean, pregFull); | 532 | Duplicate(meanDup, mean, pregFull); |
| 530 | Sub(xMeanSub, x, meanDup, pregLoop); | 533 | Sub(xMeanSub, x, meanDup, pregLoop); |
| 531 | StoreRegForDtype(xSubMeanUb, xMeanSub, pregLoop, (a * aStride)); | 534 | StoreRegForDtype(xSubMeanUb, xMeanSub, pregLoop, (a * aStride)); |
| 532 | Mul(square, xMeanSub, xMeanSub, pregLoop); | 535 | Mul(square, xMeanSub, xMeanSub, pregLoop); |
| 533 | Muls(varSum, square, n, pregLoop); | 536 | Muls(varSum, square, n, pregLoop); |
| 534 | - ReduceSum(var, varSum, pregLoop); | 537 | + Reduce<ReduceType::SUM>(var, varSum, pregLoop); |
| 535 | Muls(var, var, nCorrectionFactor, pregOne); | 538 | Muls(var, var, nCorrectionFactor, pregOne); |
| 536 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 539 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 537 | } | 540 | } |
| 538 | } | 541 | } |
| 539 | } | 542 | } |
| 540 | - __aicore__ inline void CalculateMeanVarRLessThanTwoVL(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 543 | + __aicore__ inline void CalculateMeanVarRLessThanTwoVL(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 541 | - __local_mem__ float* rstdInUb, | 544 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 542 | - __local_mem__ float* xSubMeanUb, uint16_t currentANum) | 545 | + uint16_t currentANum) |
| 543 | { | 546 | { |
| 544 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 547 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| 545 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); | 548 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); |
| @@ -575,9 +578,9 @@ private: | |||
| 575 | Muls(meanSum1, x1, n, pregFull); | 578 | Muls(meanSum1, x1, n, pregFull); |
| 576 | Muls(meanSum2, x2, n, pregTail); | 579 | Muls(meanSum2, x2, n, pregTail); |
| 577 | Add(meanSum, meanSum1, meanSum2, pregFull); | 580 | Add(meanSum, meanSum1, meanSum2, pregFull); |
| 578 | - ReduceSum(mean, meanSum, pregFull); | 581 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 579 | Muls(mean, mean, nCorrectionFactor, pregOne); | 582 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 580 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 583 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 581 | 584 | ||
| 582 | Duplicate(meanDup, mean, pregFull); | 585 | Duplicate(meanDup, mean, pregFull); |
| 583 | Sub(xMeanSub1, x1, meanDup, pregFull); | 586 | Sub(xMeanSub1, x1, meanDup, pregFull); |
| @@ -589,16 +592,16 @@ private: | |||
| 589 | Muls(varSum1, square1, n, pregFull); | 592 | Muls(varSum1, square1, n, pregFull); |
| 590 | Muls(varSum2, square2, n, pregTail); | 593 | Muls(varSum2, square2, n, pregTail); |
| 591 | Add(varSum, varSum1, varSum2, pregFull); | 594 | Add(varSum, varSum1, varSum2, pregFull); |
| 592 | - ReduceSum(var, varSum, pregFull); | 595 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 593 | Muls(var, var, nCorrectionFactor, pregOne); | 596 | Muls(var, var, nCorrectionFactor, pregOne); |
| 594 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 597 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 595 | } | 598 | } |
| 596 | } | 599 | } |
| 597 | } | 600 | } |
| 598 | template <int32_t LAST_LOOP_NUMS> | 601 | template <int32_t LAST_LOOP_NUMS> |
| 599 | - __aicore__ inline void CalculateMeanVarRCommon(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 602 | + __aicore__ inline void CalculateMeanVarRCommon(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 600 | - __local_mem__ float* rstdInUb, __local_mem__ float* xSubMeanUb, | 603 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 601 | - __local_mem__ float* tmpUb, uint16_t currentANum) | 604 | + __ubuf__ float* tmpUb, uint16_t currentANum) |
| 602 | { | 605 | { |
| 603 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 606 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| 604 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); | 607 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); |
| @@ -637,8 +640,8 @@ private: | |||
| 637 | Muls(x1, x1, n, pregFull); | 640 | Muls(x1, x1, n, pregFull); |
| 638 | Muls(x2, x2, n, pregFull); | 641 | Muls(x2, x2, n, pregFull); |
| 639 | Add(meanSum, x1, x2, pregFull); | 642 | Add(meanSum, x1, x2, pregFull); |
| 640 | - ReduceSum(mean, meanSum, pregFull); | 643 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 641 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 644 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 642 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), mean, pregOne); | 645 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), mean, pregOne); |
| 643 | } | 646 | } |
| 644 | for (uint16_t r = 0; | 647 | for (uint16_t r = 0; |
| @@ -651,8 +654,8 @@ private: | |||
| 651 | Muls(x1, x1, n, pregFull); | 654 | Muls(x1, x1, n, pregFull); |
| 652 | Muls(x2, x2, n, pregLoop); | 655 | Muls(x2, x2, n, pregLoop); |
| 653 | Add(meanSum, x1, x2, pregFull); | 656 | Add(meanSum, x1, x2, pregFull); |
| 654 | - ReduceSum(mean, meanSum, pregFull); | 657 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 655 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 658 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 656 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, | 659 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, |
| 657 | pregOne); | 660 | pregOne); |
| 658 | } | 661 | } |
| @@ -661,8 +664,8 @@ private: | |||
| 661 | LoadRegForDtype(xInUb + binaryAddRemainderCeilLoop * VL_B32, x1, pregFull, | 664 | LoadRegForDtype(xInUb + binaryAddRemainderCeilLoop * VL_B32, x1, pregFull, |
| 662 | (r * VL_B32 + a * aStride)); | 665 | (r * VL_B32 + a * aStride)); |
| 663 | Muls(x1, x1, n, pregFull); | 666 | Muls(x1, x1, n, pregFull); |
| 664 | - ReduceSum(mean, x1, pregFull); | 667 | + Reduce<ReduceType::SUM>(mean, x1, pregFull); |
| 665 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 668 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 666 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, | 669 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, |
| 667 | pregOne); | 670 | pregOne); |
| 668 | } | 671 | } |
| @@ -671,24 +674,24 @@ private: | |||
| 671 | if constexpr (LAST_LOOP_NUMS == 1) { | 674 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 672 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); | 675 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); |
| 673 | for (uint16_t a = 0; a < currentANum; a++) { | 676 | for (uint16_t a = 0; a < currentANum; a++) { |
| 674 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 677 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 675 | - ReduceSum(mean, x1, pregLast); | 678 | + Reduce<ReduceType::SUM>(mean, x1, pregLast); |
| 676 | Muls(mean, mean, nCorrectionFactor, pregOne); | 679 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 677 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 680 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 678 | } | 681 | } |
| 679 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 682 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 680 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; | 683 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; |
| 681 | MaskReg pregLast = UpdateMask<float>(lastTailNum); | 684 | MaskReg pregLast = UpdateMask<float>(lastTailNum); |
| 682 | RegTensor<float> shlReg; | 685 | RegTensor<float> shlReg; |
| 683 | for (uint16_t a = 0; a < currentANum; a++) { | 686 | for (uint16_t a = 0; a < currentANum; a++) { |
| 684 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 687 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 685 | - DataCopy(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); | 688 | + LoadAlign(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); |
| 686 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), | 689 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), |
| 687 | pregLast); | 690 | pregLast); |
| 688 | Add(x1, x1, shlReg, pregFull); | 691 | Add(x1, x1, shlReg, pregFull); |
| 689 | - ReduceSum(mean, x1, pregFull); | 692 | + Reduce<ReduceType::SUM>(mean, x1, pregFull); |
| 690 | Muls(mean, mean, nCorrectionFactor, pregOne); | 693 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 691 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 694 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 692 | } | 695 | } |
| 693 | } | 696 | } |
| 694 | } | 697 | } |
| @@ -708,7 +711,7 @@ private: | |||
| 708 | MaskReg pregLoop; | 711 | MaskReg pregLoop; |
| 709 | 712 | ||
| 710 | for (uint16_t a = 0; a < currentANum; a++) { | 713 | for (uint16_t a = 0; a < currentANum; a++) { |
| 711 | - DataCopy<float, LoadDist::DIST_BRC_B32>(mean, meanInUb + a); | 714 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(mean, meanInUb + a); |
| 712 | uint32_t sregRemainder = binaryAddRemainder; | 715 | uint32_t sregRemainder = binaryAddRemainder; |
| 713 | for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; r++) { | 716 | for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; r++) { |
| 714 | pregLoop = UpdateMask<float>(sregRemainder); | 717 | pregLoop = UpdateMask<float>(sregRemainder); |
| @@ -723,8 +726,8 @@ private: | |||
| 723 | Muls(square1, square1, n, pregFull); | 726 | Muls(square1, square1, n, pregFull); |
| 724 | Muls(square2, square2, n, pregFull); | 727 | Muls(square2, square2, n, pregFull); |
| 725 | Add(varSum, square1, square2, pregFull); | 728 | Add(varSum, square1, square2, pregFull); |
| 726 | - ReduceSum(var, varSum, pregFull); | 729 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 727 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 730 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 728 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), var, pregOne); | 731 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), var, pregOne); |
| 729 | } | 732 | } |
| 730 | for (uint16_t r = 0; | 733 | for (uint16_t r = 0; |
| @@ -745,8 +748,8 @@ private: | |||
| 745 | Muls(square1, square1, n, pregFull); | 748 | Muls(square1, square1, n, pregFull); |
| 746 | Muls(square2, square2, n, pregLoop); | 749 | Muls(square2, square2, n, pregLoop); |
| 747 | Add(varSum, square1, square2, pregFull); | 750 | Add(varSum, square1, square2, pregFull); |
| 748 | - ReduceSum(var, varSum, pregFull); | 751 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 749 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 752 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 750 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), var, | 753 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), var, |
| 751 | pregOne); | 754 | pregOne); |
| 752 | } | 755 | } |
| @@ -759,8 +762,8 @@ private: | |||
| 759 | (r * VL_B32 + a * aStride)); | 762 | (r * VL_B32 + a * aStride)); |
| 760 | Mul(square1, xMeanSub1, xMeanSub1, pregFull); | 763 | Mul(square1, xMeanSub1, xMeanSub1, pregFull); |
| 761 | Muls(square1, square1, n, pregFull); | 764 | Muls(square1, square1, n, pregFull); |
| 762 | - ReduceSum(var, square1, pregFull); | 765 | + Reduce<ReduceType::SUM>(var, square1, pregFull); |
| 763 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 766 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 764 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), var, | 767 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), var, |
| 765 | pregOne); | 768 | pregOne); |
| 766 | } | 769 | } |
| @@ -769,31 +772,30 @@ private: | |||
| 769 | if constexpr (LAST_LOOP_NUMS == 1) { | 772 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 770 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNumTmp); | 773 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNumTmp); |
| 771 | for (uint16_t a = 0; a < currentANum; a++) { | 774 | for (uint16_t a = 0; a < currentANum; a++) { |
| 772 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 775 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 773 | - ReduceSum(var, x1, pregLast); | 776 | + Reduce<ReduceType::SUM>(var, x1, pregLast); |
| 774 | Muls(var, var, nCorrectionFactor, pregOne); | 777 | Muls(var, var, nCorrectionFactor, pregOne); |
| 775 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 778 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 776 | } | 779 | } |
| 777 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 780 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 778 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; | 781 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; |
| 779 | MaskReg pregLast = UpdateMask<float>(lastTailNum); | 782 | MaskReg pregLast = UpdateMask<float>(lastTailNum); |
| 780 | RegTensor<float> shlReg; | 783 | RegTensor<float> shlReg; |
| 781 | for (uint16_t a = 0; a < currentANum; a++) { | 784 | for (uint16_t a = 0; a < currentANum; a++) { |
| 782 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 785 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 783 | - DataCopy(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); | 786 | + LoadAlign(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); |
| 784 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), | 787 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), |
| 785 | pregLast); | 788 | pregLast); |
| 786 | Add(x1, x1, shlReg, pregFull); | 789 | Add(x1, x1, shlReg, pregFull); |
| 787 | - ReduceSum(var, x1, pregFull); | 790 | + Reduce<ReduceType::SUM>(var, x1, pregFull); |
| 788 | Muls(var, var, nCorrectionFactor, pregOne); | 791 | Muls(var, var, nCorrectionFactor, pregOne); |
| 789 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 792 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 790 | } | 793 | } |
| 791 | } | 794 | } |
| 792 | } | 795 | } |
| 793 | } | 796 | } |
| 794 | 797 | ||
| 795 | - __aicore__ inline void CalculateRstdVF(__local_mem__ float* rstdOutUb, __local_mem__ float* tmpUb, | 798 | + __aicore__ inline void CalculateRstdVF(__ubuf__ float* rstdOutUb, __ubuf__ float* tmpUb, uint16_t currentANum) |
| 796 | - uint16_t currentANum) | ||
| 797 | { | 799 | { |
| 798 | float epsilonLocal = tl_->epsilon; | 800 | float epsilonLocal = tl_->epsilon; |
| 799 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); | 801 | uint16_t aLoop = static_cast<uint16_t>((currentANum + VL_B32 - 1) / VL_B32); |
| @@ -805,32 +807,31 @@ private: | |||
| 805 | MaskReg pregLoop; | 807 | MaskReg pregLoop; |
| 806 | 808 | ||
| 807 | for (uint16_t a = 0; a < aLoop; a++) { | 809 | for (uint16_t a = 0; a < aLoop; a++) { |
| 808 | - DataCopy<float, LoadDist::DIST_NORM>(varReg, rstdOutUb + a * VL_B32); | 810 | + LoadAlign<float, LoadDist::DIST_NORM>(varReg, rstdOutUb + a * VL_B32); |
| 809 | pregLoop = UpdateMask<float>(sreg); | 811 | pregLoop = UpdateMask<float>(sreg); |
| 810 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(varReg, rstdReg, pregLoop, epsilonLocal); | 812 | NormCommon::ComputeRstdNewtonRaphsonReg<false>(varReg, rstdReg, pregLoop, epsilonLocal); |
| 811 | if constexpr (!IsOutRstd) { | 813 | if constexpr (!IsOutRstd) { |
| 812 | // variance stays in rstdOutUb for output, compute rstd to tmpUb | 814 | // variance stays in rstdOutUb for output, compute rstd to tmpUb |
| 813 | - DataCopy<float, StoreDist::DIST_NORM>(tmpUb + a * VL_B32, rstdReg, pregLoop); | 815 | + StoreAlign<float, StoreDist::DIST_NORM>(tmpUb + a * VL_B32, rstdReg, pregLoop); |
| 814 | } else { | 816 | } else { |
| 815 | // compute rstd in-place to rstdOutUb for output | 817 | // compute rstd in-place to rstdOutUb for output |
| 816 | - DataCopy<float, StoreDist::DIST_NORM>(rstdOutUb + a * VL_B32, rstdReg, pregLoop); | 818 | + StoreAlign<float, StoreDist::DIST_NORM>(rstdOutUb + a * VL_B32, rstdReg, pregLoop); |
| 817 | } | 819 | } |
| 818 | } | 820 | } |
| 819 | } | 821 | } |
| 820 | } | 822 | } |
| 821 | 823 | ||
| 822 | template <bool hasGammaFlag, bool hasBetaFlag> | 824 | template <bool hasGammaFlag, bool hasBetaFlag> |
| 823 | - __aicore__ inline void CalculateNormalizeVF(__local_mem__ float* xSubMeanUb, __local_mem__ U* betaInUb, | 825 | + __aicore__ inline void CalculateNormalizeVF(__ubuf__ float* xSubMeanUb, __ubuf__ U* betaInUb, __ubuf__ U* gammaInUb, |
| 824 | - __local_mem__ U* gammaInUb, __local_mem__ T* yOutUb, | 826 | + __ubuf__ T* yOutUb, __ubuf__ float* rstdOutUb, uint16_t currentANum) |
| 825 | - __local_mem__ float* rstdOutUb, uint16_t currentANum) | ||
| 826 | { | 827 | { |
| 827 | uint32_t reduceNum = tl_->r; | 828 | uint32_t reduceNum = tl_->r; |
| 828 | uint32_t aStride = tl_->rAlign; | 829 | uint32_t aStride = tl_->rAlign; |
| 829 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; | 830 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; |
| 830 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; | 831 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; |
| 831 | uint16_t remainderLoop = currentANum - remainderA; | 832 | uint16_t remainderLoop = currentANum - remainderA; |
| 832 | - __local_mem__ float* rstdOutUbPair = rstdOutUb + 1; | 833 | + __ubuf__ float* rstdOutUbPair = rstdOutUb + 1; |
| 833 | - __local_mem__ float* rstdOutUbRemainder = rstdOutUb + remainderA; | 834 | + __ubuf__ float* rstdOutUbRemainder = rstdOutUb + remainderA; |
| 834 | 835 | ||
| 835 | __VEC_SCOPE__ | 836 | __VEC_SCOPE__ |
| 836 | { | 837 | { |
| @@ -850,8 +851,8 @@ private: | |||
| 850 | MaskReg pregLoop; | 851 | MaskReg pregLoop; |
| 851 | 852 | ||
| 852 | for (uint16_t a = 0; a < static_cast<uint16_t>(currentANum / static_cast<uint16_t>(NUM_TWO)); a++) { | 853 | for (uint16_t a = 0; a < static_cast<uint16_t>(currentANum / static_cast<uint16_t>(NUM_TWO)); a++) { |
| 853 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); | 854 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); |
| 854 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); | 855 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); |
| 855 | uint32_t sreg0 = reduceNum; | 856 | uint32_t sreg0 = reduceNum; |
| 856 | for (uint16_t r = 0; r < loopCount; r++) { | 857 | for (uint16_t r = 0; r < loopCount; r++) { |
| 857 | pregLoop = UpdateMask<float>(sreg0); | 858 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -866,8 +867,8 @@ private: | |||
| 866 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); | 867 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); |
| 867 | } | 868 | } |
| 868 | if constexpr (hasGammaFlag && hasBetaFlag) { | 869 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 869 | - FusedMulDstAdd(y1, gamma, beta, pregLoop); | 870 | + MulDstAdd(y1, gamma, beta, pregLoop); |
| 870 | - FusedMulDstAdd(y2, gamma, beta, pregLoop); | 871 | + MulDstAdd(y2, gamma, beta, pregLoop); |
| 871 | } else { | 872 | } else { |
| 872 | if constexpr (hasGammaFlag) { | 873 | if constexpr (hasGammaFlag) { |
| 873 | Mul(y1, y1, gamma, pregLoop); | 874 | Mul(y1, y1, gamma, pregLoop); |
| @@ -883,7 +884,7 @@ private: | |||
| 883 | } | 884 | } |
| 884 | } | 885 | } |
| 885 | for (uint16_t a = 0; a < remainderLoop; a++) { | 886 | for (uint16_t a = 0; a < remainderLoop; a++) { |
| 886 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); | 887 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); |
| 887 | uint32_t sreg1 = reduceNum; | 888 | uint32_t sreg1 = reduceNum; |
| 888 | for (uint16_t r = 0; r < loopCount; r++) { | 889 | for (uint16_t r = 0; r < loopCount; r++) { |
| 889 | pregLoop = UpdateMask<float>(sreg1); | 890 | pregLoop = UpdateMask<float>(sreg1); |
| @@ -896,7 +897,7 @@ private: | |||
| 896 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); | 897 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); |
| 897 | } | 898 | } |
| 898 | if constexpr (hasGammaFlag && hasBetaFlag) { | 899 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 899 | - FusedMulDstAdd(yRemainder, gamma, beta, pregLoop); | 900 | + MulDstAdd(yRemainder, gamma, beta, pregLoop); |
| 900 | } else { | 901 | } else { |
| 901 | if constexpr (hasGammaFlag) { | 902 | if constexpr (hasGammaFlag) { |
| 902 | Mul(yRemainder, yRemainder, gamma, pregLoop); | 903 | Mul(yRemainder, yRemainder, gamma, pregLoop); |
| @@ -912,10 +913,9 @@ private: | |||
| 912 | } | 913 | } |
| 913 | 914 | ||
| 914 | template <bool hasGammaFlag, bool hasBetaFlag> | 915 | template <bool hasGammaFlag, bool hasBetaFlag> |
| 915 | - __aicore__ inline void CalculateNormalizeVFFullB(__local_mem__ float* xSubMeanUb, __local_mem__ U* betaInUb, | 916 | + __aicore__ inline void CalculateNormalizeVFFullB(__ubuf__ float* xSubMeanUb, __ubuf__ U* betaInUb, |
| 916 | - __local_mem__ U* gammaInUb, __local_mem__ T* yOutUb, | 917 | + __ubuf__ U* gammaInUb, __ubuf__ T* yOutUb, |
| 917 | - __local_mem__ float* rstdOutUb, uint16_t currentANum, | 918 | + __ubuf__ float* rstdOutUb, uint16_t currentANum, int64_t aOffset) |
| 918 | - int64_t aOffset) | ||
| 919 | { | 919 | { |
| 920 | uint32_t reduceNum = tl_->r; | 920 | uint32_t reduceNum = tl_->r; |
| 921 | uint32_t aStride = tl_->rAlign; | 921 | uint32_t aStride = tl_->rAlign; |
| @@ -939,47 +939,46 @@ private: | |||
| 939 | uint32_t numColAlignTwoGamma = NUM_TWO * gammaBetaStride; | 939 | uint32_t numColAlignTwoGamma = NUM_TWO * gammaBetaStride; |
| 940 | uint32_t numColAlignTwoSecond = b * aStride; | 940 | uint32_t numColAlignTwoSecond = b * aStride; |
| 941 | 941 | ||
| 942 | - __local_mem__ float* rstdOutUbPair = rstdOutUb + 1; | 942 | + __ubuf__ float* rstdOutUbPair = rstdOutUb + 1; |
| 943 | - __local_mem__ float* rstdOutUbRemainder = rstdOutUb + firstRemainderA; | 943 | + __ubuf__ float* rstdOutUbRemainder = rstdOutUb + firstRemainderA; |
| 944 | - __local_mem__ float* rstdOutUbSecondRemainder = rstdOutUb + firstEnd; | 944 | + __ubuf__ float* rstdOutUbSecondRemainder = rstdOutUb + firstEnd; |
| 945 | - __local_mem__ float* rstdOutUbThirdRemainder = rstdOutUb + firstEnd + b * secondLoopNum + thirdRemainderA; | 945 | + __ubuf__ float* rstdOutUbThirdRemainder = rstdOutUb + firstEnd + b * secondLoopNum + thirdRemainderA; |
| 946 | 946 | ||
| 947 | - __local_mem__ U* gammaInUbOne = gammaInUb + firstStart * gammaBetaStride; | 947 | + __ubuf__ U* gammaInUbOne = gammaInUb + firstStart * gammaBetaStride; |
| 948 | - __local_mem__ U* gammaInUbTwo = gammaInUb + (firstStart + 1) * gammaBetaStride; | 948 | + __ubuf__ U* gammaInUbTwo = gammaInUb + (firstStart + 1) * gammaBetaStride; |
| 949 | - __local_mem__ U* betaInUbOne = betaInUb + firstStart * gammaBetaStride; | 949 | + __ubuf__ U* betaInUbOne = betaInUb + firstStart * gammaBetaStride; |
| 950 | - __local_mem__ U* betaInUbTwo = betaInUb + (firstStart + 1) * gammaBetaStride; | 950 | + __ubuf__ U* betaInUbTwo = betaInUb + (firstStart + 1) * gammaBetaStride; |
| 951 | 951 | ||
| 952 | uint16_t firstPairLoopNum = firstEnd / NUM_TWO; | 952 | uint16_t firstPairLoopNum = firstEnd / NUM_TWO; |
| 953 | uint16_t bPairLoopNum = b / NUM_TWO; | 953 | uint16_t bPairLoopNum = b / NUM_TWO; |
| 954 | uint16_t thirdPairLoopNum = thirdEnd / NUM_TWO; | 954 | uint16_t thirdPairLoopNum = thirdEnd / NUM_TWO; |
| 955 | 955 | ||
| 956 | - __local_mem__ float* xSubMeanUbNext = xSubMeanUb + aStride; | 956 | + __ubuf__ float* xSubMeanUbNext = xSubMeanUb + aStride; |
| 957 | - __local_mem__ T* yOutUbNext = yOutUb + aStride; | 957 | + __ubuf__ T* yOutUbNext = yOutUb + aStride; |
| 958 | - __local_mem__ float* xSubMeanUbFirstRem = xSubMeanUb + firstRemainderA * aStride; | 958 | + __ubuf__ float* xSubMeanUbFirstRem = xSubMeanUb + firstRemainderA * aStride; |
| 959 | - __local_mem__ T* yOutUbFirstRem = yOutUb + aStride * firstRemainderA; | 959 | + __ubuf__ T* yOutUbFirstRem = yOutUb + aStride * firstRemainderA; |
| 960 | - __local_mem__ U* gammaInUbFirstRem = gammaInUb + (firstStart + firstRemainderA) * gammaBetaStride; | 960 | + __ubuf__ U* gammaInUbFirstRem = gammaInUb + (firstStart + firstRemainderA) * gammaBetaStride; |
| 961 | - __local_mem__ U* betaInUbFirstRem = betaInUb + (firstStart + firstRemainderA) * gammaBetaStride; | 961 | + __ubuf__ U* betaInUbFirstRem = betaInUb + (firstStart + firstRemainderA) * gammaBetaStride; |
| 962 | 962 | ||
| 963 | - __local_mem__ float* xSubMeanUbFirstEnd = xSubMeanUb + firstEnd * aStride; | 963 | + __ubuf__ float* xSubMeanUbFirstEnd = xSubMeanUb + firstEnd * aStride; |
| 964 | - __local_mem__ float* xSubMeanUbFirstEndNext = xSubMeanUb + aStride + firstEnd * aStride; | 964 | + __ubuf__ float* xSubMeanUbFirstEndNext = xSubMeanUb + aStride + firstEnd * aStride; |
| 965 | - __local_mem__ T* yOutUbFirstEnd = yOutUb + firstEnd * aStride; | 965 | + __ubuf__ T* yOutUbFirstEnd = yOutUb + firstEnd * aStride; |
| 966 | - __local_mem__ T* yOutUbFirstEndNext = yOutUb + aStride + firstEnd * aStride; | 966 | + __ubuf__ T* yOutUbFirstEndNext = yOutUb + aStride + firstEnd * aStride; |
| 967 | - __local_mem__ float* xSubMeanUbSecondRem = xSubMeanUb + (firstEnd + secondRemainderA) * aStride; | 967 | + __ubuf__ float* xSubMeanUbSecondRem = xSubMeanUb + (firstEnd + secondRemainderA) * aStride; |
| 968 | - __local_mem__ T* yOutUbSecondRem = yOutUb + (firstEnd + secondRemainderA) * aStride; | 968 | + __ubuf__ T* yOutUbSecondRem = yOutUb + (firstEnd + secondRemainderA) * aStride; |
| 969 | - __local_mem__ U* gammaInUbNext = gammaInUb + gammaBetaStride; | 969 | + __ubuf__ U* gammaInUbNext = gammaInUb + gammaBetaStride; |
| 970 | - __local_mem__ U* betaInUbNext = betaInUb + gammaBetaStride; | 970 | + __ubuf__ U* betaInUbNext = betaInUb + gammaBetaStride; |
| 971 | - __local_mem__ U* gammaInUbSecondRem = gammaInUb + secondRemainderA * gammaBetaStride; | 971 | + __ubuf__ U* gammaInUbSecondRem = gammaInUb + secondRemainderA * gammaBetaStride; |
| 972 | - __local_mem__ U* betaInUbSecondRem = betaInUb + secondRemainderA * gammaBetaStride; | 972 | + __ubuf__ U* betaInUbSecondRem = betaInUb + secondRemainderA * gammaBetaStride; |
| 973 | 973 | ||
| 974 | - __local_mem__ float* xSubMeanUbThird = xSubMeanUb + (firstEnd + b * secondLoopNum) * aStride; | 974 | + __ubuf__ float* xSubMeanUbThird = xSubMeanUb + (firstEnd + b * secondLoopNum) * aStride; |
| 975 | - __local_mem__ float* xSubMeanUbThirdNext = xSubMeanUb + (firstEnd + b * secondLoopNum + 1) * aStride; | 975 | + __ubuf__ float* xSubMeanUbThirdNext = xSubMeanUb + (firstEnd + b * secondLoopNum + 1) * aStride; |
| 976 | - __local_mem__ T* yOutUbThird = yOutUb + (firstEnd + b * secondLoopNum) * aStride; | 976 | + __ubuf__ T* yOutUbThird = yOutUb + (firstEnd + b * secondLoopNum) * aStride; |
| 977 | - __local_mem__ T* yOutUbThirdNext = yOutUb + (firstEnd + b * secondLoopNum + 1) * aStride; | 977 | + __ubuf__ T* yOutUbThirdNext = yOutUb + (firstEnd + b * secondLoopNum + 1) * aStride; |
| 978 | - __local_mem__ float* xSubMeanUbThirdRem = xSubMeanUb + | 978 | + __ubuf__ float* xSubMeanUbThirdRem = xSubMeanUb + (firstEnd + b * secondLoopNum + thirdRemainderA) * aStride; |
| 979 | - (firstEnd + b * secondLoopNum + thirdRemainderA) * aStride; | 979 | + __ubuf__ T* yOutUbThirdRem = yOutUb + (firstEnd + b * secondLoopNum + thirdRemainderA) * aStride; |
| 980 | - __local_mem__ T* yOutUbThirdRem = yOutUb + (firstEnd + b * secondLoopNum + thirdRemainderA) * aStride; | 980 | + __ubuf__ U* gammaInUbThirdRem = gammaInUb + thirdRemainderA * gammaBetaStride; |
| 981 | - __local_mem__ U* gammaInUbThirdRem = gammaInUb + thirdRemainderA * gammaBetaStride; | 981 | + __ubuf__ U* betaInUbThirdRem = betaInUb + thirdRemainderA * gammaBetaStride; |
| 982 | - __local_mem__ U* betaInUbThirdRem = betaInUb + thirdRemainderA * gammaBetaStride; | ||
| 983 | 982 | ||
| 984 | __VEC_SCOPE__ | 983 | __VEC_SCOPE__ |
| 985 | { | 984 | { |
| @@ -1003,8 +1002,8 @@ private: | |||
| 1003 | MaskReg pregLoop; | 1002 | MaskReg pregLoop; |
| 1004 | 1003 | ||
| 1005 | for (uint16_t a = 0; a < firstPairLoopNum; a++) { | 1004 | for (uint16_t a = 0; a < firstPairLoopNum; a++) { |
| 1006 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); | 1005 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); |
| 1007 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); | 1006 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); |
| 1008 | uint32_t sreg0 = reduceNum; | 1007 | uint32_t sreg0 = reduceNum; |
| 1009 | for (uint16_t r = 0; r < loopCount; r++) { | 1008 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1010 | pregLoop = UpdateMask<float>(sreg0); | 1009 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1027,8 +1026,8 @@ private: | |||
| 1027 | LoadTensorForDtypeTIn<U>(betaInUbTwo, beta2, pregLoop, gammaRegAddr); | 1026 | LoadTensorForDtypeTIn<U>(betaInUbTwo, beta2, pregLoop, gammaRegAddr); |
| 1028 | } | 1027 | } |
| 1029 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1028 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1030 | - FusedMulDstAdd(y1, gamma1, beta1, pregLoop); | 1029 | + MulDstAdd(y1, gamma1, beta1, pregLoop); |
| 1031 | - FusedMulDstAdd(y2, gamma2, beta2, pregLoop); | 1030 | + MulDstAdd(y2, gamma2, beta2, pregLoop); |
| 1032 | } else { | 1031 | } else { |
| 1033 | if constexpr (hasGammaFlag) { | 1032 | if constexpr (hasGammaFlag) { |
| 1034 | Mul(y1, y1, gamma1, pregLoop); | 1033 | Mul(y1, y1, gamma1, pregLoop); |
| @@ -1044,7 +1043,7 @@ private: | |||
| 1044 | } | 1043 | } |
| 1045 | } | 1044 | } |
| 1046 | for (uint16_t a = 0; a < firstRemainderLoop; a++) { | 1045 | for (uint16_t a = 0; a < firstRemainderLoop; a++) { |
| 1047 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); | 1046 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); |
| 1048 | uint32_t sreg0 = reduceNum; | 1047 | uint32_t sreg0 = reduceNum; |
| 1049 | for (uint16_t r = 0; r < loopCount; r++) { | 1048 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1050 | pregLoop = UpdateMask<float>(sreg0); | 1049 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1060,7 +1059,7 @@ private: | |||
| 1060 | LoadTensorForDtypeTIn<U>(betaInUbFirstRem, betaRemainder, pregLoop, gammaRegAddr); | 1059 | LoadTensorForDtypeTIn<U>(betaInUbFirstRem, betaRemainder, pregLoop, gammaRegAddr); |
| 1061 | } | 1060 | } |
| 1062 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1061 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1063 | - FusedMulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); | 1062 | + MulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); |
| 1064 | } else { | 1063 | } else { |
| 1065 | if constexpr (hasGammaFlag) { | 1064 | if constexpr (hasGammaFlag) { |
| 1066 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); | 1065 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); |
| @@ -1097,8 +1096,8 @@ private: | |||
| 1097 | 1096 | ||
| 1098 | for (uint16_t loop = 0; loop < secondLoopNum; loop++) { | 1097 | for (uint16_t loop = 0; loop < secondLoopNum; loop++) { |
| 1099 | for (uint16_t a = 0; a < bPairLoopNum; a++) { | 1098 | for (uint16_t a = 0; a < bPairLoopNum; a++) { |
| 1100 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + firstEnd + b * loop + a * NUM_TWO); | 1099 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + firstEnd + b * loop + a * NUM_TWO); |
| 1101 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + firstEnd + b * loop + a * NUM_TWO); | 1100 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + firstEnd + b * loop + a * NUM_TWO); |
| 1102 | uint32_t sreg0 = reduceNum; | 1101 | uint32_t sreg0 = reduceNum; |
| 1103 | for (uint16_t r = 0; r < loopCount; r++) { | 1102 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1104 | pregLoop = UpdateMask<float>(sreg0); | 1103 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1121,8 +1120,8 @@ private: | |||
| 1121 | LoadTensorForDtypeTIn<U>(betaInUbNext, beta2, pregLoop, gammaRegAddr); | 1120 | LoadTensorForDtypeTIn<U>(betaInUbNext, beta2, pregLoop, gammaRegAddr); |
| 1122 | } | 1121 | } |
| 1123 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1122 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1124 | - FusedMulDstAdd(y1, gamma1, beta1, pregLoop); | 1123 | + MulDstAdd(y1, gamma1, beta1, pregLoop); |
| 1125 | - FusedMulDstAdd(y2, gamma2, beta2, pregLoop); | 1124 | + MulDstAdd(y2, gamma2, beta2, pregLoop); |
| 1126 | } else { | 1125 | } else { |
| 1127 | if constexpr (hasGammaFlag) { | 1126 | if constexpr (hasGammaFlag) { |
| 1128 | Mul(y1, y1, gamma1, pregLoop); | 1127 | Mul(y1, y1, gamma1, pregLoop); |
| @@ -1138,8 +1137,8 @@ private: | |||
| 1138 | } | 1137 | } |
| 1139 | } | 1138 | } |
| 1140 | for (uint16_t a = 0; a < secondRemainderLoop; a++) { | 1139 | for (uint16_t a = 0; a < secondRemainderLoop; a++) { |
| 1141 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, | 1140 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, |
| 1142 | - rstdOutUbSecondRemainder + b * loop + secondRemainderA); | 1141 | + rstdOutUbSecondRemainder + b * loop + secondRemainderA); |
| 1143 | uint32_t sreg0 = reduceNum; | 1142 | uint32_t sreg0 = reduceNum; |
| 1144 | for (uint16_t r = 0; r < loopCount; r++) { | 1143 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1145 | pregLoop = UpdateMask<float>(sreg0); | 1144 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1157,7 +1156,7 @@ private: | |||
| 1157 | LoadTensorForDtypeTIn<U>(betaInUbSecondRem, betaRemainder, pregLoop, gammaRegAddr); | 1156 | LoadTensorForDtypeTIn<U>(betaInUbSecondRem, betaRemainder, pregLoop, gammaRegAddr); |
| 1158 | } | 1157 | } |
| 1159 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1158 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1160 | - FusedMulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); | 1159 | + MulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); |
| 1161 | } else { | 1160 | } else { |
| 1162 | if constexpr (hasGammaFlag) { | 1161 | if constexpr (hasGammaFlag) { |
| 1163 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); | 1162 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); |
| @@ -1194,9 +1193,10 @@ private: | |||
| 1194 | MaskReg pregLoop; | 1193 | MaskReg pregLoop; |
| 1195 | 1194 | ||
| 1196 | for (uint16_t a = 0; a < thirdPairLoopNum; a++) { | 1195 | for (uint16_t a = 0; a < thirdPairLoopNum; a++) { |
| 1197 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + firstEnd + b * secondLoopNum + a * NUM_TWO); | 1196 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt1, |
| 1198 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt2, | 1197 | + rstdOutUb + firstEnd + b * secondLoopNum + a * NUM_TWO); |
| 1199 | - rstdOutUbPair + firstEnd + b * secondLoopNum + a * NUM_TWO); | 1198 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt2, |
| 1199 | + rstdOutUbPair + firstEnd + b * secondLoopNum + a * NUM_TWO); | ||
| 1200 | uint32_t sreg0 = reduceNum; | 1200 | uint32_t sreg0 = reduceNum; |
| 1201 | for (uint16_t r = 0; r < loopCount; r++) { | 1201 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1202 | pregLoop = UpdateMask<float>(sreg0); | 1202 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1219,8 +1219,8 @@ private: | |||
| 1219 | LoadTensorForDtypeTIn<U>(betaInUbNext, beta2, pregLoop, gammaRegAddr); | 1219 | LoadTensorForDtypeTIn<U>(betaInUbNext, beta2, pregLoop, gammaRegAddr); |
| 1220 | } | 1220 | } |
| 1221 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1221 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1222 | - FusedMulDstAdd(y1, gamma1, beta1, pregLoop); | 1222 | + MulDstAdd(y1, gamma1, beta1, pregLoop); |
| 1223 | - FusedMulDstAdd(y2, gamma2, beta2, pregLoop); | 1223 | + MulDstAdd(y2, gamma2, beta2, pregLoop); |
| 1224 | } else { | 1224 | } else { |
| 1225 | if constexpr (hasGammaFlag) { | 1225 | if constexpr (hasGammaFlag) { |
| 1226 | Mul(y1, y1, gamma1, pregLoop); | 1226 | Mul(y1, y1, gamma1, pregLoop); |
| @@ -1236,7 +1236,7 @@ private: | |||
| 1236 | } | 1236 | } |
| 1237 | } | 1237 | } |
| 1238 | for (uint16_t a = 0; a < thirdRemainderLoop; a++) { | 1238 | for (uint16_t a = 0; a < thirdRemainderLoop; a++) { |
| 1239 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbThirdRemainder); | 1239 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbThirdRemainder); |
| 1240 | uint32_t sreg0 = reduceNum; | 1240 | uint32_t sreg0 = reduceNum; |
| 1241 | for (uint16_t r = 0; r < loopCount; r++) { | 1241 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1242 | pregLoop = UpdateMask<float>(sreg0); | 1242 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1252,7 +1252,7 @@ private: | |||
| 1252 | LoadTensorForDtypeTIn<U>(betaInUbThirdRem, betaRemainder, pregLoop, gammaRegAddr); | 1252 | LoadTensorForDtypeTIn<U>(betaInUbThirdRem, betaRemainder, pregLoop, gammaRegAddr); |
| 1253 | } | 1253 | } |
| 1254 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1254 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1255 | - FusedMulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); | 1255 | + MulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); |
| 1256 | } else { | 1256 | } else { |
| 1257 | if constexpr (hasGammaFlag) { | 1257 | if constexpr (hasGammaFlag) { |
| 1258 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); | 1258 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); |
| @@ -1268,9 +1268,9 @@ private: | |||
| 1268 | } | 1268 | } |
| 1269 | 1269 | ||
| 1270 | template <bool hasGammaFlag, bool hasBetaFlag> | 1270 | template <bool hasGammaFlag, bool hasBetaFlag> |
| 1271 | - __aicore__ inline void CalculateNormalizeVFNotFullB(__local_mem__ float* xSubMeanUb, __local_mem__ U* betaInUb, | 1271 | + __aicore__ inline void CalculateNormalizeVFNotFullB(__ubuf__ float* xSubMeanUb, __ubuf__ U* betaInUb, |
| 1272 | - __local_mem__ U* gammaInUb, __local_mem__ T* yOutUb, | 1272 | + __ubuf__ U* gammaInUb, __ubuf__ T* yOutUb, |
| 1273 | - __local_mem__ float* rstdOutUb, uint16_t currentANum) | 1273 | + __ubuf__ float* rstdOutUb, uint16_t currentANum) |
| 1274 | { | 1274 | { |
| 1275 | uint32_t reduceNum = tl_->r; | 1275 | uint32_t reduceNum = tl_->r; |
| 1276 | uint32_t aStride = tl_->rAlign; | 1276 | uint32_t aStride = tl_->rAlign; |
| @@ -1279,23 +1279,23 @@ private: | |||
| 1279 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; | 1279 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; |
| 1280 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; | 1280 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; |
| 1281 | uint16_t remainderLoop = currentANum - remainderA; | 1281 | uint16_t remainderLoop = currentANum - remainderA; |
| 1282 | - __local_mem__ float* rstdOutUbPair = rstdOutUb + 1; | 1282 | + __ubuf__ float* rstdOutUbPair = rstdOutUb + 1; |
| 1283 | - __local_mem__ float* rstdOutUbRemainder = rstdOutUb + remainderA; | 1283 | + __ubuf__ float* rstdOutUbRemainder = rstdOutUb + remainderA; |
| 1284 | 1284 | ||
| 1285 | uint32_t numColAlignTwo = NUM_TWO * aStride; | 1285 | uint32_t numColAlignTwo = NUM_TWO * aStride; |
| 1286 | uint32_t numColAlignTwoGamma = NUM_TWO * gammaBetaStride; | 1286 | uint32_t numColAlignTwoGamma = NUM_TWO * gammaBetaStride; |
| 1287 | 1287 | ||
| 1288 | uint16_t pairLoopNum = currentANum / static_cast<uint16_t>(NUM_TWO); | 1288 | uint16_t pairLoopNum = currentANum / static_cast<uint16_t>(NUM_TWO); |
| 1289 | 1289 | ||
| 1290 | - __local_mem__ float* xSubMeanUbNext = xSubMeanUb + aStride; | 1290 | + __ubuf__ float* xSubMeanUbNext = xSubMeanUb + aStride; |
| 1291 | - __local_mem__ T* yOutUbNext = yOutUb + aStride; | 1291 | + __ubuf__ T* yOutUbNext = yOutUb + aStride; |
| 1292 | - __local_mem__ U* gammaInUbNext = gammaInUb + gammaBetaStride; | 1292 | + __ubuf__ U* gammaInUbNext = gammaInUb + gammaBetaStride; |
| 1293 | - __local_mem__ U* betaInUbNext = betaInUb + gammaBetaStride; | 1293 | + __ubuf__ U* betaInUbNext = betaInUb + gammaBetaStride; |
| 1294 | 1294 | ||
| 1295 | - __local_mem__ float* xSubMeanUbRem = xSubMeanUb + remainderA * aStride; | 1295 | + __ubuf__ float* xSubMeanUbRem = xSubMeanUb + remainderA * aStride; |
| 1296 | - __local_mem__ T* yOutUbRem = yOutUb + aStride * remainderA; | 1296 | + __ubuf__ T* yOutUbRem = yOutUb + aStride * remainderA; |
| 1297 | - __local_mem__ U* gammaInUbRem = gammaInUb + remainderA * gammaBetaStride; | 1297 | + __ubuf__ U* gammaInUbRem = gammaInUb + remainderA * gammaBetaStride; |
| 1298 | - __local_mem__ U* betaInUbRem = betaInUb + remainderA * gammaBetaStride; | 1298 | + __ubuf__ U* betaInUbRem = betaInUb + remainderA * gammaBetaStride; |
| 1299 | 1299 | ||
| 1300 | __VEC_SCOPE__ | 1300 | __VEC_SCOPE__ |
| 1301 | { | 1301 | { |
| @@ -1319,8 +1319,8 @@ private: | |||
| 1319 | MaskReg pregLoop; | 1319 | MaskReg pregLoop; |
| 1320 | 1320 | ||
| 1321 | for (uint16_t a = 0; a < pairLoopNum; a++) { | 1321 | for (uint16_t a = 0; a < pairLoopNum; a++) { |
| 1322 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); | 1322 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); |
| 1323 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); | 1323 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); |
| 1324 | uint32_t sreg0 = reduceNum; | 1324 | uint32_t sreg0 = reduceNum; |
| 1325 | for (uint16_t r = 0; r < loopCount; r++) { | 1325 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1326 | pregLoop = UpdateMask<float>(sreg0); | 1326 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1343,8 +1343,8 @@ private: | |||
| 1343 | LoadTensorForDtypeTIn<U>(betaInUbNext, beta2, pregLoop, gammaRegAddr); | 1343 | LoadTensorForDtypeTIn<U>(betaInUbNext, beta2, pregLoop, gammaRegAddr); |
| 1344 | } | 1344 | } |
| 1345 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1345 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1346 | - FusedMulDstAdd(y1, gamma1, beta1, pregLoop); | 1346 | + MulDstAdd(y1, gamma1, beta1, pregLoop); |
| 1347 | - FusedMulDstAdd(y2, gamma2, beta2, pregLoop); | 1347 | + MulDstAdd(y2, gamma2, beta2, pregLoop); |
| 1348 | } else { | 1348 | } else { |
| 1349 | if constexpr (hasGammaFlag) { | 1349 | if constexpr (hasGammaFlag) { |
| 1350 | Mul(y1, y1, gamma1, pregLoop); | 1350 | Mul(y1, y1, gamma1, pregLoop); |
| @@ -1360,7 +1360,7 @@ private: | |||
| 1360 | } | 1360 | } |
| 1361 | } | 1361 | } |
| 1362 | for (uint16_t a = 0; a < remainderLoop; a++) { | 1362 | for (uint16_t a = 0; a < remainderLoop; a++) { |
| 1363 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); | 1363 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); |
| 1364 | uint32_t sreg0 = reduceNum; | 1364 | uint32_t sreg0 = reduceNum; |
| 1365 | for (uint16_t r = 0; r < loopCount; r++) { | 1365 | for (uint16_t r = 0; r < loopCount; r++) { |
| 1366 | pregLoop = UpdateMask<float>(sreg0); | 1366 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -1376,7 +1376,7 @@ private: | |||
| 1376 | LoadTensorForDtypeTIn<U>(betaInUbRem, betaRemainder, pregLoop, gammaRegAddr); | 1376 | LoadTensorForDtypeTIn<U>(betaInUbRem, betaRemainder, pregLoop, gammaRegAddr); |
| 1377 | } | 1377 | } |
| 1378 | if constexpr (hasGammaFlag && hasBetaFlag) { | 1378 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 1379 | - FusedMulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); | 1379 | + MulDstAdd(yRemainder, gammaRemainder, betaRemainder, pregLoop); |
| 1380 | } else { | 1380 | } else { |
| 1381 | if constexpr (hasGammaFlag) { | 1381 | if constexpr (hasGammaFlag) { |
| 1382 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); | 1382 | Mul(yRemainder, yRemainder, gammaRemainder, pregLoop); |
| @@ -1393,10 +1393,10 @@ private: | |||
| 1393 | 1393 | ||
| 1394 | __aicore__ inline void CastMeanRstd(int64_t currentANum) | 1394 | __aicore__ inline void CastMeanRstd(int64_t currentANum) |
| 1395 | { | 1395 | { |
| 1396 | - __local_mem__ float* meanInAddr = (__local_mem__ float*)meanOutUb_.GetPhyAddr(); | 1396 | + __ubuf__ float* meanInAddr = (__ubuf__ float*)meanOutUb_.GetPhyAddr(); |
| 1397 | - __local_mem__ float* rstdInAddr = (__local_mem__ float*)rstdOutUb_.GetPhyAddr(); | 1397 | + __ubuf__ float* rstdInAddr = (__ubuf__ float*)rstdOutUb_.GetPhyAddr(); |
| 1398 | - __local_mem__ M* meanOutAddr = (__local_mem__ M*)meanOutUb_.GetPhyAddr(); | 1398 | + __ubuf__ M* meanOutAddr = (__ubuf__ M*)meanOutUb_.GetPhyAddr(); |
| 1399 | - __local_mem__ M* rstdOutAddr = (__local_mem__ M*)rstdOutUb_.GetPhyAddr(); | 1399 | + __ubuf__ M* rstdOutAddr = (__ubuf__ M*)rstdOutUb_.GetPhyAddr(); |
| 1400 | 1400 | ||
| 1401 | uint32_t castCount = static_cast<uint32_t>(currentANum); | 1401 | uint32_t castCount = static_cast<uint32_t>(currentANum); |
| 1402 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_B32 - 1) / VL_B32); | 1402 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_B32 - 1) / VL_B32); |
| @@ -1409,14 +1409,12 @@ private: | |||
| 1409 | MicroAPI::MaskReg pregLoop; | 1409 | MicroAPI::MaskReg pregLoop; |
| 1410 | for (uint16_t i = 0; i < castLoops; i++) { | 1410 | for (uint16_t i = 0; i < castLoops; i++) { |
| 1411 | pregLoop = MicroAPI::UpdateMask<float>(castCount); | 1411 | pregLoop = MicroAPI::UpdateMask<float>(castCount); |
| 1412 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, meanInAddr + VL_B32 * i); | 1412 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, meanInAddr + VL_B32 * i); |
| 1413 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_rstd, rstdInAddr + VL_B32 * i); | 1413 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_rstd, rstdInAddr + VL_B32 * i); |
| 1414 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 1414 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 1415 | Cast<M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); | 1415 | Cast<M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); |
| 1416 | - DataCopy<M, StoreDist::DIST_PACK_B32>(((__local_mem__ M*)meanOutAddr + i * VL_B16), output_mean, | 1416 | + StoreAlign<M, StoreDist::DIST_PACK_B32>(((__ubuf__ M*)meanOutAddr + i * VL_B16), output_mean, pregLoop); |
| 1417 | - pregLoop); | 1417 | + StoreAlign<M, StoreDist::DIST_PACK_B32>(((__ubuf__ M*)rstdOutAddr + i * VL_B16), output_rstd, pregLoop); |
| 1418 | - DataCopy<M, StoreDist::DIST_PACK_B32>(((__local_mem__ M*)rstdOutAddr + i * VL_B16), output_rstd, | ||
| 1419 | - pregLoop); | ||
| 1420 | } | 1418 | } |
| 1421 | } | 1419 | } |
| 1422 | } | 1420 | } |
| @@ -29,6 +29,7 @@ using AscendC::MicroAPI::MemType; | |||
| 29 | using AscendC::MicroAPI::RegTensor; | 29 | using AscendC::MicroAPI::RegTensor; |
| 30 | using AscendC::MicroAPI::StoreDist; | 30 | using AscendC::MicroAPI::StoreDist; |
| 31 | using AscendC::MicroAPI::UpdateMask; | 31 | using AscendC::MicroAPI::UpdateMask; |
| 32 | +using AscendC::Reg::StoreAlign; | ||
| 32 | 33 | ||
| 33 | constexpr static LayerNormConfig hasGammaBetaConfig = { | 34 | constexpr static LayerNormConfig hasGammaBetaConfig = { |
| 34 | false, | 35 | false, |
| @@ -296,10 +297,10 @@ private: | |||
| 296 | 297 | ||
| 297 | __aicore__ inline void CastBatchMeanLastout(uint64_t currentANum) | 298 | __aicore__ inline void CastBatchMeanLastout(uint64_t currentANum) |
| 298 | { | 299 | { |
| 299 | - __local_mem__ float* batchMeanInAddr = (__local_mem__ float*)batchMeanOutUb.GetPhyAddr(); | 300 | + __ubuf__ float* batchMeanInAddr = (__ubuf__ float*)batchMeanOutUb.GetPhyAddr(); |
| 300 | - __local_mem__ float* batchLastoutInAddr = (__local_mem__ float*)batchLastoutOutUb.GetPhyAddr(); | 301 | + __ubuf__ float* batchLastoutInAddr = (__ubuf__ float*)batchLastoutOutUb.GetPhyAddr(); |
| 301 | - __local_mem__ M* batchMeanOutAddr = (__local_mem__ M*)batchMeanOutUb.GetPhyAddr(); | 302 | + __ubuf__ M* batchMeanOutAddr = (__ubuf__ M*)batchMeanOutUb.GetPhyAddr(); |
| 302 | - __local_mem__ M* batchLastoutOutAddr = (__local_mem__ M*)batchLastoutOutUb.GetPhyAddr(); | 303 | + __ubuf__ M* batchLastoutOutAddr = (__ubuf__ M*)batchLastoutOutUb.GetPhyAddr(); |
| 303 | 304 | ||
| 304 | uint32_t castCount = static_cast<uint32_t>(currentANum); | 305 | uint32_t castCount = static_cast<uint32_t>(currentANum); |
| 305 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); | 306 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_F32 - 1) / VL_F32); |
| @@ -312,15 +313,15 @@ private: | |||
| 312 | MicroAPI::MaskReg pregLoop; | 313 | MicroAPI::MaskReg pregLoop; |
| 313 | for (uint16_t i = 0; i < castLoops; i++) { | 314 | for (uint16_t i = 0; i < castLoops; i++) { |
| 314 | pregLoop = MicroAPI::UpdateMask<float>(castCount); | 315 | pregLoop = MicroAPI::UpdateMask<float>(castCount); |
| 315 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); | 316 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, batchMeanInAddr + VL_F32 * i); |
| 316 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_lastout, | 317 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_lastout, |
| 317 | - batchLastoutInAddr + VL_F32 * i); | 318 | + batchLastoutInAddr + VL_F32 * i); |
| 318 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 319 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 319 | Cast<M, float, castTraitB322B16>(output_lastout, input_lastout, pregLoop); | 320 | Cast<M, float, castTraitB322B16>(output_lastout, input_lastout, pregLoop); |
| 320 | - DataCopy<M, StoreDist::DIST_PACK_B32>(((__local_mem__ M*)batchMeanOutAddr + i * VL_MEAN), output_mean, | 321 | + StoreAlign<M, StoreDist::DIST_PACK_B32>(((__ubuf__ M*)batchMeanOutAddr + i * VL_MEAN), output_mean, |
| 321 | - pregLoop); | 322 | + pregLoop); |
| 322 | - DataCopy<M, StoreDist::DIST_PACK_B32>(((__local_mem__ M*)batchLastoutOutAddr + i * VL_MEAN), | 323 | + StoreAlign<M, StoreDist::DIST_PACK_B32>(((__ubuf__ M*)batchLastoutOutAddr + i * VL_MEAN), |
| 323 | - output_lastout, pregLoop); | 324 | + output_lastout, pregLoop); |
| 324 | } | 325 | } |
| 325 | } | 326 | } |
| 326 | } | 327 | } |
| @@ -30,6 +30,9 @@ using AscendC::MicroAPI::MemType; | |||
| 30 | using AscendC::MicroAPI::RegTensor; | 30 | using AscendC::MicroAPI::RegTensor; |
| 31 | using AscendC::MicroAPI::StoreDist; | 31 | using AscendC::MicroAPI::StoreDist; |
| 32 | using AscendC::MicroAPI::UpdateMask; | 32 | using AscendC::MicroAPI::UpdateMask; |
| 33 | +using AscendC::Reg::LoadAlign; | ||
| 34 | +using AscendC::Reg::Reduce; | ||
| 35 | +using AscendC::Reg::StoreAlign; | ||
| 33 | using NormCommon::NormCommonRegbase::LoadRegForDtype; | 36 | using NormCommon::NormCommonRegbase::LoadRegForDtype; |
| 34 | using NormCommon::NormCommonRegbase::StoreRegForDtype; | 37 | using NormCommon::NormCommonRegbase::StoreRegForDtype; |
| 35 | 38 | ||
| @@ -137,11 +140,11 @@ private: | |||
| 137 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); | 140 | rstdOutUb_ = rstdQueue_.AllocTensor<float>(); |
| 138 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); | 141 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); |
| 139 | 142 | ||
| 140 | - __local_mem__ T* xInUbAddr = (__local_mem__ T*)xInUb.GetPhyAddr(); | 143 | + __ubuf__ T* xInUbAddr = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 141 | - __local_mem__ float* meanOutUbAddr = (__local_mem__ float*)meanOutUb_.GetPhyAddr(); | 144 | + __ubuf__ float* meanOutUbAddr = (__ubuf__ float*)meanOutUb_.GetPhyAddr(); |
| 142 | - __local_mem__ float* rstdOutUbAddr = (__local_mem__ float*)rstdOutUb_.GetPhyAddr(); | 145 | + __ubuf__ float* rstdOutUbAddr = (__ubuf__ float*)rstdOutUb_.GetPhyAddr(); |
| 143 | - __local_mem__ float* xSubMeanUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr(); | 146 | + __ubuf__ float* xSubMeanUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr(); |
| 144 | - __local_mem__ float* tmpUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr() + elemNum_; | 147 | + __ubuf__ float* tmpUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr() + elemNum_; |
| 145 | 148 | ||
| 146 | if (tl_->rAlign <= VL_B32) { | 149 | if (tl_->rAlign <= VL_B32) { |
| 147 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); | 150 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); |
| @@ -156,9 +159,9 @@ private: | |||
| 156 | CalculateRstdVF(rstdOutUbAddr, currentANum); | 159 | CalculateRstdVF(rstdOutUbAddr, currentANum); |
| 157 | 160 | ||
| 158 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); | 161 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); |
| 159 | - __local_mem__ U* gammaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr(); | 162 | + __ubuf__ U* gammaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr(); |
| 160 | - __local_mem__ U* betaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAlign; | 163 | + __ubuf__ U* betaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAlign; |
| 161 | - __local_mem__ T* yOutUbAddr = (__local_mem__ T*)yOutUb.GetPhyAddr(); | 164 | + __ubuf__ T* yOutUbAddr = (__ubuf__ T*)yOutUb.GetPhyAddr(); |
| 162 | if (hasGamma_ && hasBeta_) { | 165 | if (hasGamma_ && hasBeta_) { |
| 163 | CalculateNormalizeVF<true, true>(xSubMeanUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, rstdOutUbAddr, | 166 | CalculateNormalizeVF<true, true>(xSubMeanUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, rstdOutUbAddr, |
| 164 | currentANum); | 167 | currentANum); |
| @@ -176,8 +179,8 @@ private: | |||
| 176 | yQueue_.EnQue(yOutUb); | 179 | yQueue_.EnQue(yOutUb); |
| 177 | } | 180 | } |
| 178 | 181 | ||
| 179 | - __aicore__ inline void CalculateMeanVarRLessThanVL(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 182 | + __aicore__ inline void CalculateMeanVarRLessThanVL(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 180 | - __local_mem__ float* rstdInUb, __local_mem__ float* xSubMeanUb, | 183 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 181 | uint16_t currentANum) | 184 | uint16_t currentANum) |
| 182 | { | 185 | { |
| 183 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 186 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| @@ -202,25 +205,25 @@ private: | |||
| 202 | for (uint16_t a = 0; a < currentANum; a++) { | 205 | for (uint16_t a = 0; a < currentANum; a++) { |
| 203 | LoadRegForDtype(xInUb, x, pregLoop, (a * aStride)); | 206 | LoadRegForDtype(xInUb, x, pregLoop, (a * aStride)); |
| 204 | Muls(meanSum, x, n, pregLoop); | 207 | Muls(meanSum, x, n, pregLoop); |
| 205 | - ReduceSum(mean, meanSum, pregLoop); | 208 | + Reduce<ReduceType::SUM>(mean, meanSum, pregLoop); |
| 206 | Muls(mean, mean, nCorrectionFactor, pregOne); | 209 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 207 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 210 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 208 | 211 | ||
| 209 | Duplicate(meanDup, mean, pregFull); | 212 | Duplicate(meanDup, mean, pregFull); |
| 210 | Sub(xMeanSub, x, meanDup, pregLoop); | 213 | Sub(xMeanSub, x, meanDup, pregLoop); |
| 211 | StoreRegForDtype(xSubMeanUb, xMeanSub, pregLoop, (a * aStride)); | 214 | StoreRegForDtype(xSubMeanUb, xMeanSub, pregLoop, (a * aStride)); |
| 212 | Mul(square, xMeanSub, xMeanSub, pregLoop); | 215 | Mul(square, xMeanSub, xMeanSub, pregLoop); |
| 213 | Muls(varSum, square, n, pregLoop); | 216 | Muls(varSum, square, n, pregLoop); |
| 214 | - ReduceSum(var, varSum, pregLoop); | 217 | + Reduce<ReduceType::SUM>(var, varSum, pregLoop); |
| 215 | Muls(var, var, nCorrectionFactor, pregOne); | 218 | Muls(var, var, nCorrectionFactor, pregOne); |
| 216 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 219 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 217 | } | 220 | } |
| 218 | } | 221 | } |
| 219 | } | 222 | } |
| 220 | 223 | ||
| 221 | - __aicore__ inline void CalculateMeanVarRLessThanTwoVL(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 224 | + __aicore__ inline void CalculateMeanVarRLessThanTwoVL(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 222 | - __local_mem__ float* rstdInUb, | 225 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 223 | - __local_mem__ float* xSubMeanUb, uint16_t currentANum) | 226 | + uint16_t currentANum) |
| 224 | { | 227 | { |
| 225 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 228 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| 226 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); | 229 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); |
| @@ -256,9 +259,9 @@ private: | |||
| 256 | Muls(meanSum1, x1, n, pregFull); | 259 | Muls(meanSum1, x1, n, pregFull); |
| 257 | Muls(meanSum2, x2, n, pregTail); | 260 | Muls(meanSum2, x2, n, pregTail); |
| 258 | Add(meanSum, meanSum1, meanSum2, pregFull); | 261 | Add(meanSum, meanSum1, meanSum2, pregFull); |
| 259 | - ReduceSum(mean, meanSum, pregFull); | 262 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 260 | Muls(mean, mean, nCorrectionFactor, pregOne); | 263 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 261 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 264 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 262 | 265 | ||
| 263 | Duplicate(meanDup, mean, pregFull); | 266 | Duplicate(meanDup, mean, pregFull); |
| 264 | Sub(xMeanSub1, x1, meanDup, pregFull); | 267 | Sub(xMeanSub1, x1, meanDup, pregFull); |
| @@ -270,17 +273,17 @@ private: | |||
| 270 | Muls(varSum1, square1, n, pregFull); | 273 | Muls(varSum1, square1, n, pregFull); |
| 271 | Muls(varSum2, square2, n, pregTail); | 274 | Muls(varSum2, square2, n, pregTail); |
| 272 | Add(varSum, varSum1, varSum2, pregFull); | 275 | Add(varSum, varSum1, varSum2, pregFull); |
| 273 | - ReduceSum(var, varSum, pregFull); | 276 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 274 | Muls(var, var, nCorrectionFactor, pregOne); | 277 | Muls(var, var, nCorrectionFactor, pregOne); |
| 275 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 278 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 276 | } | 279 | } |
| 277 | } | 280 | } |
| 278 | } | 281 | } |
| 279 | 282 | ||
| 280 | template <int32_t LAST_LOOP_NUMS> | 283 | template <int32_t LAST_LOOP_NUMS> |
| 281 | - __aicore__ inline void CalculateMeanVarRCommon(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 284 | + __aicore__ inline void CalculateMeanVarRCommon(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 282 | - __local_mem__ float* rstdInUb, __local_mem__ float* xSubMeanUb, | 285 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 283 | - __local_mem__ float* tmpUb, uint16_t currentANum) | 286 | + __ubuf__ float* tmpUb, uint16_t currentANum) |
| 284 | { | 287 | { |
| 285 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 288 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| 286 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); | 289 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); |
| @@ -319,8 +322,8 @@ private: | |||
| 319 | Muls(x1, x1, n, pregFull); | 322 | Muls(x1, x1, n, pregFull); |
| 320 | Muls(x2, x2, n, pregFull); | 323 | Muls(x2, x2, n, pregFull); |
| 321 | Add(meanSum, x1, x2, pregFull); | 324 | Add(meanSum, x1, x2, pregFull); |
| 322 | - ReduceSum(mean, meanSum, pregFull); | 325 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 323 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 326 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 324 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), mean, pregOne); | 327 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), mean, pregOne); |
| 325 | } | 328 | } |
| 326 | for (uint16_t r = 0; | 329 | for (uint16_t r = 0; |
| @@ -333,8 +336,8 @@ private: | |||
| 333 | Muls(x1, x1, n, pregFull); | 336 | Muls(x1, x1, n, pregFull); |
| 334 | Muls(x2, x2, n, pregLoop); | 337 | Muls(x2, x2, n, pregLoop); |
| 335 | Add(meanSum, x1, x2, pregFull); | 338 | Add(meanSum, x1, x2, pregFull); |
| 336 | - ReduceSum(mean, meanSum, pregFull); | 339 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 337 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 340 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 338 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, | 341 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, |
| 339 | pregOne); | 342 | pregOne); |
| 340 | } | 343 | } |
| @@ -343,8 +346,8 @@ private: | |||
| 343 | LoadRegForDtype(xInUb + binaryAddRemainderCeilLoop * VL_B32, x1, pregFull, | 346 | LoadRegForDtype(xInUb + binaryAddRemainderCeilLoop * VL_B32, x1, pregFull, |
| 344 | (r * VL_B32 + a * aStride)); | 347 | (r * VL_B32 + a * aStride)); |
| 345 | Muls(x1, x1, n, pregFull); | 348 | Muls(x1, x1, n, pregFull); |
| 346 | - ReduceSum(mean, x1, pregFull); | 349 | + Reduce<ReduceType::SUM>(mean, x1, pregFull); |
| 347 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 350 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 348 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, | 351 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, |
| 349 | pregOne); | 352 | pregOne); |
| 350 | } | 353 | } |
| @@ -353,24 +356,24 @@ private: | |||
| 353 | if constexpr (LAST_LOOP_NUMS == 1) { | 356 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 354 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); | 357 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); |
| 355 | for (uint16_t a = 0; a < currentANum; a++) { | 358 | for (uint16_t a = 0; a < currentANum; a++) { |
| 356 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 359 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 357 | - ReduceSum(mean, x1, pregLast); | 360 | + Reduce<ReduceType::SUM>(mean, x1, pregLast); |
| 358 | Muls(mean, mean, nCorrectionFactor, pregOne); | 361 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 359 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 362 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 360 | } | 363 | } |
| 361 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 364 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 362 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; | 365 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; |
| 363 | MaskReg pregLast = UpdateMask<float>(lastTailNum); | 366 | MaskReg pregLast = UpdateMask<float>(lastTailNum); |
| 364 | RegTensor<float> shlReg; | 367 | RegTensor<float> shlReg; |
| 365 | for (uint16_t a = 0; a < currentANum; a++) { | 368 | for (uint16_t a = 0; a < currentANum; a++) { |
| 366 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 369 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 367 | - DataCopy(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); | 370 | + LoadAlign(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); |
| 368 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), | 371 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), |
| 369 | pregLast); | 372 | pregLast); |
| 370 | Add(x1, x1, shlReg, pregFull); | 373 | Add(x1, x1, shlReg, pregFull); |
| 371 | - ReduceSum(mean, x1, pregFull); | 374 | + Reduce<ReduceType::SUM>(mean, x1, pregFull); |
| 372 | Muls(mean, mean, nCorrectionFactor, pregOne); | 375 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 373 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 376 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 374 | } | 377 | } |
| 375 | } | 378 | } |
| 376 | } | 379 | } |
| @@ -390,7 +393,7 @@ private: | |||
| 390 | MaskReg pregLoop; | 393 | MaskReg pregLoop; |
| 391 | 394 | ||
| 392 | for (uint16_t a = 0; a < currentANum; a++) { | 395 | for (uint16_t a = 0; a < currentANum; a++) { |
| 393 | - DataCopy<float, LoadDist::DIST_BRC_B32>(mean, meanInUb + a); | 396 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(mean, meanInUb + a); |
| 394 | uint32_t sregRemainder = binaryAddRemainder; | 397 | uint32_t sregRemainder = binaryAddRemainder; |
| 395 | for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; r++) { | 398 | for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; r++) { |
| 396 | pregLoop = UpdateMask<float>(sregRemainder); | 399 | pregLoop = UpdateMask<float>(sregRemainder); |
| @@ -405,8 +408,8 @@ private: | |||
| 405 | Muls(square1, square1, n, pregFull); | 408 | Muls(square1, square1, n, pregFull); |
| 406 | Muls(square2, square2, n, pregFull); | 409 | Muls(square2, square2, n, pregFull); |
| 407 | Add(varSum, square1, square2, pregFull); | 410 | Add(varSum, square1, square2, pregFull); |
| 408 | - ReduceSum(var, varSum, pregFull); | 411 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 409 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 412 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 410 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), var, pregOne); | 413 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), var, pregOne); |
| 411 | } | 414 | } |
| 412 | for (uint16_t r = 0; | 415 | for (uint16_t r = 0; |
| @@ -427,8 +430,8 @@ private: | |||
| 427 | Muls(square1, square1, n, pregFull); | 430 | Muls(square1, square1, n, pregFull); |
| 428 | Muls(square2, square2, n, pregLoop); | 431 | Muls(square2, square2, n, pregLoop); |
| 429 | Add(varSum, square1, square2, pregFull); | 432 | Add(varSum, square1, square2, pregFull); |
| 430 | - ReduceSum(var, varSum, pregFull); | 433 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 431 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 434 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 432 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), var, | 435 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), var, |
| 433 | pregOne); | 436 | pregOne); |
| 434 | } | 437 | } |
| @@ -441,8 +444,8 @@ private: | |||
| 441 | (r * VL_B32 + a * aStride)); | 444 | (r * VL_B32 + a * aStride)); |
| 442 | Mul(square1, xMeanSub1, xMeanSub1, pregFull); | 445 | Mul(square1, xMeanSub1, xMeanSub1, pregFull); |
| 443 | Muls(square1, square1, n, pregFull); | 446 | Muls(square1, square1, n, pregFull); |
| 444 | - ReduceSum(var, square1, pregFull); | 447 | + Reduce<ReduceType::SUM>(var, square1, pregFull); |
| 445 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 448 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 446 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), var, | 449 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), var, |
| 447 | pregOne); | 450 | pregOne); |
| 448 | } | 451 | } |
| @@ -451,30 +454,30 @@ private: | |||
| 451 | if constexpr (LAST_LOOP_NUMS == 1) { | 454 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 452 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNumTmp); | 455 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNumTmp); |
| 453 | for (uint16_t a = 0; a < currentANum; a++) { | 456 | for (uint16_t a = 0; a < currentANum; a++) { |
| 454 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 457 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 455 | - ReduceSum(var, x1, pregLast); | 458 | + Reduce<ReduceType::SUM>(var, x1, pregLast); |
| 456 | Muls(var, var, nCorrectionFactor, pregOne); | 459 | Muls(var, var, nCorrectionFactor, pregOne); |
| 457 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 460 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 458 | } | 461 | } |
| 459 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 462 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 460 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; | 463 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; |
| 461 | MaskReg pregLast = UpdateMask<float>(lastTailNum); | 464 | MaskReg pregLast = UpdateMask<float>(lastTailNum); |
| 462 | RegTensor<float> shlReg; | 465 | RegTensor<float> shlReg; |
| 463 | for (uint16_t a = 0; a < currentANum; a++) { | 466 | for (uint16_t a = 0; a < currentANum; a++) { |
| 464 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 467 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 465 | - DataCopy(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); | 468 | + LoadAlign(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); |
| 466 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), | 469 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), |
| 467 | pregLast); | 470 | pregLast); |
| 468 | Add(x1, x1, shlReg, pregFull); | 471 | Add(x1, x1, shlReg, pregFull); |
| 469 | - ReduceSum(var, x1, pregFull); | 472 | + Reduce<ReduceType::SUM>(var, x1, pregFull); |
| 470 | Muls(var, var, nCorrectionFactor, pregOne); | 473 | Muls(var, var, nCorrectionFactor, pregOne); |
| 471 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 474 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 472 | } | 475 | } |
| 473 | } | 476 | } |
| 474 | } | 477 | } |
| 475 | } | 478 | } |
| 476 | 479 | ||
| 477 | - __aicore__ inline void CalculateRstdVF(__local_mem__ float* rstdOutUb, uint16_t currentANum) | 480 | + __aicore__ inline void CalculateRstdVF(__ubuf__ float* rstdOutUb, uint16_t currentANum) |
| 478 | { | 481 | { |
| 479 | float epsilonLocal = tl_->epsilon; | 482 | float epsilonLocal = tl_->epsilon; |
| 480 | NormCommon::ComputeRstdNewtonRaphson<false>(rstdOutUb, rstdOutUb, static_cast<uint32_t>(currentANum), | 483 | NormCommon::ComputeRstdNewtonRaphson<false>(rstdOutUb, rstdOutUb, static_cast<uint32_t>(currentANum), |
| @@ -482,17 +485,16 @@ private: | |||
| 482 | } | 485 | } |
| 483 | 486 | ||
| 484 | template <bool hasGammaFlag, bool hasBetaFlag> | 487 | template <bool hasGammaFlag, bool hasBetaFlag> |
| 485 | - __aicore__ inline void CalculateNormalizeVF(__local_mem__ float* xSubMeanUb, __local_mem__ U* betaInUb, | 488 | + __aicore__ inline void CalculateNormalizeVF(__ubuf__ float* xSubMeanUb, __ubuf__ U* betaInUb, __ubuf__ U* gammaInUb, |
| 486 | - __local_mem__ U* gammaInUb, __local_mem__ T* yOutUb, | 489 | + __ubuf__ T* yOutUb, __ubuf__ float* rstdOutUb, uint16_t currentANum) |
| 487 | - __local_mem__ float* rstdOutUb, uint16_t currentANum) | ||
| 488 | { | 490 | { |
| 489 | uint32_t reduceNum = tl_->r; | 491 | uint32_t reduceNum = tl_->r; |
| 490 | uint32_t aStride = tl_->rAlign; | 492 | uint32_t aStride = tl_->rAlign; |
| 491 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; | 493 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; |
| 492 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; | 494 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; |
| 493 | uint16_t remainderLoop = currentANum - remainderA; | 495 | uint16_t remainderLoop = currentANum - remainderA; |
| 494 | - __local_mem__ float* rstdOutUbPair = rstdOutUb + 1; | 496 | + __ubuf__ float* rstdOutUbPair = rstdOutUb + 1; |
| 495 | - __local_mem__ float* rstdOutUbRemainder = rstdOutUb + remainderA; | 497 | + __ubuf__ float* rstdOutUbRemainder = rstdOutUb + remainderA; |
| 496 | 498 | ||
| 497 | __VEC_SCOPE__ | 499 | __VEC_SCOPE__ |
| 498 | { | 500 | { |
| @@ -512,8 +514,8 @@ private: | |||
| 512 | MaskReg pregLoop; | 514 | MaskReg pregLoop; |
| 513 | 515 | ||
| 514 | for (uint16_t a = 0; a < static_cast<uint16_t>(currentANum / static_cast<uint16_t>(NUM_TWO)); a++) { | 516 | for (uint16_t a = 0; a < static_cast<uint16_t>(currentANum / static_cast<uint16_t>(NUM_TWO)); a++) { |
| 515 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); | 517 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); |
| 516 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); | 518 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); |
| 517 | uint32_t sreg0 = reduceNum; | 519 | uint32_t sreg0 = reduceNum; |
| 518 | for (uint16_t r = 0; r < loopCount; r++) { | 520 | for (uint16_t r = 0; r < loopCount; r++) { |
| 519 | pregLoop = UpdateMask<float>(sreg0); | 521 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -528,8 +530,8 @@ private: | |||
| 528 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); | 530 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); |
| 529 | } | 531 | } |
| 530 | if constexpr (hasGammaFlag && hasBetaFlag) { | 532 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 531 | - FusedMulDstAdd(y1, gamma, beta, pregLoop); | 533 | + MulDstAdd(y1, gamma, beta, pregLoop); |
| 532 | - FusedMulDstAdd(y2, gamma, beta, pregLoop); | 534 | + MulDstAdd(y2, gamma, beta, pregLoop); |
| 533 | } else { | 535 | } else { |
| 534 | if constexpr (hasGammaFlag) { | 536 | if constexpr (hasGammaFlag) { |
| 535 | Mul(y1, y1, gamma, pregLoop); | 537 | Mul(y1, y1, gamma, pregLoop); |
| @@ -545,7 +547,7 @@ private: | |||
| 545 | } | 547 | } |
| 546 | } | 548 | } |
| 547 | for (uint16_t a = 0; a < remainderLoop; a++) { | 549 | for (uint16_t a = 0; a < remainderLoop; a++) { |
| 548 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); | 550 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); |
| 549 | uint32_t sreg1 = reduceNum; | 551 | uint32_t sreg1 = reduceNum; |
| 550 | for (uint16_t r = 0; r < loopCount; r++) { | 552 | for (uint16_t r = 0; r < loopCount; r++) { |
| 551 | pregLoop = UpdateMask<float>(sreg1); | 553 | pregLoop = UpdateMask<float>(sreg1); |
| @@ -558,7 +560,7 @@ private: | |||
| 558 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); | 560 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); |
| 559 | } | 561 | } |
| 560 | if constexpr (hasGammaFlag && hasBetaFlag) { | 562 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 561 | - FusedMulDstAdd(yRemainder, gamma, beta, pregLoop); | 563 | + MulDstAdd(yRemainder, gamma, beta, pregLoop); |
| 562 | } else { | 564 | } else { |
| 563 | if constexpr (hasGammaFlag) { | 565 | if constexpr (hasGammaFlag) { |
| 564 | Mul(yRemainder, yRemainder, gamma, pregLoop); | 566 | Mul(yRemainder, yRemainder, gamma, pregLoop); |
| @@ -575,10 +577,10 @@ private: | |||
| 575 | 577 | ||
| 576 | __aicore__ inline void CastMeanRstd(int64_t currentANum) | 578 | __aicore__ inline void CastMeanRstd(int64_t currentANum) |
| 577 | { | 579 | { |
| 578 | - __local_mem__ float* meanInAddr = (__local_mem__ float*)meanOutUb_.GetPhyAddr(); | 580 | + __ubuf__ float* meanInAddr = (__ubuf__ float*)meanOutUb_.GetPhyAddr(); |
| 579 | - __local_mem__ float* rstdInAddr = (__local_mem__ float*)rstdOutUb_.GetPhyAddr(); | 581 | + __ubuf__ float* rstdInAddr = (__ubuf__ float*)rstdOutUb_.GetPhyAddr(); |
| 580 | - __local_mem__ M* meanOutAddr = (__local_mem__ M*)meanOutUb_.GetPhyAddr(); | 582 | + __ubuf__ M* meanOutAddr = (__ubuf__ M*)meanOutUb_.GetPhyAddr(); |
| 581 | - __local_mem__ M* rstdOutAddr = (__local_mem__ M*)rstdOutUb_.GetPhyAddr(); | 583 | + __ubuf__ M* rstdOutAddr = (__ubuf__ M*)rstdOutUb_.GetPhyAddr(); |
| 582 | 584 | ||
| 583 | uint32_t castCount = static_cast<uint32_t>(currentANum); | 585 | uint32_t castCount = static_cast<uint32_t>(currentANum); |
| 584 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_B32 - 1) / VL_B32); | 586 | uint16_t castLoops = static_cast<uint32_t>((castCount + VL_B32 - 1) / VL_B32); |
| @@ -591,14 +593,12 @@ private: | |||
| 591 | MicroAPI::MaskReg pregLoop; | 593 | MicroAPI::MaskReg pregLoop; |
| 592 | for (uint16_t i = 0; i < castLoops; i++) { | 594 | for (uint16_t i = 0; i < castLoops; i++) { |
| 593 | pregLoop = MicroAPI::UpdateMask<float>(castCount); | 595 | pregLoop = MicroAPI::UpdateMask<float>(castCount); |
| 594 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, meanInAddr + VL_B32 * i); | 596 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_mean, meanInAddr + VL_B32 * i); |
| 595 | - MicroAPI::DataCopy<float, MicroAPI::LoadDist::DIST_NORM>(input_rstd, rstdInAddr + VL_B32 * i); | 597 | + MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_NORM>(input_rstd, rstdInAddr + VL_B32 * i); |
| 596 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 598 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 597 | Cast<M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); | 599 | Cast<M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); |
| 598 | - DataCopy<M, StoreDist::DIST_PACK_B32>(((__local_mem__ M*)meanOutAddr + i * VL_B16), output_mean, | 600 | + StoreAlign<M, StoreDist::DIST_PACK_B32>(((__ubuf__ M*)meanOutAddr + i * VL_B16), output_mean, pregLoop); |
| 599 | - pregLoop); | 601 | + StoreAlign<M, StoreDist::DIST_PACK_B32>(((__ubuf__ M*)rstdOutAddr + i * VL_B16), output_rstd, pregLoop); |
| 600 | - DataCopy<M, StoreDist::DIST_PACK_B32>(((__local_mem__ M*)rstdOutAddr + i * VL_B16), output_rstd, | ||
| 601 | - pregLoop); | ||
| 602 | } | 602 | } |
| 603 | } | 603 | } |
| 604 | } | 604 | } |
| @@ -31,6 +31,8 @@ using AscendC::MicroAPI::MaskReg; | |||
| 31 | using AscendC::MicroAPI::RegTensor; | 31 | using AscendC::MicroAPI::RegTensor; |
| 32 | using AscendC::MicroAPI::StoreDist; | 32 | using AscendC::MicroAPI::StoreDist; |
| 33 | using AscendC::MicroAPI::UpdateMask; | 33 | using AscendC::MicroAPI::UpdateMask; |
| 34 | +using AscendC::Reg::LoadAlign; | ||
| 35 | +using AscendC::Reg::StoreAlign; | ||
| 34 | 36 | ||
| 35 | template <typename T, typename U, typename M, bool IsOutRstd> | 37 | template <typename T, typename U, typename M, bool IsOutRstd> |
| 36 | class LayerNormV3WelfordMultiReduce { | 38 | class LayerNormV3WelfordMultiReduce { |
| @@ -242,13 +244,13 @@ private: | |||
| 242 | { | 244 | { |
| 243 | int64_t r0Aligned = td_->r0Align; | 245 | int64_t r0Aligned = td_->r0Align; |
| 244 | 246 | ||
| 245 | - __local_mem__ float* meanAddr = (__local_mem__ float*)meanTensor_.GetPhyAddr() + cacheCount_; | 247 | + __ubuf__ float* meanAddr = (__ubuf__ float*)meanTensor_.GetPhyAddr() + cacheCount_; |
| 246 | - __local_mem__ float* rstdAddr; | 248 | + __ubuf__ float* rstdAddr; |
| 247 | 249 | ||
| 248 | if constexpr (IsOutRstd) { | 250 | if constexpr (IsOutRstd) { |
| 249 | // variance is in varianceTensor_, compute rstd and store to lastoutTensor_ (for GM output) | 251 | // variance is in varianceTensor_, compute rstd and store to lastoutTensor_ (for GM output) |
| 250 | - __local_mem__ float* varAddr = (__local_mem__ float*)varianceTensor_.GetPhyAddr() + cacheCount_; | 252 | + __ubuf__ float* varAddr = (__ubuf__ float*)varianceTensor_.GetPhyAddr() + cacheCount_; |
| 251 | - rstdAddr = (__local_mem__ float*)lastoutTensor_.GetPhyAddr() + cacheCount_; | 253 | + rstdAddr = (__ubuf__ float*)lastoutTensor_.GetPhyAddr() + cacheCount_; |
| 252 | { | 254 | { |
| 253 | __VEC_SCOPE__ | 255 | __VEC_SCOPE__ |
| 254 | { | 256 | { |
| @@ -256,14 +258,14 @@ private: | |||
| 256 | RegTensor<float> rstdRegTmp; | 258 | RegTensor<float> rstdRegTmp; |
| 257 | MaskReg pregAll = CreateMask<float, MaskPattern::ALL>(); | 259 | MaskReg pregAll = CreateMask<float, MaskPattern::ALL>(); |
| 258 | MaskReg pregOne = CreateMask<float, MaskPattern::VL1>(); | 260 | MaskReg pregOne = CreateMask<float, MaskPattern::VL1>(); |
| 259 | - DataCopy<float, LoadDist::DIST_BRC_B32>(varReg, varAddr); | 261 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(varReg, varAddr); |
| 260 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdRegTmp, pregAll, td_->epsilon); | 262 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdRegTmp, pregAll, td_->epsilon); |
| 261 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr, rstdRegTmp, pregOne); | 263 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr, rstdRegTmp, pregOne); |
| 262 | } | 264 | } |
| 263 | } | 265 | } |
| 264 | } else { | 266 | } else { |
| 265 | - __local_mem__ float* varAddr = (__local_mem__ float*)lastoutTensor_.GetPhyAddr() + cacheCount_; | 267 | + __ubuf__ float* varAddr = (__ubuf__ float*)lastoutTensor_.GetPhyAddr() + cacheCount_; |
| 266 | - rstdAddr = (__local_mem__ float*)rstdTensor_.GetPhyAddr() + cacheCount_; | 268 | + rstdAddr = (__ubuf__ float*)rstdTensor_.GetPhyAddr() + cacheCount_; |
| 267 | { | 269 | { |
| 268 | __VEC_SCOPE__ | 270 | __VEC_SCOPE__ |
| 269 | { | 271 | { |
| @@ -271,21 +273,21 @@ private: | |||
| 271 | RegTensor<float> rstdRegTmp; | 273 | RegTensor<float> rstdRegTmp; |
| 272 | MaskReg pregAll = CreateMask<float, MaskPattern::ALL>(); | 274 | MaskReg pregAll = CreateMask<float, MaskPattern::ALL>(); |
| 273 | MaskReg pregOne = CreateMask<float, MaskPattern::VL1>(); | 275 | MaskReg pregOne = CreateMask<float, MaskPattern::VL1>(); |
| 274 | - DataCopy<float, LoadDist::DIST_BRC_B32>(varReg, varAddr); | 276 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(varReg, varAddr); |
| 275 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdRegTmp, pregAll, td_->epsilon); | 277 | NormCommon::ComputeRstdNewtonRaphsonReg(varReg, rstdRegTmp, pregAll, td_->epsilon); |
| 276 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr, rstdRegTmp, pregOne); | 278 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdAddr, rstdRegTmp, pregOne); |
| 277 | } | 279 | } |
| 278 | } | 280 | } |
| 279 | } | 281 | } |
| 280 | 282 | ||
| 281 | - __local_mem__ U* gammaAddr; | 283 | + __ubuf__ U* gammaAddr; |
| 282 | - __local_mem__ U* betaAddr; | 284 | + __ubuf__ U* betaAddr; |
| 283 | if (td_->r1ComputeFactor > 1) { | 285 | if (td_->r1ComputeFactor > 1) { |
| 284 | - gammaAddr = (__local_mem__ U*)gammaPackedLocal_.GetPhyAddr(); | 286 | + gammaAddr = (__ubuf__ U*)gammaPackedLocal_.GetPhyAddr(); |
| 285 | - betaAddr = (__local_mem__ U*)betaPackedLocal_.GetPhyAddr(); | 287 | + betaAddr = (__ubuf__ U*)betaPackedLocal_.GetPhyAddr(); |
| 286 | } else { | 288 | } else { |
| 287 | - gammaAddr = (__local_mem__ U*)gammaResident_.GetPhyAddr(); | 289 | + gammaAddr = (__ubuf__ U*)gammaResident_.GetPhyAddr(); |
| 288 | - betaAddr = (__local_mem__ U*)betaResident_.GetPhyAddr(); | 290 | + betaAddr = (__ubuf__ U*)betaResident_.GetPhyAddr(); |
| 289 | } | 291 | } |
| 290 | 292 | ||
| 291 | for (int64_t r1Loop = 0; r1Loop < td_->loopR1outer; r1Loop++) { | 293 | for (int64_t r1Loop = 0; r1Loop < td_->loopR1outer; r1Loop++) { |
| @@ -313,8 +315,8 @@ private: | |||
| 313 | 315 | ||
| 314 | LocalTensor<T> yTensor = outQueueY_.template AllocTensor<T>(); | 316 | LocalTensor<T> yTensor = outQueueY_.template AllocTensor<T>(); |
| 315 | 317 | ||
| 316 | - __local_mem__ T* xUbAddr = (__local_mem__ T*)xTensor.GetPhyAddr(); | 318 | + __ubuf__ T* xUbAddr = (__ubuf__ T*)xTensor.GetPhyAddr(); |
| 317 | - __local_mem__ T* yUbAddr = (__local_mem__ T*)yTensor.GetPhyAddr(); | 319 | + __ubuf__ T* yUbAddr = (__ubuf__ T*)yTensor.GetPhyAddr(); |
| 318 | NormalizeCutR1VF(xUbAddr, yUbAddr, gammaAddr, betaAddr, meanAddr, rstdAddr, r0Aligned, curR1); | 320 | NormalizeCutR1VF(xUbAddr, yUbAddr, gammaAddr, betaAddr, meanAddr, rstdAddr, r0Aligned, curR1); |
| 319 | 321 | ||
| 320 | inQueueX_.FreeTensor(xTensor); | 322 | inQueueX_.FreeTensor(xTensor); |
| @@ -331,9 +333,8 @@ private: | |||
| 331 | } | 333 | } |
| 332 | } | 334 | } |
| 333 | 335 | ||
| 334 | - __aicore__ inline void NormalizeCutR1VF(__local_mem__ T* xAddr, __local_mem__ T* yOutAddr, | 336 | + __aicore__ inline void NormalizeCutR1VF(__ubuf__ T* xAddr, __ubuf__ T* yOutAddr, __ubuf__ U* gammaAddr, |
| 335 | - __local_mem__ U* gammaAddr, __local_mem__ U* betaAddr, | 337 | + __ubuf__ U* betaAddr, __ubuf__ float* meanAddr, __ubuf__ float* rstdAddr, |
| 336 | - __local_mem__ float* meanAddr, __local_mem__ float* rstdAddr, | ||
| 337 | int64_t r0Aligned, int64_t curR1) | 338 | int64_t r0Aligned, int64_t curR1) |
| 338 | { | 339 | { |
| 339 | int64_t r1ComputeFactor = td_->r1ComputeFactor; | 340 | int64_t r1ComputeFactor = td_->r1ComputeFactor; |
| @@ -343,8 +344,8 @@ private: | |||
| 343 | return; | 344 | return; |
| 344 | } | 345 | } |
| 345 | 346 | ||
| 346 | - __local_mem__ U* gammaPackedAddr = (__local_mem__ U*)gammaPackedLocal_.GetPhyAddr(); | 347 | + __ubuf__ U* gammaPackedAddr = (__ubuf__ U*)gammaPackedLocal_.GetPhyAddr(); |
| 347 | - __local_mem__ U* betaPackedAddr = (__local_mem__ U*)betaPackedLocal_.GetPhyAddr(); | 348 | + __ubuf__ U* betaPackedAddr = (__ubuf__ U*)betaPackedLocal_.GetPhyAddr(); |
| 348 | 349 | ||
| 349 | uint32_t packedLen = static_cast<uint32_t>(r1ComputeFactor * r0Aligned); | 350 | uint32_t packedLen = static_cast<uint32_t>(r1ComputeFactor * r0Aligned); |
| 350 | int64_t mainLoops = curR1 / r1ComputeFactor; | 351 | int64_t mainLoops = curR1 / r1ComputeFactor; |
| @@ -364,8 +365,8 @@ private: | |||
| 364 | RegTensor<float> yReg; | 365 | RegTensor<float> yReg; |
| 365 | MaskReg pregFull; | 366 | MaskReg pregFull; |
| 366 | 367 | ||
| 367 | - DataCopy<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr); | 368 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr); |
| 368 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr); | 369 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr); |
| 369 | 370 | ||
| 370 | uint32_t fullSreg = packedLen; | 371 | uint32_t fullSreg = packedLen; |
| 371 | pregFull = UpdateMask<float>(fullSreg); | 372 | pregFull = UpdateMask<float>(fullSreg); |
| @@ -395,8 +396,8 @@ private: | |||
| 395 | RegTensor<float> yReg; | 396 | RegTensor<float> yReg; |
| 396 | MaskReg pregTail; | 397 | MaskReg pregTail; |
| 397 | 398 | ||
| 398 | - DataCopy<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr); | 399 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr); |
| 399 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr); | 400 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr); |
| 400 | 401 | ||
| 401 | uint32_t fullSreg = tailLen; | 402 | uint32_t fullSreg = tailLen; |
| 402 | pregTail = UpdateMask<float>(fullSreg); | 403 | pregTail = UpdateMask<float>(fullSreg); |
| @@ -413,10 +414,9 @@ private: | |||
| 413 | } | 414 | } |
| 414 | } | 415 | } |
| 415 | 416 | ||
| 416 | - __aicore__ inline void NormalizeCutR1VFSingle(__local_mem__ T* xAddr, __local_mem__ T* yOutAddr, | 417 | + __aicore__ inline void NormalizeCutR1VFSingle(__ubuf__ T* xAddr, __ubuf__ T* yOutAddr, __ubuf__ U* gammaAddr, |
| 417 | - __local_mem__ U* gammaAddr, __local_mem__ U* betaAddr, | 418 | + __ubuf__ U* betaAddr, __ubuf__ float* meanAddr, |
| 418 | - __local_mem__ float* meanAddr, __local_mem__ float* rstdAddr, | 419 | + __ubuf__ float* rstdAddr, int64_t r0Aligned, int64_t curR1) |
| 419 | - int64_t r0Aligned, int64_t curR1) | ||
| 420 | { | 420 | { |
| 421 | uint32_t r0Num = static_cast<uint32_t>(td_->r0); | 421 | uint32_t r0Num = static_cast<uint32_t>(td_->r0); |
| 422 | uint16_t loopCount = static_cast<uint16_t>((r0Num + VL_B32 - 1) / VL_B32); | 422 | uint16_t loopCount = static_cast<uint16_t>((r0Num + VL_B32 - 1) / VL_B32); |
| @@ -431,8 +431,8 @@ private: | |||
| 431 | RegTensor<float> yReg; | 431 | RegTensor<float> yReg; |
| 432 | MaskReg pregLoop; | 432 | MaskReg pregLoop; |
| 433 | 433 | ||
| 434 | - DataCopy<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr); | 434 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(meanReg, meanAddr); |
| 435 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr); | 435 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rstdReg, rstdAddr); |
| 436 | 436 | ||
| 437 | uint32_t sreg = r0Num; | 437 | uint32_t sreg = r0Num; |
| 438 | for (uint16_t r = 0; r < loopCount; r++) { | 438 | for (uint16_t r = 0; r < loopCount; r++) { |
| @@ -459,28 +459,28 @@ private: | |||
| 459 | } | 459 | } |
| 460 | 460 | ||
| 461 | template <typename DType> | 461 | template <typename DType> |
| 462 | - __aicore__ inline void LoadTensorForDtype(RegTensor<float>& dst, __local_mem__ DType* src, MaskReg& preg, | 462 | + __aicore__ inline void LoadTensorForDtype(RegTensor<float>& dst, __ubuf__ DType* src, MaskReg& preg, |
| 463 | uint32_t offset) | 463 | uint32_t offset) |
| 464 | { | 464 | { |
| 465 | if constexpr (IsSameType<DType, float>::value) { | 465 | if constexpr (IsSameType<DType, float>::value) { |
| 466 | - DataCopy<float, LoadDist::DIST_NORM>(dst, src + offset); | 466 | + LoadAlign<float, LoadDist::DIST_NORM>(dst, src + offset); |
| 467 | } else { | 467 | } else { |
| 468 | RegTensor<DType> tmp; | 468 | RegTensor<DType> tmp; |
| 469 | - DataCopy<DType, LoadDist::DIST_UNPACK_B16>(tmp, src + offset); | 469 | + LoadAlign<DType, LoadDist::DIST_UNPACK_B16>(tmp, src + offset); |
| 470 | Cast<float, DType, castTraitB162B32>(dst, tmp, preg); | 470 | Cast<float, DType, castTraitB162B32>(dst, tmp, preg); |
| 471 | } | 471 | } |
| 472 | } | 472 | } |
| 473 | 473 | ||
| 474 | template <typename DType> | 474 | template <typename DType> |
| 475 | - __aicore__ inline void StoreTensorForDtype(__local_mem__ DType* dst, RegTensor<float>& src, MaskReg& preg, | 475 | + __aicore__ inline void StoreTensorForDtype(__ubuf__ DType* dst, RegTensor<float>& src, MaskReg& preg, |
| 476 | uint32_t offset) | 476 | uint32_t offset) |
| 477 | { | 477 | { |
| 478 | if constexpr (IsSameType<DType, float>::value) { | 478 | if constexpr (IsSameType<DType, float>::value) { |
| 479 | - DataCopy<DType, StoreDist::DIST_NORM>(dst + offset, src, preg); | 479 | + StoreAlign<DType, StoreDist::DIST_NORM>(dst + offset, src, preg); |
| 480 | } else { | 480 | } else { |
| 481 | RegTensor<DType> tmp; | 481 | RegTensor<DType> tmp; |
| 482 | Cast<DType, float, castTraitB322B16>(tmp, src, preg); | 482 | Cast<DType, float, castTraitB322B16>(tmp, src, preg); |
| 483 | - DataCopy<DType, StoreDist::DIST_PACK_B32>(dst + offset, tmp, preg); | 483 | + StoreAlign<DType, StoreDist::DIST_PACK_B32>(dst + offset, tmp, preg); |
| 484 | } | 484 | } |
| 485 | } | 485 | } |
| 486 | 486 | ||
| @@ -609,4 +609,4 @@ private: | |||
| 609 | 609 | ||
| 610 | } // namespace LayerNormV3 | 610 | } // namespace LayerNormV3 |
| 611 | 611 | ||
| 612 | -#endif // LAYER_NORM_V3_WELFORD_MULTI_REDUCE_H | 612 | +#endif // LAYER_NORM_V3_WELFORD_MULTI_REDUCE_H |
| @@ -21,6 +21,7 @@ | |||
| 21 | 21 | ||
| 22 | namespace LayerNormV4 { | 22 | namespace LayerNormV4 { |
| 23 | using namespace AscendC; | 23 | using namespace AscendC; |
| 24 | +using AscendC::Reg::StoreAlign; | ||
| 24 | 25 | ||
| 25 | constexpr static int64_t BLOCK_SIZE = 32; | 26 | constexpr static int64_t BLOCK_SIZE = 32; |
| 26 | constexpr static uint32_t FLOAT_BYTES = 4; | 27 | constexpr static uint32_t FLOAT_BYTES = 4; |
| @@ -49,9 +50,9 @@ constexpr static AscendC::MicroAPI::CastTrait castTraitB322B16 = { | |||
| 49 | }; | 50 | }; |
| 50 | 51 | ||
| 51 | template <typename M> | 52 | template <typename M> |
| 52 | -__aicore__ inline void CastBatchMeanRstdToDtype(__local_mem__ float* batchMeanInAddr, | 53 | +__aicore__ inline void CastBatchMeanRstdToDtype(__ubuf__ float* batchMeanInAddr, __ubuf__ float* batchRstdInAddr, |
| 53 | - __local_mem__ float* batchRstdInAddr, __local_mem__ M* batchMeanOutAddr, | 54 | + __ubuf__ M* batchMeanOutAddr, __ubuf__ M* batchRstdOutAddr, |
| 54 | - __local_mem__ M* batchRstdOutAddr, uint64_t currentANum) | 55 | + uint64_t currentANum) |
| 55 | { | 56 | { |
| 56 | constexpr uint32_t VL_F32 = AscendC::VECTOR_REG_WIDTH / sizeof(float); | 57 | constexpr uint32_t VL_F32 = AscendC::VECTOR_REG_WIDTH / sizeof(float); |
| 57 | constexpr uint32_t VL_MEAN = AscendC::VECTOR_REG_WIDTH / sizeof(M); | 58 | constexpr uint32_t VL_MEAN = AscendC::VECTOR_REG_WIDTH / sizeof(M); |
| @@ -66,16 +67,16 @@ __aicore__ inline void CastBatchMeanRstdToDtype(__local_mem__ float* batchMeanIn | |||
| 66 | AscendC::MicroAPI::MaskReg pregLoop; | 67 | AscendC::MicroAPI::MaskReg pregLoop; |
| 67 | for (uint16_t i = 0; i < castLoops; i++) { | 68 | for (uint16_t i = 0; i < castLoops; i++) { |
| 68 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(castCount); | 69 | pregLoop = AscendC::MicroAPI::UpdateMask<float>(castCount); |
| 69 | - AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(input_mean, | 70 | + AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(input_mean, |
| 70 | - batchMeanInAddr + VL_F32 * i); | 71 | + batchMeanInAddr + VL_F32 * i); |
| 71 | - AscendC::MicroAPI::DataCopy<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(input_rstd, | 72 | + AscendC::MicroAPI::LoadAlign<float, AscendC::MicroAPI::LoadDist::DIST_NORM>(input_rstd, |
| 72 | - batchRstdInAddr + VL_F32 * i); | 73 | + batchRstdInAddr + VL_F32 * i); |
| 73 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); | 74 | Cast<M, float, castTraitB322B16>(output_mean, input_mean, pregLoop); |
| 74 | Cast<M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); | 75 | Cast<M, float, castTraitB322B16>(output_rstd, input_rstd, pregLoop); |
| 75 | - DataCopy<M, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(batchMeanOutAddr + i * VL_MEAN, output_mean, | 76 | + StoreAlign<M, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(batchMeanOutAddr + i * VL_MEAN, output_mean, |
| 76 | - pregLoop); | 77 | + pregLoop); |
| 77 | - DataCopy<M, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(batchRstdOutAddr + i * VL_MEAN, output_rstd, | 78 | + StoreAlign<M, AscendC::MicroAPI::StoreDist::DIST_PACK_B32>(batchRstdOutAddr + i * VL_MEAN, output_rstd, |
| 78 | - pregLoop); | 79 | + pregLoop); |
| 79 | } | 80 | } |
| 80 | } | 81 | } |
| 81 | } | 82 | } |
| @@ -250,10 +250,9 @@ private: | |||
| 250 | { | 250 | { |
| 251 | if constexpr (!IsSameType<M, float>::value) { | 251 | if constexpr (!IsSameType<M, float>::value) { |
| 252 | // float to bfloat16 or float16, input continue and output each repeat have only half value | 252 | // float to bfloat16 or float16, input continue and output each repeat have only half value |
| 253 | - CastBatchMeanRstdToDtype<M>((__local_mem__ float*)batchMeanOutUb.GetPhyAddr(), | 253 | + CastBatchMeanRstdToDtype<M>( |
| 254 | - (__local_mem__ float*)batchRstdOutUb.GetPhyAddr(), | 254 | + (__ubuf__ float*)batchMeanOutUb.GetPhyAddr(), (__ubuf__ float*)batchRstdOutUb.GetPhyAddr(), |
| 255 | - (__local_mem__ M*)batchMeanOutUb.GetPhyAddr(), | 255 | + (__ubuf__ M*)batchMeanOutUb.GetPhyAddr(), (__ubuf__ M*)batchRstdOutUb.GetPhyAddr(), currentANum); |
| 256 | - (__local_mem__ M*)batchRstdOutUb.GetPhyAddr(), currentANum); | ||
| 257 | batchMeanQueue.EnQue(batchMeanOutUb); | 256 | batchMeanQueue.EnQue(batchMeanOutUb); |
| 258 | batchRstdQueue.EnQue(batchRstdOutUb); | 257 | batchRstdQueue.EnQue(batchRstdOutUb); |
| 259 | LocalTensor<M> batchMeanInUb = batchMeanQueue.template DeQue<M>(); | 258 | LocalTensor<M> batchMeanInUb = batchMeanQueue.template DeQue<M>(); |
| @@ -30,6 +30,9 @@ using AscendC::MicroAPI::MemType; | |||
| 30 | using AscendC::MicroAPI::RegTensor; | 30 | using AscendC::MicroAPI::RegTensor; |
| 31 | using AscendC::MicroAPI::StoreDist; | 31 | using AscendC::MicroAPI::StoreDist; |
| 32 | using AscendC::MicroAPI::UpdateMask; | 32 | using AscendC::MicroAPI::UpdateMask; |
| 33 | +using AscendC::Reg::LoadAlign; | ||
| 34 | +using AscendC::Reg::Reduce; | ||
| 35 | +using AscendC::Reg::StoreAlign; | ||
| 33 | using NormCommon::NormCommonRegbase::LoadRegForDtype; | 36 | using NormCommon::NormCommonRegbase::LoadRegForDtype; |
| 34 | using NormCommon::NormCommonRegbase::StoreRegForDtype; | 37 | using NormCommon::NormCommonRegbase::StoreRegForDtype; |
| 35 | 38 | ||
| @@ -136,11 +139,11 @@ private: | |||
| 136 | LocalTensor<float> rstdOutUb = rstdQueue_.AllocTensor<float>(); | 139 | LocalTensor<float> rstdOutUb = rstdQueue_.AllocTensor<float>(); |
| 137 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); | 140 | LocalTensor<float> tmpTensor = tmpBuf.Get<float>(); |
| 138 | 141 | ||
| 139 | - __local_mem__ T* xInUbAddr = (__local_mem__ T*)xInUb.GetPhyAddr(); | 142 | + __ubuf__ T* xInUbAddr = (__ubuf__ T*)xInUb.GetPhyAddr(); |
| 140 | - __local_mem__ float* meanOutUbAddr = (__local_mem__ float*)meanOutUb.GetPhyAddr(); | 143 | + __ubuf__ float* meanOutUbAddr = (__ubuf__ float*)meanOutUb.GetPhyAddr(); |
| 141 | - __local_mem__ float* rstdOutUbAddr = (__local_mem__ float*)rstdOutUb.GetPhyAddr(); | 144 | + __ubuf__ float* rstdOutUbAddr = (__ubuf__ float*)rstdOutUb.GetPhyAddr(); |
| 142 | - __local_mem__ float* xSubMeanUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr(); | 145 | + __ubuf__ float* xSubMeanUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr(); |
| 143 | - __local_mem__ float* tmpUbAddr = (__local_mem__ float*)tmpTensor.GetPhyAddr() + elemNum_; | 146 | + __ubuf__ float* tmpUbAddr = (__ubuf__ float*)tmpTensor.GetPhyAddr() + elemNum_; |
| 144 | 147 | ||
| 145 | if (tl_->rAlign <= VL_B32) { | 148 | if (tl_->rAlign <= VL_B32) { |
| 146 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); | 149 | CalculateMeanVarRLessThanVL(xInUbAddr, meanOutUbAddr, rstdOutUbAddr, xSubMeanUbAddr, currentANum); |
| @@ -162,9 +165,9 @@ private: | |||
| 162 | CopyOutRstd(aOffset, currentANum); | 165 | CopyOutRstd(aOffset, currentANum); |
| 163 | 166 | ||
| 164 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); | 167 | LocalTensor<T> yOutUb = yQueue_.AllocTensor<T>(); |
| 165 | - __local_mem__ U* gammaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr(); | 168 | + __ubuf__ U* gammaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr(); |
| 166 | - __local_mem__ U* betaInUbAddr = (__local_mem__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAlign; | 169 | + __ubuf__ U* betaInUbAddr = (__ubuf__ U*)gammaBetaInUb_.GetPhyAddr() + tl_->rAlign; |
| 167 | - __local_mem__ T* yOutUbAddr = (__local_mem__ T*)yOutUb.GetPhyAddr(); | 170 | + __ubuf__ T* yOutUbAddr = (__ubuf__ T*)yOutUb.GetPhyAddr(); |
| 168 | if (hasGamma_ && hasBeta_) { | 171 | if (hasGamma_ && hasBeta_) { |
| 169 | CalculateNormalizeVF<true, true>(xSubMeanUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, rstdOutUbAddr, | 172 | CalculateNormalizeVF<true, true>(xSubMeanUbAddr, betaInUbAddr, gammaInUbAddr, yOutUbAddr, rstdOutUbAddr, |
| 170 | currentANum); | 173 | currentANum); |
| @@ -183,8 +186,8 @@ private: | |||
| 183 | yQueue_.EnQue(yOutUb); | 186 | yQueue_.EnQue(yOutUb); |
| 184 | } | 187 | } |
| 185 | 188 | ||
| 186 | - __aicore__ inline void CalculateMeanVarRLessThanVL(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 189 | + __aicore__ inline void CalculateMeanVarRLessThanVL(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 187 | - __local_mem__ float* rstdInUb, __local_mem__ float* xSubMeanUb, | 190 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 188 | uint16_t currentANum) | 191 | uint16_t currentANum) |
| 189 | { | 192 | { |
| 190 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 193 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| @@ -209,25 +212,25 @@ private: | |||
| 209 | for (uint16_t a = 0; a < currentANum; a++) { | 212 | for (uint16_t a = 0; a < currentANum; a++) { |
| 210 | LoadRegForDtype(xInUb, x, pregLoop, (a * aStride)); | 213 | LoadRegForDtype(xInUb, x, pregLoop, (a * aStride)); |
| 211 | Muls(meanSum, x, n, pregLoop); | 214 | Muls(meanSum, x, n, pregLoop); |
| 212 | - ReduceSum(mean, meanSum, pregLoop); | 215 | + Reduce<ReduceType::SUM>(mean, meanSum, pregLoop); |
| 213 | Muls(mean, mean, nCorrectionFactor, pregOne); | 216 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 214 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 217 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 215 | 218 | ||
| 216 | Duplicate(meanDup, mean, pregFull); | 219 | Duplicate(meanDup, mean, pregFull); |
| 217 | Sub(xMeanSub, x, meanDup, pregLoop); | 220 | Sub(xMeanSub, x, meanDup, pregLoop); |
| 218 | StoreRegForDtype(xSubMeanUb, xMeanSub, pregLoop, (a * aStride)); | 221 | StoreRegForDtype(xSubMeanUb, xMeanSub, pregLoop, (a * aStride)); |
| 219 | Mul(square, xMeanSub, xMeanSub, pregLoop); | 222 | Mul(square, xMeanSub, xMeanSub, pregLoop); |
| 220 | Muls(varSum, square, n, pregLoop); | 223 | Muls(varSum, square, n, pregLoop); |
| 221 | - ReduceSum(var, varSum, pregLoop); | 224 | + Reduce<ReduceType::SUM>(var, varSum, pregLoop); |
| 222 | Muls(var, var, nCorrectionFactor, pregOne); | 225 | Muls(var, var, nCorrectionFactor, pregOne); |
| 223 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 226 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 224 | } | 227 | } |
| 225 | } | 228 | } |
| 226 | } | 229 | } |
| 227 | 230 | ||
| 228 | - __aicore__ inline void CalculateMeanVarRLessThanTwoVL(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 231 | + __aicore__ inline void CalculateMeanVarRLessThanTwoVL(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 229 | - __local_mem__ float* rstdInUb, | 232 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 230 | - __local_mem__ float* xSubMeanUb, uint16_t currentANum) | 233 | + uint16_t currentANum) |
| 231 | { | 234 | { |
| 232 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 235 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| 233 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); | 236 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); |
| @@ -263,9 +266,9 @@ private: | |||
| 263 | Muls(meanSum1, x1, n, pregFull); | 266 | Muls(meanSum1, x1, n, pregFull); |
| 264 | Muls(meanSum2, x2, n, pregTail); | 267 | Muls(meanSum2, x2, n, pregTail); |
| 265 | Add(meanSum, meanSum1, meanSum2, pregFull); | 268 | Add(meanSum, meanSum1, meanSum2, pregFull); |
| 266 | - ReduceSum(mean, meanSum, pregFull); | 269 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 267 | Muls(mean, mean, nCorrectionFactor, pregOne); | 270 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 268 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 271 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 269 | 272 | ||
| 270 | Duplicate(meanDup, mean, pregFull); | 273 | Duplicate(meanDup, mean, pregFull); |
| 271 | Sub(xMeanSub1, x1, meanDup, pregFull); | 274 | Sub(xMeanSub1, x1, meanDup, pregFull); |
| @@ -277,17 +280,17 @@ private: | |||
| 277 | Muls(varSum1, square1, n, pregFull); | 280 | Muls(varSum1, square1, n, pregFull); |
| 278 | Muls(varSum2, square2, n, pregTail); | 281 | Muls(varSum2, square2, n, pregTail); |
| 279 | Add(varSum, varSum1, varSum2, pregFull); | 282 | Add(varSum, varSum1, varSum2, pregFull); |
| 280 | - ReduceSum(var, varSum, pregFull); | 283 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 281 | Muls(var, var, nCorrectionFactor, pregOne); | 284 | Muls(var, var, nCorrectionFactor, pregOne); |
| 282 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 285 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 283 | } | 286 | } |
| 284 | } | 287 | } |
| 285 | } | 288 | } |
| 286 | 289 | ||
| 287 | template <int32_t LAST_LOOP_NUMS> | 290 | template <int32_t LAST_LOOP_NUMS> |
| 288 | - __aicore__ inline void CalculateMeanVarRCommon(__local_mem__ T* xInUb, __local_mem__ float* meanInUb, | 291 | + __aicore__ inline void CalculateMeanVarRCommon(__ubuf__ T* xInUb, __ubuf__ float* meanInUb, |
| 289 | - __local_mem__ float* rstdInUb, __local_mem__ float* xSubMeanUb, | 292 | + __ubuf__ float* rstdInUb, __ubuf__ float* xSubMeanUb, |
| 290 | - __local_mem__ float* tmpUb, uint16_t currentANum) | 293 | + __ubuf__ float* tmpUb, uint16_t currentANum) |
| 291 | { | 294 | { |
| 292 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); | 295 | uint32_t reduceNum = static_cast<uint32_t>(tl_->r); |
| 293 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); | 296 | float n = static_cast<float>(1.0) / static_cast<float>(tl_->powerOfTwoForR); |
| @@ -326,8 +329,8 @@ private: | |||
| 326 | Muls(x1, x1, n, pregFull); | 329 | Muls(x1, x1, n, pregFull); |
| 327 | Muls(x2, x2, n, pregFull); | 330 | Muls(x2, x2, n, pregFull); |
| 328 | Add(meanSum, x1, x2, pregFull); | 331 | Add(meanSum, x1, x2, pregFull); |
| 329 | - ReduceSum(mean, meanSum, pregFull); | 332 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 330 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 333 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 331 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), mean, pregOne); | 334 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), mean, pregOne); |
| 332 | } | 335 | } |
| 333 | for (uint16_t r = 0; | 336 | for (uint16_t r = 0; |
| @@ -340,8 +343,8 @@ private: | |||
| 340 | Muls(x1, x1, n, pregFull); | 343 | Muls(x1, x1, n, pregFull); |
| 341 | Muls(x2, x2, n, pregLoop); | 344 | Muls(x2, x2, n, pregLoop); |
| 342 | Add(meanSum, x1, x2, pregFull); | 345 | Add(meanSum, x1, x2, pregFull); |
| 343 | - ReduceSum(mean, meanSum, pregFull); | 346 | + Reduce<ReduceType::SUM>(mean, meanSum, pregFull); |
| 344 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 347 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 345 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, | 348 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), mean, |
| 346 | pregOne); | 349 | pregOne); |
| 347 | } | 350 | } |
| @@ -350,8 +353,8 @@ private: | |||
| 350 | LoadRegForDtype(xInUb + binaryAddRemainderCeilLoop * VL_B32, x1, pregFull, | 353 | LoadRegForDtype(xInUb + binaryAddRemainderCeilLoop * VL_B32, x1, pregFull, |
| 351 | (r * VL_B32 + a * aStride)); | 354 | (r * VL_B32 + a * aStride)); |
| 352 | Muls(x1, x1, n, pregFull); | 355 | Muls(x1, x1, n, pregFull); |
| 353 | - ReduceSum(mean, x1, pregFull); | 356 | + Reduce<ReduceType::SUM>(mean, x1, pregFull); |
| 354 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 357 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 355 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, | 358 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), mean, |
| 356 | pregOne); | 359 | pregOne); |
| 357 | } | 360 | } |
| @@ -360,24 +363,24 @@ private: | |||
| 360 | if constexpr (LAST_LOOP_NUMS == 1) { | 363 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 361 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); | 364 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNum); |
| 362 | for (uint16_t a = 0; a < currentANum; a++) { | 365 | for (uint16_t a = 0; a < currentANum; a++) { |
| 363 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 366 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 364 | - ReduceSum(mean, x1, pregLast); | 367 | + Reduce<ReduceType::SUM>(mean, x1, pregLast); |
| 365 | Muls(mean, mean, nCorrectionFactor, pregOne); | 368 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 366 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 369 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 367 | } | 370 | } |
| 368 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 371 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 369 | for (uint16_t a = 0; a < currentANum; a++) { | 372 | for (uint16_t a = 0; a < currentANum; a++) { |
| 370 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; | 373 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; |
| 371 | MaskReg pregLast = UpdateMask<float>(lastTailNum); | 374 | MaskReg pregLast = UpdateMask<float>(lastTailNum); |
| 372 | RegTensor<float> shlReg; | 375 | RegTensor<float> shlReg; |
| 373 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 376 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 374 | - DataCopy(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); | 377 | + LoadAlign(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); |
| 375 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), | 378 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), |
| 376 | pregLast); | 379 | pregLast); |
| 377 | Add(x1, x1, shlReg, pregFull); | 380 | Add(x1, x1, shlReg, pregFull); |
| 378 | - ReduceSum(mean, x1, pregFull); | 381 | + Reduce<ReduceType::SUM>(mean, x1, pregFull); |
| 379 | Muls(mean, mean, nCorrectionFactor, pregOne); | 382 | Muls(mean, mean, nCorrectionFactor, pregOne); |
| 380 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); | 383 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(meanInUb + a, mean, pregOne); |
| 381 | } | 384 | } |
| 382 | } | 385 | } |
| 383 | } | 386 | } |
| @@ -397,7 +400,7 @@ private: | |||
| 397 | MaskReg pregLoop; | 400 | MaskReg pregLoop; |
| 398 | 401 | ||
| 399 | for (uint16_t a = 0; a < currentANum; a++) { | 402 | for (uint16_t a = 0; a < currentANum; a++) { |
| 400 | - DataCopy<float, LoadDist::DIST_BRC_B32>(mean, meanInUb + a); | 403 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(mean, meanInUb + a); |
| 401 | uint32_t sregRemainder = binaryAddRemainder; | 404 | uint32_t sregRemainder = binaryAddRemainder; |
| 402 | for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; r++) { | 405 | for (uint16_t r = 0; r < binaryAddRemainderFloorLoop; r++) { |
| 403 | pregLoop = UpdateMask<float>(sregRemainder); | 406 | pregLoop = UpdateMask<float>(sregRemainder); |
| @@ -412,8 +415,8 @@ private: | |||
| 412 | Muls(square1, square1, n, pregFull); | 415 | Muls(square1, square1, n, pregFull); |
| 413 | Muls(square2, square2, n, pregFull); | 416 | Muls(square2, square2, n, pregFull); |
| 414 | Add(varSum, square1, square2, pregFull); | 417 | Add(varSum, square1, square2, pregFull); |
| 415 | - ReduceSum(var, varSum, pregFull); | 418 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 416 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 419 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 417 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), var, pregOne); | 420 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + r), var, pregOne); |
| 418 | } | 421 | } |
| 419 | for (uint16_t r = 0; | 422 | for (uint16_t r = 0; |
| @@ -434,8 +437,8 @@ private: | |||
| 434 | Muls(square1, square1, n, pregFull); | 437 | Muls(square1, square1, n, pregFull); |
| 435 | Muls(square2, square2, n, pregLoop); | 438 | Muls(square2, square2, n, pregLoop); |
| 436 | Add(varSum, square1, square2, pregFull); | 439 | Add(varSum, square1, square2, pregFull); |
| 437 | - ReduceSum(var, varSum, pregFull); | 440 | + Reduce<ReduceType::SUM>(var, varSum, pregFull); |
| 438 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 441 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 439 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), var, | 442 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderFloorLoop), var, |
| 440 | pregOne); | 443 | pregOne); |
| 441 | } | 444 | } |
| @@ -448,8 +451,8 @@ private: | |||
| 448 | (r * VL_B32 + a * aStride)); | 451 | (r * VL_B32 + a * aStride)); |
| 449 | Mul(square1, xMeanSub1, xMeanSub1, pregFull); | 452 | Mul(square1, xMeanSub1, xMeanSub1, pregFull); |
| 450 | Muls(square1, square1, n, pregFull); | 453 | Muls(square1, square1, n, pregFull); |
| 451 | - ReduceSum(var, square1, pregFull); | 454 | + Reduce<ReduceType::SUM>(var, square1, pregFull); |
| 452 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>( | 455 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>( |
| 453 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), var, | 456 | tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + binaryAddRemainderCeilLoop + r), var, |
| 454 | pregOne); | 457 | pregOne); |
| 455 | } | 458 | } |
| @@ -458,41 +461,40 @@ private: | |||
| 458 | if constexpr (LAST_LOOP_NUMS == 1) { | 461 | if constexpr (LAST_LOOP_NUMS == 1) { |
| 459 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNumTmp); | 462 | MaskReg pregLast = UpdateMask<float>(lastBinaryAddNumTmp); |
| 460 | for (uint16_t a = 0; a < currentANum; a++) { | 463 | for (uint16_t a = 0; a < currentANum; a++) { |
| 461 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 464 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 462 | - ReduceSum(var, x1, pregLast); | 465 | + Reduce<ReduceType::SUM>(var, x1, pregLast); |
| 463 | Muls(var, var, nCorrectionFactor, pregOne); | 466 | Muls(var, var, nCorrectionFactor, pregOne); |
| 464 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 467 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 465 | } | 468 | } |
| 466 | } else if constexpr (LAST_LOOP_NUMS == 2) { | 469 | } else if constexpr (LAST_LOOP_NUMS == 2) { |
| 467 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; | 470 | uint32_t lastTailNum = lastBinaryAddNum - VL_B32; |
| 468 | MaskReg pregLast = UpdateMask<float>(lastTailNum); | 471 | MaskReg pregLast = UpdateMask<float>(lastTailNum); |
| 469 | RegTensor<float> shlReg; | 472 | RegTensor<float> shlReg; |
| 470 | for (uint16_t a = 0; a < currentANum; a++) { | 473 | for (uint16_t a = 0; a < currentANum; a++) { |
| 471 | - DataCopy(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); | 474 | + LoadAlign(x1, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign)); |
| 472 | - DataCopy(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); | 475 | + LoadAlign(x2, tmpUb + static_cast<uint32_t>(a * lastBinaryAddNumAlign + VL_B32)); |
| 473 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), | 476 | ShiftLefts((RegTensor<uint32_t>&)shlReg, (RegTensor<uint32_t>&)x2, static_cast<int16_t>(0), |
| 474 | pregLast); | 477 | pregLast); |
| 475 | Add(x1, x1, shlReg, pregFull); | 478 | Add(x1, x1, shlReg, pregFull); |
| 476 | - ReduceSum(var, x1, pregFull); | 479 | + Reduce<ReduceType::SUM>(var, x1, pregFull); |
| 477 | Muls(var, var, nCorrectionFactor, pregOne); | 480 | Muls(var, var, nCorrectionFactor, pregOne); |
| 478 | - DataCopy<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); | 481 | + StoreAlign<float, StoreDist::DIST_FIRST_ELEMENT_B32>(rstdInUb + a, var, pregOne); |
| 479 | } | 482 | } |
| 480 | } | 483 | } |
| 481 | } | 484 | } |
| 482 | } | 485 | } |
| 483 | 486 | ||
| 484 | template <bool hasGammaFlag, bool hasBetaFlag> | 487 | template <bool hasGammaFlag, bool hasBetaFlag> |
| 485 | - __aicore__ inline void CalculateNormalizeVF(__local_mem__ float* xSubMeanUb, __local_mem__ U* betaInUb, | 488 | + __aicore__ inline void CalculateNormalizeVF(__ubuf__ float* xSubMeanUb, __ubuf__ U* betaInUb, __ubuf__ U* gammaInUb, |
| 486 | - __local_mem__ U* gammaInUb, __local_mem__ T* yOutUb, | 489 | + __ubuf__ T* yOutUb, __ubuf__ float* rstdOutUb, uint16_t currentANum) |
| 487 | - __local_mem__ float* rstdOutUb, uint16_t currentANum) | ||
| 488 | { | 490 | { |
| 489 | uint32_t reduceNum = tl_->r; | 491 | uint32_t reduceNum = tl_->r; |
| 490 | uint32_t aStride = tl_->rAlign; | 492 | uint32_t aStride = tl_->rAlign; |
| 491 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; | 493 | uint16_t loopCount = (reduceNum + VL_B32 - 1) / VL_B32; |
| 492 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; | 494 | uint32_t remainderA = currentANum / NUM_TWO * NUM_TWO; |
| 493 | uint16_t remainderLoop = currentANum - remainderA; | 495 | uint16_t remainderLoop = currentANum - remainderA; |
| 494 | - __local_mem__ float* rstdOutUbPair = rstdOutUb + 1; | 496 | + __ubuf__ float* rstdOutUbPair = rstdOutUb + 1; |
| 495 | - __local_mem__ float* rstdOutUbRemainder = rstdOutUb + remainderA; | 497 | + __ubuf__ float* rstdOutUbRemainder = rstdOutUb + remainderA; |
| 496 | 498 | ||
| 497 | __VEC_SCOPE__ | 499 | __VEC_SCOPE__ |
| 498 | { | 500 | { |
| @@ -512,8 +514,8 @@ private: | |||
| 512 | MaskReg pregLoop; | 514 | MaskReg pregLoop; |
| 513 | 515 | ||
| 514 | for (uint16_t a = 0; a < static_cast<uint16_t>(currentANum / static_cast<uint16_t>(NUM_TWO)); a++) { | 516 | for (uint16_t a = 0; a < static_cast<uint16_t>(currentANum / static_cast<uint16_t>(NUM_TWO)); a++) { |
| 515 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); | 517 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt1, rstdOutUb + a * NUM_TWO); |
| 516 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); | 518 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrt2, rstdOutUbPair + a * NUM_TWO); |
| 517 | uint32_t sreg0 = reduceNum; | 519 | uint32_t sreg0 = reduceNum; |
| 518 | for (uint16_t r = 0; r < loopCount; r++) { | 520 | for (uint16_t r = 0; r < loopCount; r++) { |
| 519 | pregLoop = UpdateMask<float>(sreg0); | 521 | pregLoop = UpdateMask<float>(sreg0); |
| @@ -528,8 +530,8 @@ private: | |||
| 528 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); | 530 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); |
| 529 | } | 531 | } |
| 530 | if constexpr (hasGammaFlag && hasBetaFlag) { | 532 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 531 | - FusedMulDstAdd(y1, gamma, beta, pregLoop); | 533 | + MulDstAdd(y1, gamma, beta, pregLoop); |
| 532 | - FusedMulDstAdd(y2, gamma, beta, pregLoop); | 534 | + MulDstAdd(y2, gamma, beta, pregLoop); |
| 533 | } else { | 535 | } else { |
| 534 | if constexpr (hasGammaFlag) { | 536 | if constexpr (hasGammaFlag) { |
| 535 | Mul(y1, y1, gamma, pregLoop); | 537 | Mul(y1, y1, gamma, pregLoop); |
| @@ -545,7 +547,7 @@ private: | |||
| 545 | } | 547 | } |
| 546 | } | 548 | } |
| 547 | for (uint16_t a = 0; a < remainderLoop; a++) { | 549 | for (uint16_t a = 0; a < remainderLoop; a++) { |
| 548 | - DataCopy<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); | 550 | + LoadAlign<float, LoadDist::DIST_BRC_B32>(rsqrtRemainder, rstdOutUbRemainder); |
| 549 | uint32_t sreg1 = reduceNum; | 551 | uint32_t sreg1 = reduceNum; |
| 550 | for (uint16_t r = 0; r < loopCount; r++) { | 552 | for (uint16_t r = 0; r < loopCount; r++) { |
| 551 | pregLoop = UpdateMask<float>(sreg1); | 553 | pregLoop = UpdateMask<float>(sreg1); |
| @@ -558,7 +560,7 @@ private: | |||
| 558 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); | 560 | LoadRegForDtype(betaInUb, beta, pregLoop, (r * VL_B32)); |
| 559 | } | 561 | } |
| 560 | if constexpr (hasGammaFlag && hasBetaFlag) { | 562 | if constexpr (hasGammaFlag && hasBetaFlag) { |
| 561 | - FusedMulDstAdd(yRemainder, gamma, beta, pregLoop); | 563 | + MulDstAdd(yRemainder, gamma, beta, pregLoop); |
| 562 | } else { | 564 | } else { |
| 563 | if constexpr (hasGammaFlag) { | 565 | if constexpr (hasGammaFlag) { |
| 564 | Mul(yRemainder, yRemainder, gamma, pregLoop); | 566 | Mul(yRemainder, yRemainder, gamma, pregLoop); |
| @@ -23,6 +23,7 @@ | |||
| 23 | 23 | ||
| 24 | namespace LayerNormV4 { | 24 | namespace LayerNormV4 { |
| 25 | using namespace AscendC; | 25 | using namespace AscendC; |
| 26 | +using AscendC::Reg::StoreAlign; | ||
| 26 | 27 | ||
| 27 | namespace LayerNormV4Regbase { | 28 | namespace LayerNormV4Regbase { |
| 28 | __aicore__ inline constexpr uint32_t GetVRegSize() | 29 | __aicore__ inline constexpr uint32_t GetVRegSize() |
| @@ -183,9 +184,9 @@ private: | |||
| 183 | { | 184 | { |
| 184 | if constexpr (!IsSameType<M, float>::value) { | 185 | if constexpr (!IsSameType<M, float>::value) { |
| 185 | // float to bfloat16 or float16, input continue and output each repeat have only half value | 186 | // float to bfloat16 or float16, input continue and output each repeat have only half value |
| 186 | - CastBatchMeanRstdToDtype<M>( | 187 | + CastBatchMeanRstdToDtype<M>((__ubuf__ float*)meanTensor.GetPhyAddr(), |
| 187 | - (__local_mem__ float*)meanTensor.GetPhyAddr(), (__local_mem__ float*)rstdTensor.GetPhyAddr(), | 188 | + (__ubuf__ float*)rstdTensor.GetPhyAddr(), (__ubuf__ M*)meanTensor.GetPhyAddr(), |
| 188 | - (__local_mem__ M*)meanTensor.GetPhyAddr(), (__local_mem__ M*)rstdTensor.GetPhyAddr(), cacheCount); | 189 | + (__ubuf__ M*)rstdTensor.GetPhyAddr(), cacheCount); |
| 189 | outQueueMean.EnQue(meanTensor); | 190 | outQueueMean.EnQue(meanTensor); |
| 190 | outQueueRstd.EnQue(rstdTensor); | 191 | outQueueRstd.EnQue(rstdTensor); |
| 191 | meanTensor = outQueueMean.template DeQue<float>(); | 192 | meanTensor = outQueueMean.template DeQue<float>(); |
| @@ -243,16 +244,16 @@ private: | |||
| 243 | uint16_t loopTimes = (elemCnt + VL_B32 - 1) / VL_B32; | 244 | uint16_t loopTimes = (elemCnt + VL_B32 - 1) / VL_B32; |
| 244 | __VEC_SCOPE__ | 245 | __VEC_SCOPE__ |
| 245 | { | 246 | { |
| 246 | - __local_mem__ float* meamPtr = (__local_mem__ float*)mean.GetPhyAddr(); | 247 | + __ubuf__ float* meamPtr = (__ubuf__ float*)mean.GetPhyAddr(); |
| 247 | - __local_mem__ float* variancePtr = (__local_mem__ float*)variance.GetPhyAddr(); | 248 | + __ubuf__ float* variancePtr = (__ubuf__ float*)variance.GetPhyAddr(); |
| 248 | uint32_t count = static_cast<uint32_t>(elemCnt); | 249 | uint32_t count = static_cast<uint32_t>(elemCnt); |
| 249 | AscendC::MicroAPI::RegTensor<float> xReg; | 250 | AscendC::MicroAPI::RegTensor<float> xReg; |
| 250 | AscendC::MicroAPI::MaskReg pMask; | 251 | AscendC::MicroAPI::MaskReg pMask; |
| 251 | Duplicate(xReg, 0.0f); | 252 | Duplicate(xReg, 0.0f); |
| 252 | for (uint16_t i = 0; i < loopTimes; ++i) { | 253 | for (uint16_t i = 0; i < loopTimes; ++i) { |
| 253 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); | 254 | pMask = AscendC::MicroAPI::UpdateMask<float>(count); |
| 254 | - DataCopy((__local_mem__ float*)meamPtr + i * VL_B32, xReg, pMask); | 255 | + StoreAlign((__ubuf__ float*)meamPtr + i * VL_B32, xReg, pMask); |
| 255 | - DataCopy((__local_mem__ float*)variancePtr + i * VL_B32, xReg, pMask); | 256 | + StoreAlign((__ubuf__ float*)variancePtr + i * VL_B32, xReg, pMask); |
| 256 | } | 257 | } |
| 257 | } | 258 | } |
| 258 | } | 259 | } |